VFM-VAE: Vision Foundation Models Can Be Good Tokenizers for Latent Diffusion Models

Título: VFM-VAE: Modelos Fundacionais de Visão Podem Ser Bons Tokenizadores para Modelos de Difusão Latente

1. O Problema

Os Modelos de Difusão Latente (LDMs) dependem criticamente da qualidade de seus tokenizadores visuais (geralmente Autoencoders Variacionais ou VAEs) para codificar imagens em um espaço latente compacto.

• Abordagem Atual: Trabalhos recentes tentaram melhorar a qualidade semântica dos tokenizadores alinhando os latentes do VAE com representações de Modelos Fundacionais de Visão (VFMs) (como DINOv2) através de distilação (treinamento do VAE para imitar o VFM).
• Limitação Identificada: Os autores observaram empiricamente que essa distilação baseada em alinhamento inevitavelmente enfraquece a robustez das representações aprendidas. Métricas como CKNNA (Centered Kernel Nearest-Neighbor Alignment) mostram que, sob transformações que preservam a semântica (ruído, rotação, escala), as representações distiladas tornam-se frágeis e perdem informações críticas, resultando em degradação em comparação ao VFM original.

2. Metodologia: VFM-VAE

O artigo propõe uma abordagem direta para contornar a distilação: utilizar um VFM congelado (frozen) diretamente como o codificador (encoder) do tokenizador, construindo apenas um decodificador especializado para reconstruir a imagem.

Arquitetura Proposta:

1. Encoder (VFM Congelado):

   • Em vez de treinar um encoder do zero, o modelo utiliza um VFM pré-treinado e congelado (ex: SigLIP2-Large).
   • Extrai características de múltiplas escalas (camadas rasas, intermediárias e finais) para capturar desde detalhes espaciais finos até semântica de alto nível.
   • Um projeto leve (Attention Projection) comprime essas características multi-escala em um espaço latente de baixa dimensão, preservando a riqueza semântica.

2. Decodificador (Novo Design):

   • Como os VFMs são otimizados para semântica e não para reconstrução de pixels, o decodificador padrão falha. O VFM-VAE introduz duas inovações:
     • Fusão de Latentes Multi-Escala: Decompõe o latente em componentes globais (estilo) e espaciais (estrutura), injetando-os em diferentes estágios do decodificador.
     • Blocos de Reconstrução com Resolução Progressiva: O decodificador sintetiza a imagem em etapas (de baixa para alta resolução). Blocos iniciais usam injeção espacial para definir o layout, enquanto blocos posteriores focam em detalhes finos, controlados por um estilo global modulado.

3. Objetivos de Treinamento:

   • Combina perda de reconstrução (L1, LPIPS, Adversarial) com uma perda de regularização de representação que garante que o latente compacto mantenha a similaridade semântica com o VFM original, sem distilação agressiva.

4. Nova Métrica: SE-CKNNA

   • Os autores propõem o Semantic-Equivariant CKNNA (SE-CKNNA). Diferente do CKNNA padrão, que mede alinhamento apenas em imagens limpas, o SE-CKNNA avalia a estabilidade do alinhamento sob um conjunto de transformações que preservam a semântica (ruído, escala, rotação). Isso revela a robustez real do tokenizador.

3. Contribuições Principais

1. Uso Direto de VFM Congelado: Elimina a degradação causada pela distilação, utilizando o VFM como encoder front-end, enquanto mantém alta fidelidade de reconstrução através de um decodificador especializado.
2. Métrica SE-CKNNA: Introduz uma ferramenta diagnóstica mais rigorosa para avaliar o alinhamento entre tokenizadores e VFMs sob perturbações.
3. Alinhamento de Dupla Face (Dual-Side Alignment): Demonstra que alinhar tanto o tokenizador quanto o modelo de difusão com o VFM gera representações internas superiores no modelo de difusão, especialmente em camadas profundas.
4. Eficiência e Desempenho: O sistema alcança resultados de ponta com treinamento significativamente mais rápido.

4. Resultados Experimentais

Os experimentos foram realizados no conjunto de dados ImageNet (256x256).

• Velocidade de Convergência: O sistema VFM-VAE atinge um gFID (sem CFG) de 2.22 em apenas 80 épocas, representando um aceleração de 10x em comparação a tokenizadores anteriores (como VA-VAE).
• Qualidade Final: Com treinamento estendido para 640 épocas, o modelo atinge um gFID de 1.62, superando métodos anteriores e estabelecendo um novo estado da arte em eficiência.
• Robustez: O VFM-VAE demonstra uma queda mínima (+1.6%) entre CKNNA e SE-CKNNA, enquanto o VA-VAE sofre uma degradação de -33.2%, provando sua superioridade em manter a estrutura semântica sob perturbações.
• Análise de Camadas: A análise mostra que tokenizadores de alta qualidade facilitam o aprendizado de representação em todas as camadas do modelo de difusão. O alinhamento duplo (tokenizador + difusão) melhora ainda mais a qualidade das representações internas.
• Versatilidade: O método funciona bem com diferentes VFMs (SigLIP2, DINOv2, EVA-CLIP) e também em tarefas de geração texto-para-imagem (T2I), superando o VA-VAE em benchmarks como DPG-Bench e MJHQ-30K.

5. Significado e Impacto

O trabalho oferece evidências sólidas de que Modelos Fundacionais de Visão (VFMs) podem servir como tokenizadores visuais superiores quando suas representações nativas são preservadas (congeladas) em vez de distiladas.

• Paradigma de Treinamento: Muda o foco de "treinar um VAE para imitar um VFM" para "usar o VFM como base e aprender a reconstruir a partir dele".
• Eficiência Computacional: Acelera drasticamente o treinamento de modelos generativos de grande escala, reduzindo o tempo e custo computacional necessários para atingir alta fidelidade.
• Qualidade de Representação: Garante que o espaço latente seja semanticamente rico e robusto, o que é fundamental para a estabilidade e qualidade da geração em modelos de difusão modernos.

Em resumo, o VFM-VAE resolve o dilema entre fidelidade de reconstrução e riqueza semântica, estabelecendo uma nova base para o desenvolvimento de modelos generativos mais rápidos e de maior qualidade.