JLT: Clean-Latent Prediction in Latent Diffusion Transformers

Resumo Técnico: JLT: Predição Latente Limpa em Transformers de Difusão Latente

Declaração do Problema

Os modelos de difusão para remoção de ruído revertem tradicionalmente um processo de corrupção, mas o alvo específico que a rede neural é treinada para prever varia entre arquiteturas. Alvos comuns incluem o ruído ($\epsilon$), a velocidade ($v = x - \epsilon$) ou os dados limpos ($x$). Embora esses alvos sejam algebricamente conversíveis para um tempo de corrupção fixo, pesquisas recentes (por exemplo, JiT) sugerem que, sob capacidade finita da rede e dados finitos, a escolha do alvo de previsão altera significativamente a dificuldade do problema de regressão. Especificamente, prever dados limpos permite que um modelo se concentre em variações estruturadas de baixa dimensão no manifold, em vez de reconstruir ruído ambiente fora do manifold.

A questão central abordada por este artigo é se essa vantagem de "previsão limpa" persiste quando o modelo opera em um espaço latente comprimido em vez do espaço de pixels bruto. Na difusão latente, um Autoencoder Variacional (VAE) já removeu grande parte da variabilidade de pixels brutos. Os autores investigam se as propriedades geométricas do alvo de previsão ($x$ vs. $v$) permanecem um fator crítico mesmo quando o espaço de representação é fixo e comprimido.

Metodologia

Os autores introduzem o JLT (Transformer Latente Conjunto), um quadro experimental controlado projetado para isolar o efeito do alvo de previsão dentro de um espaço latente fixo.

1. Arquitetura e Configuração:

   • Backbone: Um Transformer de escala Base (130M parâmetros) com 12 blocos, dimensão oculta 768 e 12 cabeças de atenção, espelhando a configuração do JiT-B/16.
   • Representação: O modelo opera em códigos latentes FLUX.2 VAE congelados. Isso garante que o espaço de representação seja idêntico em todas as variantes comparadas.
   • Treinamento: Os modelos são treinados por 250K passos (200 épocas) no ImageNet 256 $\times$ 256.
   • Variáveis Controladas: O estudo mantém estritamente constantes a representação (VAE FLUX.2), a escala da arquitetura, as configurações do otimizador, o tamanho do lote e os protocolos de amostragem.

2. Comparação Experimental:
   Os autores comparam duas variantes principais sob as mesmas configurações:

   • JLT (Predição Latente Limpa): A rede é supervisionada diretamente para prever o vetor latente limpo $x$.
   • DiT (Predição de Velocidade): Uma linha de base correspondente onde a rede é supervisionada para prever o vetor de velocidade $v = x - \epsilon$.
   • Variantes: Os experimentos são conduzidos em dois tamanhos de patch de grade VAE, denotados como /1 e /2 (por exemplo, JLT-B/1 vs. DiT-B/1).

3. Análise Teórica:
   O artigo fornece uma análise Gaussiana local para explicar a lacuna empírica. Assumindo uma aproximação linear-Gaussiana local da distribuição de dados latentes ($x \sim \mathcal{N}(0, \Sigma)$) e ruído independente ($\epsilon \sim \mathcal{N}(0, I)$):

   • Estrutura de Covariância: A covariância do alvo de velocidade é $Cov(y_v) = \Sigma + I$, enquanto o alvo limpo é $Cov(y_x) = \Sigma$. Isso implica que a previsão de velocidade adiciona um piso unitário isotrópico a cada direção. Em direções onde os dados limpos têm baixa variância (fracamente suportados pelo manifold), a previsão de velocidade força o modelo a aprender um alvo de variância unitária, efetivamente amplificando o ruído nessas direções.
   • Ambiguidade Condicional: A análise mostra que a variância residual de Bayes para a previsão de velocidade é escalada por um fator de $1/(1-t)^2$ em relação à previsão limpa. À medida que o autovalor $\lambda_i$ da covariância dos dados se aproxima de zero, o coeficiente para a previsão limpa tende a 0 (atenuando a direção), enquanto o coeficiente de velocidade tende a $-1/(1-t)$ (amplificando a direção).

Principais Resultados

Os experimentos no ImageNet 256 $\times$ 256 demonstram uma lacuna de desempenho consistente e significativa favorecendo a predição latente limpa:

• Desempenho Quantitativo:
  • JLT-B/1 (Limpo) alcança um FID-50K de 2,50 (com orientação sem classificador), comparado a 6,56 para o DiT-B/1 correspondente (Velocidade).
  • JLT-B/2 (Limpo) alcança um FID-50K de 14,81, comparado a 28,71 para o DiT-B/2.
  • A Pontuação de Inception (IS) também mostra uma melhoria substancial (por exemplo, 232,51 vs. 132,12 para as variantes /1).
• Dinâmica de Treinamento: As curvas de treinamento indicam que os modelos latentes limpos entram no regime de baixo FID mais cedo (por volta de 100K passos) e mantêm uma margem clara sobre os modelos de velocidade durante todo o treinamento.
• Robustez: A vantagem mantém-se em diferentes granularidades de tokenização (/1 e /2), sugerindo que o benefício não é um artefato de um tamanho de patch específico.

Significado e Alegações

O artigo afirma que a escolha do alvo de previsão na difusão latente é uma escolha de modelagem geométrica e não meramente uma parametrização algébrica.

1. Geometria do Alvo sobre Compressão: Os resultados refutam a noção de que a compressão latente sozinha explica a eficiência da difusão latente. A lacuna de desempenho existe dentro do mesmo espaço latente fixo, provando que a "geometria do alvo" (a distribuição do sinal supervisionado) é um fator distinto da "geometria de representação" (o tokenizador VAE).
2. Mecanismo de Melhoria: Os autores argumentam que a previsão limpa é superior porque amortece direções latentes de baixa variância (que são frequentemente ruído ou características de manifold fracamente suportadas), enquanto a previsão de velocidade amplifica essas direções ao herdar um piso de covariância isotrópico.
3. Escopo das Alegações: Os autores são modestos em suas alegações teóricas. Eles não afirmam que a previsão limpa é globalmente ótima para todos os tokenizadores, agendamentos de ruído ou amostradores. Em vez disso, fornecem um mecanismo Gaussiano local que é consistente com a lacuna empírica medida. O estudo é apresentado como uma ablação controlada para isolar o efeito do alvo, em vez de uma proposta para uma nova arquitetura de backbone ou uma substituição universal para todos os objetivos de difusão.

Em conclusão, o artigo estabelece que, mesmo em espaços latentes comprimidos, prever diretamente o estado latente limpo oferece uma vantagem geométrica sobre a previsão de velocidade, levando a uma qualidade de geração significativamente maior no ImageNet.