Resumo Técnico: Alinhamento da Geometria Latente para Fluxo Esférico em Geração de Imagens

Declaração do Problema

O fluxo latente (e os modelos de difusão) para geração de imagens opera tipicamente transportando ruído gaussiano para latentes de Autoencoders Variacionais (VAE) ao longo de caminhos lineares no espaço euclidiano. No entanto, essa abordagem ignora duas propriedades estruturais críticas dos latentes de VAE de alta dimensão:

1. Concentração de Medida: Tanto o prior gaussiano quanto os latentes de dados codificados concentram-se em cascas esféricas finas, em vez de preencherem o volume do espaço. Uma linha reta euclidiana (corda) conectando dois pontos nessas cascas passa por raios internos que nenhuma das distribuições ocupa, forçando o modelo a aprender transições através de regiões do espaço latente raramente populadas.
2. Dominância Direcional: A saída do decodificador é predominantemente sensível à direção (componente angular) de um token latente, e não ao seu comprimento (componente radial). Experimentos que trocam o raio e a direção de tokens de vizinhos da mesma classe demonstram que preservar a direção enquanto se altera o raio resulta em mudança perceptual mínima, enquanto alterar a direção causa deslocamentos semânticos significativos.

O fluxo linear padrão ignora esses fatos, alocando uma porção significativa do alvo de velocidade de treinamento (até 50-90%, dependendo do tokenizador) para movimento radial. Como o decodificador é menos sensível ao raio, essa supervisão é ineficiente, levando a uma convergência mais lenta e qualidade de imagem subótima (FID) em comparação com métodos que respeitam a geometria subjacente.

Metodologia

Os autores propõem o Fluxo Esférico, um método que alinha a geometria latente com a estrutura esférica dos dados sem modificar a arquitetura de difusão ou exigir codificadores auxiliares. A abordagem consiste em quatro componentes-chave:

1. Projeção Esférica Token a Token:
   O método projeta cada token latente $z$ de um VAE pré-treinado sobre uma hipersfera de raio fixo $S^{d-1}(\sqrt{d})$. Isso é alcançado por meio de uma etapa de normalização $L_2$ inserida entre o codificador congelado e o decodificador:
   $$z_{i,j} \leftarrow \sqrt{d} \cdot \frac{z_{i,j}}{\|z_{i,j}\|}$$
   Isso garante que todos os tokens residam em uma esfera com um raio correspondente ao raio de concentração do prior gaussiano padrão.

2. Prior Esférico:
   Em vez de usar ruído gaussiano bruto, o ponto final do ruído $z_0$ é construído amostrando ruído gaussiano isotrópico $\epsilon \sim \mathcal{N}(0, I)$ e projetando-o sobre a esfera:
   $$z_0 = \sqrt{d} \cdot \frac{\epsilon}{\|\epsilon\|_2}$$
   Isso resulta em uma distribuição uniforme na esfera, correspondendo à geometria dos dados projetados.

3. Ajuste Fino do Decodificador:
   Para acomodar a restrição geométrica, o decodificador é ajustado finamente por algumas épocas (5 épocas nos experimentos), enquanto o codificador e a camada de projeção permanecem congelados. O objetivo de reconstrução mantém as perdas de pixel padrão (L1, LPIPS, adversarial), mas descarta o termo KL, pois o latente projetado é determinístico.

4. Transporte Slerp e Velocidade Tangente:
   O caminho de transporte entre os pontos finais de ruído e dados é substituído por Interpolação Linear Esférica (slerp), que segue o arco geodésico na esfera.

   • Caminho Geodésico: O caminho $z_t$ permanece na esfera para todo $t \in [0, 1]$.
   • Alvo de Velocidade: A velocidade condicional $u_t$ é a derivada temporal da geodésica, projetada no espaço tangente $T_{z_t}S^{d-1}$. Isso garante que a velocidade seja puramente angular, com componente radial zero por construção.
   • Treinamento e Amostragem: O modelo é treinado para prever essa velocidade tangente. Durante a inferência, o amostrador usa o mapa exponencial para integrar a velocidade, garantindo que a trajetória permaneça na esfera sem exigir correções de projeção radial a cada passo.

Contribuições Principais

O artigo identifica e quantifica um "desajuste casca-radial" no fluxo latente padrão e introduz uma intervenção geométrica para resolvê-lo:

1. Identificação do Desajuste: Os autores demonstram que caminhos lineares forçam o modelo a prever movimento radial para o qual o decodificador é insensível, desperdiçando capacidade e retardando a convergência.
2. Projeção Token a Token: Eles introduzem um método para restringir latentes de VAE pré-treinados a uma esfera de raio fixo usando apenas o ajuste fino do decodificador, evitando a necessidade de retreinar o codificador ou usar codificadores de representação auxiliares.
3. Fluxo Baseado em Slerp: Eles propõem treinar o fluxo ao longo de geodésicas slerp com alvos de velocidade tangente, garantindo que toda a trajetória respeite a variedade esférica.
4. Validação Empírica: O método demonstra melhorar o desempenho em várias famílias de tokenizadores (FLUX.2, VA-VAE, REPA-E FLUX.1) e escalas de modelo (SiT-B a SiT-XL) sem alterações arquiteturais.

Resultados

Experimentos na geração condicional por classe de ImageNet-256 demonstram melhorias consistentes:

• Melhorias no FID: O método esférico-slerp alcança pontuações FID mais baixas do que o fluxo linear padrão em todas as configurações testadas. Por exemplo, no SiT-B/2 com FLUX.2, o FID cai de 26,35 (Linear) para 20,55 (Sférico-slerp).
• Eficiência de Treinamento: O método atinge um FID alvo de 30 em aproximadamente 2,2$\times$ menos passos de treinamento em comparação com a linha de base linear padrão.
• Generalização: Ganhos são observados em diferentes escalas de orientação (CFG) e tamanhos de modelo (B e XL), indicando que o benefício é inerente à geometria latente e não ao ajuste específico de hiperparâmetros.
• Estudos de Ablação:
  • Trocar decodificadores entre pipelines padrão e esféricos degrada o desempenho em ambas as direções, confirmando que o fluxo aprendido está fortemente acoplado à geometria latente específica.
  • Usar um caminho "decomposto em cascas" (interpolando raio e direção separadamente) em latentes padrão não produz os mesmos ganhos, confirmando que remover o grau de liberdade radial inteiramente é crucial.
  • A qualidade de reconstrução (rFID) dos decodificadores esféricos ajustados finamente permanece competitiva, mostrando que a restrição geométrica não compromete severamente a capacidade do VAE de reconstruir imagens.

Significado

O artigo estabelece a geometria de suporte latente como um eixo de design crítico para modelagem generativa. Ele argumenta que a forma da distribuição latente define o espaço que o gerador deve percorrer, e ignorar essa geometria (assumindo linearidade euclidiana) leva à ineficiência.

O significado deste trabalho reside em sua capacidade de melhorar a qualidade de geração de imagens e a velocidade de treinamento sem:

• Alterar a arquitetura de difusão (por exemplo, a espinha dorsal SiT).
• Adicionar codificadores auxiliares (ao contrário do fluxo riemanniano em recursos DINOv2 congelados).
• Retreinar o codificador VAE do zero.

Ao simplesmente projetar latentes pré-treinados existentes sobre uma esfera e alinhar o caminho do fluxo às geodésicas resultantes, o método alcança resultados state-of-the-art com sobrecarga computacional mínima. Os autores sugerem que trabalhos futuros devem explorar essas restrições geométricas sob regimes de compressão mais fortes e projetar tokenizadores que correspondam nativamente a essa estrutura esférica.