Adapting Self-Supervised Representations as a Latent Space for Efficient Generation

Título: Adaptação de Representações Auto-Supervisionadas como Espaço Latente para Geração Eficiente (RepTok)

1. O Problema

Os modelos generativos modernos, como Difusão e Modelos Baseados em Fluxo (Flow Matching), têm alcançado resultados impressionantes na síntese de imagens e vídeos. No entanto, eles enfrentam desafios significativos:

• Custo Computacional: A maioria opera no espaço de pixels de alta dimensão ou em espaços latentes bidimensionais (2D) organizados em grades, o que exige grandes quantidades de operações de ponto flutuante (FLOPs) e mecanismos de atenção complexos.
• Redundância Espacial: Os espaços latentes tradicionais (como os usados em Latent Diffusion Models - LDMs) ainda mantêm uma estrutura de grade 2D, falhando em explorar as altas redundâncias espaciais inerentes às imagens naturais.
• Ineficiência em Tokens: Métodos recentes que tentam comprimir imagens em sequências 1D (como TiTok) ainda utilizam múltiplos tokens discretos (ex: 32 tokens), mantendo uma certa complexidade.

O objetivo é criar um framework que represente uma imagem inteira com um único token contínuo, eliminando a redundância espacial e reduzindo drasticamente o custo de treinamento e inferência, sem sacrificar a qualidade da reconstrução ou geração.

2. Metodologia: RepTok

O RepTok (Representation Tokenizer) é um framework que utiliza representações de modelos de visão auto-supervisionados (SSL) pré-treinados como espaço latente para geração.

• Arquitetura Base: O sistema utiliza um encoder SSL pré-treinado (ex: DINOv2, MAE, CLIP) que possui um token de classe ([cls]). Este token agrega informações de todo o patch da imagem.
• Adaptação Mínima (Fine-tuning): Em vez de treinar um novo encoder do zero, o RepTok congela a maior parte do encoder SSL e apenas ajusta (fine-tune) o embedding do token [cls]. Isso permite injetar detalhes de baixo nível (necessários para reconstrução fiel) no token sem destruir a geometria semântica suave do espaço original.
• Decodificador Generativo: Um decodificador generativo leve é treinado conjuntamente para mapear o token latente único de volta para o espaço de pixels (ou espaço latente de baixa dimensão, como no SD-VAE).
• Função de Perda (Flow Matching + Regularização):
  1. Flow Matching: O decodificador é treinado com o objetivo padrão de Flow Matching para aprender o campo vetorial entre uma distribuição prior (ruído gaussiano) e a distribuição de dados.
  2. Perda de Similaridade Cosseno: Para evitar que o token ajustado se afaste demais do espaço latente original (que é bem estruturado e suave), é adicionada uma perda de regularização baseada na similaridade cosseno entre o token ajustado e o token do encoder congelado original. Isso equilibra a necessidade de detalhes de reconstrução com a preservação da estrutura geométrica para geração.
• Geração Sem Atenção: Como a imagem é representada por um único token, não há necessidade de mecanismos de atenção (self-attention) entre tokens. O gerador pode ser uma arquitetura simples baseada em MLP-Mixer, eliminando o custo quadrático da atenção.

3. Principais Contribuições

1. Token Único Contínuo: Demonstração de que o token [cls] de modelos SSL, quando levemente adaptado, pode servir como um espaço latente compacto e expressivo para representar imagens inteiras, substituindo grades 2D ou sequências longas de tokens.
2. Eficiência Computacional Extrema: A abordagem reduz o custo de treinamento em mais de 90% em comparação com baselines baseadas em Transformers (como DiT e SiT), permitindo o uso de arquiteturas sem atenção (MLP-Mixer).
3. Reconstrução Fiel e Geração de Alta Qualidade: O método consegue reconstruir imagens com alta fidelidade (preservando detalhes de baixo nível) e gerar novas imagens com qualidade competitiva, mantendo a suavidade geométrica do espaço SSL original.
4. Escalabilidade para Texto-para-Imagem (T2I): O framework estende-se naturalmente para síntese T2I, incorporando atenção cruzada para condicionamento textual, alcançando desempenho zero-shot competitivo no MS-COCO com um orçamento de treinamento extremamente limitado (< 20 horas em 4 GPUs A100).

4. Resultados Experimentais

• Geração Condicional (ImageNet-1K):
  • O RepTok alcançou um gFID de 1.88 (com CFG) no ImageNet 256x256, competindo diretamente com modelos muito maiores e mais complexos.
  • Custo: O treinamento consumiu apenas 34.5 K PetaFLOPs, comparado a mais de 200 K PetaFLOPs para baselines como SiT-XL/2. Isso representa uma redução de custo de cerca de 98% em relação a alguns baselines.
  • Reconstrução: O modelo obteve um rFID de 1.85 e PSNR de 14.94, superando métodos que usam apenas códigos semânticos puros (como RCG).
• Síntese Texto-para-Imagem (MS-COCO):
  • Alcançou desempenho competitivo em zero-shot no MS-COCO, superando modelos estabelecidos com menos de 200k iterações de treinamento.
  • A escalabilidade dos backbones de linguagem (CLIP, InternVL, Gemma) mostrou melhoria consistente nas métricas.
• Interpolação no Espaço Latente:
  • A interpolação linear entre tokens latentes resultou em transições suaves tanto no conteúdo semântico quanto na configuração espacial (forma, tamanho, rotação), confirmando que a geometria do espaço SSL foi preservada.
• Velocidade de Inferência:
  • O tempo total de inferência foi de 0.27 segundos (vs. 2.43s do SiT-XL), dominado pelo decodificador, mas com um gerador (MLP-Mixer) extremamente rápido.

5. Significado e Impacto

O trabalho RepTok redefine a eficiência na modelagem generativa ao demonstrar que representações SSL pré-treinadas podem ser adaptadas diretamente como espaços latentes, em vez de serem usadas apenas como guias ou regularizadores.

• Mudança de Paradigma: Move o foco de arquiteturas complexas de atenção para espaços latentes ultra-compressos (1 token), onde a complexidade arquitetural é deslocada para a fase de compressão pré-treinada.
• Acessibilidade: Ao reduzir drasticamente o custo computacional (FLOPs) e o tempo de treinamento, torna viável o treinamento de modelos generativos de alta qualidade em recursos computacionais limitados.
• Generalidade: A abordagem é agnóstica ao encoder SSL, funcionando bem com DINOv2, MAE e CLIP, sugerindo um caminho robusto para futuras pesquisas em compressão e geração.

Em resumo, o RepTok prova que a combinação de representações SSL finamente ajustadas com decodificadores leves baseados em Flow Matching oferece um caminho promissor para a próxima geração de modelos generativos eficientes e de alto desempenho.