Resumo Técnico: Destilação de Transformers com Drift em Autoencoders de Representação

Declaração do Problema

Embora os Autoencoders de Representação (RAEs) tenham melhorado significativamente o desempenho de modelos generativos baseados em fluxo ao utilizar espaços latentes semanticamente ricos e agrupados por rótulos, derivados de características DINO pré-treinadas, destilar esses modelos em geradores eficientes de um ou poucos passos permanece um desafio.

Os autores identificam um desajuste geométrico fundamental: os espaços latentes de RAE exibem anisotropia severa e grandes curvaturas em suas trajetórias subjacentes de Equações Diferenciais Ordinárias (ODEs) em comparação com os espaços tradicionais de Autoencoders Variacionais (VAEs). Os métodos de destilação baseados em trajetória existentes dependem da suposição de trajetórias do professor (teacher) suaves e quase retas. As trajetórias altamente curvas nos espaços de RAE fazem com que esses métodos se tornem instáveis ou ineficientes. Além disso, embora o recentemente proposto Modelo de Drifting ofereça uma abordagem de nível de distribuição que evita o correspondência de trajetórias curvas, sua formulação original requer um Autoencoder de Máscara (MAE) auxiliar, treinado separadamente, para fornecer características concentradas, introduzindo uma sobrecarga computacional significativa.

Metodologia

O artigo propõe o Drift-RAE, um framework de destilação que adapta o paradigma do Modelo de Drifting especificamente para espaços latentes de RAE, eliminando a necessidade de extratores de características auxiliares.

1. Análise Teórica de Geometria e Dinâmica

Os autores primeiro analisam quantitativamente os espaços latentes, demonstrando que as características de RAE são substancialmente mais anisotrópicas (menor Razão de Participação e Entropia Espectral) do que as características de SD-VAE. Eles provam teoricamente (Teorema 1) que, em espaços de alta dimensão, se as amostras positivas estiverem excessivamente dispersas, o componente de atração do campo de drift tende a desaparecer, levando a um campo degenerado.

• Insight Chave: Enquanto o Modelo de Drifting original requer um MAE para concentrar características dispersas, os espaços latentes de RAE são naturalmente concentrados semanticamente. Isso permite que o Modelo de Drifting opere diretamente nos latentes de RAE sem o módulo MAE auxiliar.

2. O Framework Drift-RAE

O método destila um modelo de fluxo pré-treinado $v_\theta$ em um gerador de um passo $G_\theta(z) = z - v_\theta(z, 1)$. Em vez de corresponder às trajetórias de ODE, o gerador é treinado para minimizar um objetivo de drifting que empurra as amostras geradas em direção às distribuições de dados reais e para longe de outras amostras geradas.

• Objetivo: A perda é computada por token nas representações latentes de RAE, utilizando um campo de drifting normalizado $\tilde{V}$ derivado de amostras positivas (reais) e negativas (geradas).

3. Modificações Práticas para Estabilidade

Para alinhar o Modelo de Drifting com o framework Diffusion-GAN e melhorar a estabilidade do treinamento, os autores introduzem três modificações específicas:

• Dimensão Softmax: Os autores removem o truque de softmax bidirecional usado no Modelo de Drifting original. Eles mantêm um softmax único sobre os índices das amostras, o que se alinha melhor com a direção do gradiente do potencial log-sum-exp subjacente.
• Perturbação de Entrada: Para evitar o desaparecimento do gradiente causado por variedades (manifolds) não interceptantes entre dados reais e gerados, as amostras geradas são perturbadas com ruído amostrado de uma distribuição de Laplace antes de computar o campo de drifting.
• Desprendimento Parcial de Amostras Negativas: Para aumentar o número de amostras negativas usadas para a aproximação da distribuição sem exceder os orçamentos de memória da GPU, os autores propõem o uso de um pool maior de amostras negativas ($N_{total\_neg}$), onde apenas um subconjunto ($N_{neg}$) carrega gradientes, enquanto o restante é desprendido (detached). Isso melhora a aproximação empírica da distribuição gerada.

Principais Contribuições

1. Análise Geométrica: Uma demonstração quantitativa de que os espaços latentes de RAE possuem alta curvatura e anisotropia, tornando a destilação baseada em trajetória padrão instável.
2. Justificativa Teórica: Uma prova mostrando que os Modelos de Drifting falham em espaços excessivamente dispersos (necessitando de um MAE na formulação original), mas prosperam nos espaços semânticos naturalmente concentrados dos RAEs.
3. Drift-RAE: Um novo método de destilação que alcança geração de um passo de estado da arte em espaços de RAE sem exigir um extrator de características MAE auxiliar.
4. Melhorias de Estabilidade: Insights teóricos e empíricos sobre a modificação do objetivo de Drifting (dimensão softmax, perturbação de entrada e desprendimento de amostras) para alinhar com os princípios de Diffusion-GAN e melhorar a convergência.

Resultados Experimentais

O método foi avaliado no dataset ImageNet 256 $\times$ 256 usando backbones DiTDH-XL e DiTDH-L.

• Desempenho: O Drift-RAE alcança um FID de 1.77 (com DiTDH-XL) e 2.12 (com DiTDH-L) usando um único passo de geração.
• Eficiência: O modelo atinge esses resultados em apenas 10.000 passos de destilação (aprox. 16 épocas).
• Comparação:
  • Supera métodos de destilação anteriores em espaços de RAE (ex: MF-RAE).
  • Alcança um FID comparável e um FDDINOv2 melhorado em relação ao Modelo de Drifting original, apesar de não utilizar o custoso MAE auxiliar.
  • Permanece competitivo com geradores de um e poucos passos treinados em outros espaços latentes (ex: Espaço de Pixels, SD-VAE).

Significância e Alegações

O artigo afirma que o Drifting fornece um framework eficaz e promissor para destilar modelos generativos em espaços de representação. A principal significância reside em resolver a instabilidade da destilação de modelos baseados em RAE, aproveitando sua concentração semântica inerente, eliminando assim a dependência de extratores de características auxiliares.

Os autores notam modestamente as limitações, reconhecendo que sua análise teórica depende de modelos de alta dimensão simplificados e que treinar Modelos de Drifting do zero (sem destilação) em espaços de RAE permanece um problema em aberto. Eles também destacam que o requisito de muitas amostras positivas da mesma classe pode apresentar desafios de escalabilidade para geração de texto-para-imagem, sugerindo isso como uma direção para pesquisas futuras.