CrossFlow: One-Step Generation Across Latent and Pixel Spaces

Resumo Técnico: CrossFlow: Geração em Passo Único entre Espaços Latente e de Pixels

1. Definição do Problema

Os sistemas atuais de geração de imagem de última geração, particularmente os Modelos de Difusão Latente (LDMs) e os métodos de correspondência de fluxo (flow-matching), operam definindo o prior, o caminho de probabilidade e o alvo de predição dentro de um único espaço de representação. Embora os LDMs melhorem a eficiência computacional ao realizar a trajetória generativa em um espaço latente de autoencoder comprimido, eles dependem de um pipeline de duas etapas: denoising iterativo no espaço latente seguido por uma etapa separada de decodificação para mapear os latentes de volta para os pixels.

Esta arquitetura introduz um descompasso fundamental:

• Descompasso de Otimização: O gerador é otimizado para prever deslocamentos latentes, enquanto a qualidade final da imagem depende de um decoder treinado separadamente.
• Mudança de Distribuição (Distribution Shift): Na inferência, o decoder recebe latentes gerados que podem diferir distributionalmente dos latentes limpos produzidos pelo encoder durante o treinamento. Esse desvio pode amplificar erros.
• Limitação de Supervisão: O gerador não pode ser diretamente treinado com perdas perceptuais ou adversariais de nível de imagem, a menos que sua saída passe pelo decoder separado, enfraquecendo a supervisão direta no espaço de pixels.

O artigo questiona se o espaço de ruído e o espaço de saída devem ser idênticos, propondo um framework que preserva a eficiência das representações latentes enquanto prevê a própria imagem diretamente.

2. Metodologia: CrossFlow

O CrossFlow é uma formulação de fluxo cross-space que mapeia entradas latentes ruidosas diretamente para imagens no espaço de pixels em uma única avaliação de função. Ele resolve o descompasso de espaço derivando objetivos onde a velocidade no espaço latente é eliminada do alvo supervisionado.

Formulação Técnica Central:

• Framework Unificado: Os autores definem uma função alvo unificada $Z$ para paradigmas de geração de passo único. Métodos padrão dependem de campos de velocidade no espaço latente, o que impede a geração cross-space.
• Critérios de Design: Para permitir a supervisão cross-space, os autores estabelecem dois critérios:
  1. Eliminação de Dependências de Velocidade Latente: O objetivo deve ser estruturalmente desacoplado de velocidades no espaço latente. A saída da rede não deve ser um vetor de deslocamento latente.
  2. Condições de Contorno de Extremidade: O framework deve produzir uma imagem válida na fronteira de destino ($r=0$).
• O Objetivo CrossFlow: Ao satisfazer esses critérios, os autores derivam uma função de perda geral (Eq. 10) onde a rede $F_\theta$ prevê uma imagem no espaço de pixels $x_0$ em vez de uma velocidade latente.
  • A trajetória latente (definida por $\alpha(t)$ e $\beta(t)$) ainda dita como os latentes limpos são corrompidos para criar as entradas.
  • O alvo supervisionado é a imagem de pixels limpa.
  • O objetivo específico usado nos experimentos (Eq. 11) envolve um termo de derivada temporal computado via Produto Jacobiano-Vetor (JVP) ao longo da trajetória latente, garantindo que o modelo aprenda a dinâmica de fluxo correta sem prever velocidades latentes.

Detalhes de Implementação:

• Arquitetura: O modelo utiliza um encoder de imagem pré-treinado e fixo (VA-VAE) para mapear imagens para latentes e um gerador Vision Transformer (ViT) ($F_\theta$) para mapear latentes ruidosos diretamente para pixels.
• Pipeline de Treinamento: Uma abordagem de duas etapas é usada para garantir estabilidade:
  1. Treinar o encoder $E$ usando objetivos de autoencoder padrão (reconstrução + adversarial).
  2. Congelar $E$ e treinar o gerador $F_\theta$ do zero usando o objetivo CrossFlow.
• Função de Perda: O objetivo total combina a perda analítica CrossFlow com perdas auxiliares no espaço de pixels:
  • Perda de Reconstrução L1: Alinhamento pixel a pixel.
  • Perda Perceptual: Similaridade estrutural.
  • Perda Adversarial (GAN): Fidelidade de alta frequência.
  • Nota: O fator de ponderação analítico do objetivo de fluxo é omitido durante o cálculo da perda perceptual para evitar instabilidade no treinamento.

3. Principais Contribuições

1. Framework Cross-Space: Introdução do CrossFlow, um framework de geração de passo único onde o prior de entrada (latente) e a imagem de saída (pixel) residem em diferentes espaços de representação.
2. Objetivo de Passo Único Livre de Velocidade: Uma derivação inovadora que permite que um modelo atue tanto como um gerador de um passo (latente-para-pixel) quanto como um substituto de decoder, eliminando a necessidade de um decoder separado na inferência.
3. Supervisão Direta no Espaço de Pixels: A capacidade de treinar o gerador diretamente com perdas perceptuais e adversariais no espaço de pixels, contornando a etapa intermediária do decoder.
4. Validação Empírica: Demonstração de que um único modelo pode alcançar pontuações FID competitivas no ImageNet-1k (256×256) com uma única avaliação de função (1-NFE), superando métodos latentes de passo único padrão e igualando baselines de difusão latente de múltiplos passos.

4. Resultados Experimentais

Experimentos foram conduzidos no ImageNet-1k na resolução de 256×256 usando geração condicionada por classe.

• Desempenho Principal:

  • CrossFlow-XL (461M parâmetros) alcança um FID de 1.62 com 1-NFE.
  • Isso é competitivo com baselines de difusão latente de múltiplos passos como DiT-XL/2 (2.27 FID com 250 passos) e SiT-XL/2 (2.06 FID com 250 passos).
  • Supera significativamente outros métodos latentes de passo único (ex: iMF-XL/2 com 1.72 FID) e evita a necessidade de um decoder separado.
  • CrossFlow-B (95M) e CrossFlow-L (315M) alcançam FIDs de 2.54 e 1.69, respectivamente, demonstrando escalabilidade.

• Estudos de Ablação:

  • Escolha do Encoder: O encoder VA-VAE (32 canais) rendeu os melhores resultados (2.54 FID), superando o SD-VAE (4.08 FID) e o FLUX.2 (3.81 FID), sugerindo que a estrutura da representação latente é crítica.
  • Componentes de Perda: Remover a perda perceptual causou uma queda severa no FID para 30.21, e remover a perda GAN reduziu para 8.86. O objetivo de fluxo CrossFlow fornece o sinal estrutural, enquanto as perdas perceptuais e adversariais são vitais para a fidelidade de alta frequência.
  • Substituição de Decoder: Quando usado como um decoder para latentes gerados por um LightningDiT-B/1, o CrossFlow consistentemente melhorou o FID sobre o decoder VA-VAE padrão durante todo o treinamento.

5. Significância e Alegações

O artigo afirma que o CrossFlow consegue unir a lacuna entre a eficiência das representações latentes e a fidelidade da supervisão direta no espaço de pixels. Ao unificar o gerador e o decoder em um único modelo treinado em latentes ruidosos e alvos no espaço de pixels, o CrossFlow:

• Remove a etapa explícita de decoder na inferência.
• Expõe o modelo durante o treinamento aos latentes ruidosos/imperfeitos que ocorrem durante a geração, em vez de apenas saídas de encoder limpas.
• Demonstra que objetivos de fluxo cross-space podem combinar os benefícios computacionais da difusão latente com as melhorias de qualidade da supervisão adversarial e perceptual direta no nível de pixel.

Os autores mantêm um tom modesto, observando que, embora a representação latente e as perdas auxiliares sejam importantes, a principal contribuição é a mudança arquitetônica e de objetivo que permite a geração cross-space de passo único sem exigir módulos de decoder separados.