Resumo Técnico: Síntese Eficiente de Imagens com Sphere Latent Encoder

Declaração do Problema

A geração de imagens em poucos passos avançou significativamente por meio de métodos baseados em consistência e meanflow, no entanto, essas abordagens frequentemente sofrem com instabilidade no treinamento, sensibilidade a hiperparâmetros e conflitos inerentes nos objetivos de otimização. Uma alternativa, o Sphere Encoder [40], alcança geração de alta qualidade em poucos passos projetando representações latentes em uma hipersfera. Contudo, o framework do Sphere Encoder enfrenta duas limitações críticas:

1. Ineficiência Computacional: O processo de geração requer transições repetidas entre o espaço de pixels e o espaço latente (decodificação para uma imagem, seguida de recodificação) a cada passo de amostragem. Isso introduz uma sobrecarga computacional substancial tanto durante o treinamento quanto na inferência.
2. Conflito de Objetivos: O framework otimiza conjuntamente a reconstrução e a geração dentro de uma única arquitetura de codificador-decodificador. Pesquisas anteriores sugerem que essas tarefas favorecem representações diferentes; por exemplo, níveis de ruído mais baixos melhoram a reconstrução, mas degradam o desempenho generativo ao reduzir a cobertura da hipersfera. Essa tensão força o uso de arquiteturas maiores para manter a qualidade das amostras.

Metodologia

Os autores propõem o Sphere Latent Encoder, um framework que desacopla a reconstrução da geração para eliminar operações iterativas no espaço de pixels. A inovação central é realizar todo o processo generativo dentro de um espaço latente esférico fixo.

Arquitetura

O método consiste em três componentes principais:

1. Tokenizador de Imagem Fixo: Um autoencoder de representação pré-treinado (RAE) [43] é utilizado como codificador ($E$) e decodificador ($D$) fixos. O codificador, baseado no DINOv2, mapeia imagens para representações latentes, enquanto o decodificador reconstrói as imagens. Este componente é congelado durante o treinamento.
2. Modelo de Remoção de Ruído Latente: Uma rede de remoção de ruído baseada em transformadores separada ($G$, utilizando a arquitetura SiT [25]) opera exclusivamente no espaço latente. Ela recebe um latente esférico ruidoso como entrada e prevê o latente limpo.
3. Projeção Esférica: As representações latentes são projetadas em uma hipersfera usando uma função de esférica baseada em RMSNorm ($F$), garantindo uma estrutura latente uniforme.

Objetivo de Treinamento

Diferentemente do Sphere Encoder, que calcula perdas no espaço de pixels, o Sphere Latent Encoder desloca toda a supervisão para o espaço latente para melhorar a eficiência. O treinamento envolve dois latentes esféricos ruidosos, $v_{NOISY}$ e $v_{noisy}$, gerados adicionando ruído gaussiano em diferentes escalas.

• Perda de Reconstrução ($L_{recon}$): Combina uma perda $\ell_1$ e uma perda de similaridade cosseno entre o latente limpo previsto $G(v_{noisy})$ e o latente limpo ground truth $z$. Isso alinha o latente desruidado com a representação limpa.
• Perda de Consistência ($L_{cons}$): Impõe consistência entre previsões de diferentes níveis de ruído ($v_{NOISY}$ e $v_{noisy}$). A previsão do latente com menor ruído é tratada como um alvo fixo (via stop-gradient) para incentivar suavidade local na hipersfera.
• Remoção da Perda de Consistência Latente: Os autores removem explicitamente a perda de consistência latente usada no Sphere Encoder (que recodifica imagens decodificadas), constatando que ela adiciona sobrecarga computacional sem melhorar o desempenho em sua configuração desacoplada.

Procedimento de Amostragem

A amostragem ocorre inteiramente no espaço latente (Algoritmo 1):

1. Um latente é amostrado de uma distribuição gaussiana.
2. Ele é projetado na hipersfera.
3. O modelo de remoção de ruído $G$ refina o latente usando orientação livre de classificador.
4. O latente refinado é reprojetado, perturbado com ruído e usado como entrada para o próximo passo.
5. Esse refinamento iterativo continua por $T$ passos.
6. O decodificador é invocado apenas uma vez no final para mapear o latente final para pixels.

Os autores também propõem uma estratégia de amostragem de ruído modificada usando uma distribuição log-normal para adequar-se melhor à amostragem em poucos passos, diferindo da amostragem uniforme no Sphere Encoder.

Principais Contribuições

1. Framework Sphere Latent Encoder: Um framework generativo operando inteiramente no espaço latente esférico que elimina transições repetidas entre pixel e latente. Isso reduz o custo de inferência em aproximadamente 85% (cerca de 6,5× menos FLOPs) em comparação com o Sphere Encoder nos mesmos passos de amostragem.
2. Design Desacoplado: Ao aproveitar um autoencoder pré-treinado fixo e um modelo de remoção de ruído separado, o framework permite que a reconstrução e a geração se especializem independentemente, resolvendo o conflito de objetivos inerente à otimização conjunta.
3. Otimização Eficiente: O método evita objetivos de aproximação de primeira ordem (por exemplo, produtos de vetor-Jacobiano) exigidos por alguns métodos de fluxo em poucos passos, aprendendo em vez disso um processo direto de remoção de ruído latente em uma variedade esférica estruturada.

Resultados Experimentais

O método foi avaliado em Animal-Faces, Oxford-Flowers e ImageNet-1K (256×256).

• Desempenho vs. Sphere Encoder: Em Animal-Faces e Oxford-Flowers, o método proposto superou significativamente o Sphere Encoder tanto na qualidade de geração (FID) quanto no custo computacional (GFLOPs) em 2, 4 e 6 passos de amostragem. Por exemplo, em Animal-Faces com 4 passos, o FID melhorou de 18,23 (Sphere Encoder) para 6,89 (Nosso), enquanto os GFLOPs caíram de 4554 para 390.
• Resultados ImageNet-1K: Com 4 passos de amostragem (8 NFE com orientação), o método alcançou um FID de 2,25, uma melhoria substancial sobre os 4,02 do Sphere Encoder sob o mesmo orçamento. Também superou várias bases recentes de poucos passos (por exemplo, MeanFlow, $\alpha$-Flow) e aproximou-se do desempenho de modelos de difusão de múltiplos passos (por exemplo, DiT-XL/2) que exigem 250+ passos, mas com significativamente menos avaliações de função.
• Estudos de Ablação:
  • Distribuição de Ruído: A amostragem log-normal superou a amostragem uniforme.
  • Perdas: A combinação de perdas de reconstrução e consistência produziu os melhores resultados; adicionar a perda de consistência latente degradou o desempenho.
  • Autoencoder: A escolha do tokenizador foi crítica. O RAE [43] (com recursos DINOv2 e $d=768$) superou significativamente o FLUX-VAE e o GAE, destacando a necessidade de representações semânticas fortes e espaços latentes de alta dimensão.

Significado e Alegações

O artigo afirma que o Sphere Latent Encoder oferece um pipeline mais simples e eficiente para geração de imagens em poucos passos. Ao remover a dependência de refinamento iterativo no espaço de pixels, o método alcança:

• Maior Eficiência: Redução drástica em FLOPs e tempo de inferência em comparação com o Sphere Encoder.
• Melhor Qualidade: Qualidade de amostra aprimorada (FID mais baixo) devido ao desacoplamento da reconstrução e da geração, permitindo que o modelo de remoção de ruído foque exclusivamente na geração sem comprometer a fidelidade da reconstrução.
• Estabilidade: Um processo de treinamento que evita a instabilidade associada a produtos de vetor-Jacobiano e aos objetivos conflitantes da otimização conjunta.

Os autores posicionam este trabalho como um passo em direção à modelagem generativa escalável e eficiente que aproveita os benefícios de espaços latentes esféricos sem as penalidades computacionais das implementações anteriores. Eles reconhecem limitações, observando o foco atual na geração condicional a classes e a dependência da qualidade do autoencoder pré-treinado, o que pode restringir as capacidades de amostragem em um único passo.