Resumo Técnico: Correspondência de Distribuição de Representação para Geração Visual de Passo Único

Declaração do Problema

A modelagem generativa visa fundamentalmente igualar a distribuição de saída de um gerador à distribuição de dados. Embora modelos de difusão e de fluxo alcancem isso implicitamente ao aprender a reverter um processo de adição de ruído ao longo de muitos passos, abordagens recentes buscam fazer isso de forma explícita e direta em uma única avaliação de rede. Esses geradores de "passo único" visam igualar as distribuições de características geradas e reais sob codificadores pré-treinados congelados, sem depender de professores online, adversários ou trajetórias simuladas.

No entanto, os métodos de passo único existentes fixam suas escolhas de design conjuntamente, deixando obscuro quais componentes específicos impulsionam a qualidade. Além disso, objetivos clássicos de correspondência de distribuição, como a Máxima Discrepância de Média (MMD), foram anteriormente descartados como fracos demais para treinar geradores competitivos, e métricas de avaliação de codificador único são propensas ao "gaming" (manipulação), onde um gerador otimiza o escore de um codificador específico enquanto produz imagens visivelmente falsas.

Metodologia: Correspondência de Distribuição de Representação (RDM)

O artigo formaliza a Correspondência de Distribuição de Representação (RDM) como um paradigma definido por dois eixos de design: como as distribuições são comparadas e as representações nas quais são comparadas. Através de estudos controlados ao longo desses eixos, os autores propõem o iRDM (Improved RDM).

1. O Eixo de Comparação

• Reavaliando o MMD: Os autores demonstram que o MMD não é inerentemente fraco, mas foi previamente mal estimado. Eles propõem uma estratégia de estimativa específica:
  • Repulsão Exata Dentro do Lote (Within-Batch): O termo de repulsão (que evita o colapso de modo) é calculado exatamente sobre o lote gerado (soma de kernel $B \times B$).
  • Atração de Nyström: O termo de atração (puxando em direção aos dados) é estimado contra o conjunto de treinamento completo usando uma incorporação de média de kernel de Nyström. Em vez de reamostrar a referência a cada passo (o que introduz ruído), todo o conjunto de treinamento (ex: 1,28M de imagens do ImageNet) é comprimido uma única vez em um conjunto congelado de 4.096 marcos (landmarks). Isso fornece um resumo de variância zero da distribuição de dados.
• Tamanho do Lote (Batch Size): O tamanho do lote gerado foi identificado como uma variável crítica. Embora lotes maiores melhorem a estimativa, eles reduzem o número de passos do otimizador para um orçamento de computação fixo. Os autores encontram um ótimo acima de 2.048 (especificamente 5.120 para ImageNet e 10.240 para FLUX), excedendo em muito a prática padrão. O cache de gradiente é usado para gerenciar a memória.
• Tarefas Condicionais: Para texto-para-imagem, o método iguala a lei conjunta das características de legenda e imagem, em vez de apenas a marginal da imagem. Isso garante que a fidelidade do prompt faça parte do objetivo.

2. O Eixo de Representação

• O Problema dos Codificadores Únicos: Um único codificador pode ser "manipulado" (gamed). Um gerador pode reduzir a distância abaixo do escore dos dados reais naquele codificador específico enquanto permanece visivelmente falso (ex: corrigindo algumas classes, mas deixando artefatos em outras).
• Bateria de Codificadores Diversos: Para evitar a manipulação, o iRDM iguala-se contra um painel diversificado de 14 codificadores pré-treinados (incluindo DINOv2, SigLIP, CLIP, etc.).
• Otimização Restrita: Em vez de uma ponderação uniforme, os autores utilizam um controlador Lagrangiano proporcional (baseado em Stooke et al., 2020). Este controlador ajusta dinamicamente os pesos: ele aumenta o peso dos codificadores que estão "violando" o piso dos dados reais (os mais difíceis de satisfazer) e diminui o peso daqueles que já estão satisfeitos ou sofreram overfitting. Isso garante que o gerador não possa ignorar espaços de características difíceis.

3. Métrica de Avaliação: SWr14

Para avaliar o desempenho sem manipular a perda de treinamento, os autores introduzem o SWr14, uma distância de Wasserstein fatiada (Sliced-Wasserstein) média sobre os 14 codificadores.

• É independente da perda de treinamento (que usa MMD).
• Dados de validação reais são escalonados para um escore de 1,0 por construção.
• É mais difícil de manipular do que a distância de Fréchet ou o MMD.

Principais Contribuições

1. Um Modelo Unificador (RDM): O artigo formaliza a geração de passo único como um espaço de design governado por dois eixos (comparação e representação), permitindo que o teto de qualidade dos métodos seja rastreado até escolhas de design específicas, em vez de apenas ideias conceituais.
2. Uma Receita Simples no Estado da Arte (iRDM): Ao variar os eixos isoladamente, os autores identificam a configuração ideal: repulsão exata dentro do lote, atração de Nyström para uma referência de dados totais congelada, lotes grandes e novos, correspondência conjunta de imagem-texto para tarefas condicionais e otimização restrita para o equilíbrio de codificadores.
3. Uma Métrica Resistente à Manipulação: A introdução do SWr14, uma métrica de transporte ótimo sobre 14 codificadores que não compartilha maquinário com a perda de treinamento, fornecendo um padrão de avaliação robusto.

Resultados

Geração de ImageNet-256

• Escore SWr14: O iRDM atinge um SWr14 de estado da arte de 1,30, superando significativamente o melhor gerador de passo único anterior (pMF-H FD-SIM com 2,05) e aproximando-se do piso dos dados reais de 1,0.
• Preferência Humana: Avaliado via PickScore (um proxy de preferência humana aprendido que não foi usado no treinamento), o iRDM é preferido em relação ao melhor gerador de passo único anterior (FD-SIM) em 71,2% das amostras pareadas. Notavelmente, é o primeiro modelo de passo único a superar fotografias reais não utilizadas no treinamento em PickScore (taxa de vitória de 63,6%).

Pós-Treinamento de Texto-para-Imagem (FLUX.2)

• A receita foi aplicada para pós-treinar o FLUX.2 [klein] de quatro passos (4 bilhões de parâmetros) em um gerador de passo único.
• GenEval: O modelo iRDM de um passo supera o professor de quatro passos, alcançando 0,826 vs. 0,794.
• PickScore: O modelo de um passo alcança 22,76 vs. 22,58 para a versão de quatro passos.
• Eficiência: Estes resultados foram alcançados em aproximadamente 90 horas de GPU H200 sem um professor online ou modelo de recompensa.

Significância e Alegações

O artigo afirma que a Correspondência de Distribuição de Representação é um paradigma viável e livre de professor para geração de alta qualidade em um passo. A significância reside em:

• Desmistificar o MMD: Mostrar que o MMD, quando estimado corretamente (atração de Nyström + repulsão exata) com lotes grandes, é um objetivo poderoso.
• Prevenir o "Reward Hacking": Demonstrar que a correspondência de codificador único é insuficiente e que uma otimização equilibrada e restrita sobre uma bateria diversificada de codificadores é necessária para garantir realismo genuíno.
• Eficiência: Provar que um modelo de difusão de quatro passos pode ser destilado em um modelo de um passo que não apenas iguala, mas supera o professor em alinhamento (GenEval) e preferência humana (PickScore), utilizando um simples objetivo de correspondência de distribuição, sem a complexidade de treinamento adversário ou destilação de trajetória.

Os autores reconhecem modestamente que permanece uma lacuna entre o melhor gerador de passo único (SWr14 1,30) e uma nova amostra real (1,0), sugerindo trabalhos futuros em kernels multiescala e acoplamentos condicionais mais ricos.