MacTok: Robust Continuous Tokenization for Image Generation

1. O Problema

O campo de geração de imagens baseada em modelos de difusão e autoregressivos depende criticamente de tokenizadores de imagem para comprimir pixels em representações latentes. Existem duas abordagens principais:

• Discretos (ex: VQ-VAE): Usam um código finito, mas sofrem com erros de quantização e fluxo de gradiente limitado.
• Contínuos (ex: KL-VAE): Usam espaços latentes contínuos regularizados por divergência KL, permitindo representações suaves. No entanto, sob forte compressão (poucos tokens), eles sofrem de colapso posterior (posterior collapse).

O colapso posterior ocorre quando o KL-VAE, sob forte regularização, ignora as características informativas da imagem e força a distribuição latente a se assemelhar excessivamente a uma prior gaussiana isotrópica. Isso resulta em latentes que não carregam semântica útil, levando a reconstruções pobres e geração de baixa fidelidade. Métodos anteriores tentaram resolver isso ajustando pesos de hiperparâmetros, mas sem uma solução fundamental.

2. Metodologia: MacTok

Os autores propõem o MacTok (Masked Augmenting 1D Continuous Tokenizer), um tokenizador contínuo 1D que combate o colapso posterior através de duas estratégias principais: Mascaramento de Imagem e Alinhamento de Representação.

Arquitetura

• Utiliza uma arquitetura baseada em Vision Transformer (ViT) tanto para o codificador quanto para o decodificador.
• Processa tokens de imagem e tokens latentes (compactos) em uma sequência unificada.
• O espaço latente é modelado como uma variável aleatória gaussiana com regularização KL.

Estratégias Chave

A. Mascaramento de Imagem (Image Masking)
Para forçar o modelo a codificar informações robustas mesmo com poucos tokens, o MacTok treina o modelo para reconstruir imagens completas a partir de entradas parciais (mascaradas).

• Mascaramento Aleatório: Substitui aleatoriamente um subconjunto de patches da imagem por tokens de máscara antes da codificação. Isso força os tokens latentes a inferir o conteúdo ausente a partir do contexto visível.
• Mascaramento Semântico Guiado por DINO: Diferente do aleatório, este método usa as características de um modelo pré-treinado DINOv2 para identificar e mascarar as regiões mais informativas da imagem (com base na similaridade entre o token de classificação e os patches). Isso força o modelo a aprender a reconstruir estruturas semânticas críticas a partir de evidências visuais incompletas, transferindo priores semânticos do espaço de imagem para o espaço latente.

B. Alinhamento de Representação Global e Local
Para garantir que o espaço latente seja semanticamente estruturado e discriminativo:

• Alinhamento Local: Os tokens latentes expandidos são alinhados com os patches de características do DINOv2.
• Alinhamento Global: A média dos tokens latentes é alinhada com o token de classificação (cls token) do DINOv2.
• Isso atua como um objetivo auxiliar que organiza o espaço latente, mantendo a coerência estrutural e a preservação de detalhes mesmo com alta compressão.

C. Função de Perda
O treinamento combina perda de reconstrução, perda perceptual, adversarial (GAN), regularização KL e a nova perda de alinhamento de representação.

3. Contribuições Principais

1. Identificação e Solução do Colapso: Demonstrar que o colapso posterior em tokenizadores contínuos é uma forma de super-regularização e que o mascaramento de reconstrução (especialmente o guiado semanticamente) estabiliza o aprendizado latente sob forte compressão.
2. Arquitetura MacTok: Introdução de um tokenizador contínuo 1D que combina mascaramento de imagem (aleatório e semântico) com alinhamento de representação global/local, prevenindo o colapso sem sacrificar a fidelidade.
3. Eficiência Extrema: Conseguir alta fidelidade de geração e reconstrução usando apenas 64 ou 128 tokens, reduzindo o uso de tokens em até 64x comparado a métodos tradicionais que usam 1024 tokens.

4. Resultados Experimentais

Os experimentos foram realizados no dataset ImageNet (resoluções 256x256 e 512x512), utilizando geradores como SiT-XL e LightningDiT.

• Qualidade de Geração (gFID):
  • Em 256x256: Alcançou um gFID de 1.44 (com SiT-XL e 128 tokens).
  • Em 512x512: Alcançou um gFID de 1.52 (com SiT-XL e 128 tokens), estabelecendo um novo estado da arte (SOTA) para essa resolução com tão poucos tokens.
• Reconstrução (rFID):
  • Obteve rFID de 0.43 (256x256, 128 tokens) e 0.79 (512x512, 128 tokens), superando tokenizadores baseados em VQ que exigem muito mais tokens para resultados similares.
• Comparação: O MacTok supera significativamente métodos contínuos anteriores (como SoftVQ-VAE, MAETok) e compete com os melhores modelos discretos, mas com a vantagem da suavidade do espaço contínuo e eficiência de tokens.
• Análise de Espaço Latente: Visualizações mostram que, ao contrário do colapso posterior (distribuição isotrópica e sem estrutura), o MacTok aprende um espaço latente bem estruturado onde conceitos semânticos similares se agrupam, permitindo até mesmo classificação linear precisa.

5. Significado e Impacto

O MacTok representa um avanço significativo na tokenização eficiente para geração de imagens. Ao resolver o problema crônico do colapso posterior em modelos variacionais contínuos através de mascaramento e alinhamento semântico, ele permite:

• Redução drástica de custo computacional: Menos tokens significam sequências mais curtas para modelos autoregressivos ou difusivos, acelerando o treinamento e a inferência.
• Alta Fidelidade com Compressão: Demonstra que é possível manter detalhes finos e coerência semântica com uma representação latente extremamente compacta (1D).
• Viabilidade de Modelos Contínuos: Reabilita os tokenizadores contínuos como uma alternativa viável e superior aos discretos em cenários onde a suavidade do espaço latente e a eficiência de tokens são críticas.

Em resumo, o MacTok prova que a combinação de mascaramento de dados (para forçar a robustez) e alinhamento semântico (para estruturar o espaço) é a chave para criar tokenizadores contínuos de alta performance, permitindo a próxima geração de modelos de geração de imagem mais rápidos e eficientes.