Resumo Técnico: Autoencoders de Representação Laminar para Difusão Eficiente (FlatDINO)

Definição do Problema

Avanços recentes em geração de imagens demonstraram que modelos de difusão podem operar diretamente sobre características de patches de aprendizado autossupervisionado (SSL) (ex: de DINOv2), em vez de latentes no espaço de pixels. Esta abordagem, exemplificada pelo RAE (Zheng et al., 2025), evita o trade-off entre reconstrução e geração inerente aos VAEs treinados em pixels, aproveitando a estrutura semanticamente significativa já codificada nas características de SSL.

No entanto, um gargalo computacional significativo permanece. Codificadores SSL padrão como o DINOv2 produzem grades bidimensionais densas de embeddings de patches de alta dimensão (ex: 256 patches de 768 dimensões para uma imagem de 256×256). Esta representação é comparável em tamanho ao sinal original, não oferecendo vantagem computacional sobre a difusão no espaço de pixels. A grade densa contém redundância espacial significativa porque vizinhos de patches compartilham conteúdo semântico substancial, tornando a difusão nessas características desnecessariamente cara em termos de FLOPs e memória.

Metodologia

Os autores propõem o FlatDINO, um autoencoder variacional (VAE) projetado para comprimir os embeddings de patches densos e bidimensionais do DINOv2 em uma sequência compacta unidimensional de tokens contínuos.

1. Arquitetura

• Codificador (Encoder): Um Vision Transformer (ViT-B) que recebe os 256 embeddings de patches de um DINOv2 ViT-B/14 congelado (com registros) como entrada. O codificador pré-pendiciona $T$ tokens de registro aprendíveis à sequência de patches. Após o processamento, apenas esses $T$ tokens de registro são retidos como a representação latente, descartando a estrutura espacial 2D original.
• Decodificador (Decoder): Um ViT-L que utiliza tokens de registro aprendíveis como queries para atender à sequência latente comprimida e reconstruir os 256 embeddings de patches originais do DINOv2.
• Objetivo de Treinamento: O modelo é treinado usando um objetivo de VAE padrão com uma perda de reconstrução de Erro Quadrático Médio (MSE) e uma penalidade de divergência KL. O peso da KL ($\beta$) é normalizado pela dimensionalidade latente total ($T \times d$) para garantir uma pressão de regularização consistente entre diferentes configurações.
• Geração de Latente: Um modelo generativo (especificamente um DiT-XL com flow matching) é treinado no espaço latente comprimido do FlatDINO. O modelo aprende um campo de velocidade para transportar amostras de uma priori Gaussiana padrão para a distribuição de dados ao longo de caminhos retos.

2. Escolhas de Design Chave

• Compressão 1D: Inspirado pelo TiTok (Yu et al., 2024a), o método comprime a grade 2D em uma sequência 1D. Os autores descobriram que 32 tokens fornecem um trade-off favorável entre fidelidade de reconstrução e comprimento de sequência.
• Redução da Dimensão de Características: Enquanto a contagem de tokens é reduzida de 256 para 32 (uma redução de 8×), a dimensão de características por token também é comprimida (ex: de 768 para 128) para reduzir ainda mais a dimensionalidade total.
• Flow Matching: O modelo generativo utiliza flow matching em vez de difusão padrão, operando diretamente na sequência latente 1D com embeddings posicionais aprendidos.

Principais Contribuições

1. FlatDINO: O primeiro método a comprimir características de patches autossupervisionadas em uma representação latente unidimensional, descartando efetivamente a estrutura espacial 2D original enquanto preserva o conteúdo semântico.
2. Ganhos de Eficiência: Demonstra que modelos de difusão treinados nos latentes FlatDINO alcançam qualidade de geração competitiva enquanto reduzem significativamente os custos computacionais. Especificamente, a redução de 8× no comprimento da sequência leva a uma redução de 8× em FLOPs por passagem forward e até uma redução de 4,5× em FLOPs por passo de treinamento em comparação com a difusão nas características não comprimidas do DINOv2.
3. Análise de Tokenização: Através de estudos de ablação de tokens, os autores mostram que o modelo aprende campos receptivos localizados e em forma de "blob" para a maioria dos tokens, particionando efetivamente a imagem espacialmente. Eles identificam uma "transição de fase" onde a redução de tokens abaixo de um certo limiar (ex: de 32 para 16) força o modelo a adotar estratégias de codificação menos eficientes (listras horizontais), levando a uma queda abrupta na qualidade.

Resultados

Experimentos foram conduzidos no ImageNet 256×256 usando uma arquitetura DiT-XL.

• Qualidade de Geração: Um DiT-XL treinado em latentes FlatDINO (32×128) alcança um gFID de 1.85 com classifier-free guidance (CFG) e 3.34 sem CFG. Isso é competitivo com o RAE (que alcança 1.41 gFID com CFG em características não comprimidas) e outros modelos de difusão latente de última geração, apesar da compressão significativa.
• Eficiência Computacional:
  • Inferência: O FlatDINO requer aproximadamente 8× menos FLOPs por passagem forward comparado a operar nos patches completos do DINOv2.
  • Treinamento: Devido ao comprimento de sequência reduzido e à inclusão do overhead do codificador, o FlatDINO requer 4,5× menos FLOPs por passo de treinamento para modelos DiT-XL em comparação com a baseline RAE.
• Reconstrução: A configuração 32×128 alcança um rFID de 0.77 (comparado a 0.62 para a baseline RAE não comprimida), indicando alta fidelidade na reconstrução do espaço de características.
• Ablação de Tokens: O estudo revela que 32 tokens são suficientes para manter campos receptivos localizados. Reduzir para 16 tokens causa uma mudança qualitativa para codificação de listras horizontais e uma queda significativa na qualidade de reconstrução.

Significância e Alegações

O artigo posiciona o FlatDINO como um passo em direção à síntese de imagem semanticamente fundamentada e eficiente. Os autores alegam que, ao explorar a redundância espacial inerente às grades de patches 2D, é possível alcançar uma redução de 8× na contagem de tokens sem sacrificar os benefícios de convergência de operar em um espaço latente semanticamente organizado.

O trabalho sugere que a difusão em representações semânticas comprimidas é uma abordagem prática para cenários com restrição de recursos. Os autores reconhecem que seus resultados atuais são preliminares e que o modelo não convergiu totalmente (treinado por 600 épocas vs. 800 para o RAE). Eles atribuem a lacuna restante na qualidade de geração à duração insuficiente de treinamento e à necessidade de receitas de difusão especificamente adaptadas para latentes semânticos comprimidos. Trabalhos futuros são propostos para investigar cronogramas de treinamento estendidos e autoencoders SSL otimizados para fechar essa lacuna.