Resumo Técnico: Aprendizado por Reforço Guiado por Discriminador para Flow Matching

Declaração do Problema

Modelos baseados em fluxo e de correspondência de escore (score-matching) tornaram-se o paradigma dominante para modelagem generativa devido aos seus objetivos de treinamento baseados em regressão e livres de simulação. No entanto, esses modelos são tipicamente treinados em um pipeline de múltiplos estágios: primeiro, um modelo base é ajustado a uma distribuição de dados $q$ usando fluxo ou correspondência de escore (FSM); segundo, o aprendizado por reforço (RL) é aplicado para inclinar o modelo em direção a regiões desejáveis definidas por uma recompensa $r(x)$.

O artigo identifica um descompasso estrutural neste paradigma. Embora as perdas de FSM minimizem o erro de regressão $\ell_2$ sobre campos de velocidade ou escore sob marginais de tempo de treinamento, elas frequentemente falham em recuperar propriedades intrínsecas da distribuição de dados — como realismo visual, estrutura de objeto coerente e movimento fisicamente correto — que estão teoricamente presentes nos dados. Consequentemente, os praticantes dependem de RL baseado em preferências para recuperar essas propriedades. Esta abordagem é problemática por duas razões:

1. Falha Estrutural: É intrigante que os modelos não consigam aprender propriedades do mundo real diretamente dos dados via perdas de correspondência, mas requeiram RL para recuperá-las.
2. Custo e Conflação: O treinamento de recompensas normalmente exige dados de preferência humana caros, o que conflaciona o realismo objetivo dos dados com inclinações subjetivas de anotadores.

Os autores argumentam que a falha dos objetivos de correspondência em capturar essas propriedades decorre de um "descompasso entre treino e teste" (onde os erros se acumulam ao longo da distribuição de rollout do modelo, em vez das marginais de treinamento) e um "descompasso geométrico" (onde pequenos erros no espaço de velocidade podem levar a grandes erros no espaço de qualidade da amostra/recompensa).

Metodologia: Aprendizado por Reforço Guiado por Discriminador (DRL)

O artigo propõe o Aprendizado por Reforço Guiado por Discriminador (DRL), um método para corrigir modelos de flow matching sem depender de dados de preferência humana. A ideia central é que a recompensa ótima para alinhar um modelo base $p_{base}$ com uma distribuição alvo $q$ é a razão de densidade de log $\log(q(x)/p_{base}(x))$. Como essa razão é intratável, o DRL a estima usando um discriminador.

O pipeline do DRL consiste em dois estágios:

Estágio 1: Estimativa de Recompensa no Espaço de Representação

Em vez de treinar um discriminador diretamente no espaço de pixels (o que é estatisticamente difícil e semanticamente não confiável), o DRL treina um discriminador em um espaço de representação de aprendizado auto-supervisionado (SSL) pré-treinado.

• Um codificador congelado $\phi$ (ex: DINOv2) mapeia imagens para um espaço de características $Z$.
• Um discriminador $D$ é treinado para distinguir entre amostras de dados reais $x \sim q$ e amostras do modelo base $x_{fake} \sim p_{base}$ usando suas representações codificadas $\phi(x)$.
• A função de recompensa é definida como o logito deste discriminador: $\hat{r}(x) = \log(D(\phi(x)) / (1 - D(\phi(x))))$.
• Justificativa Teórica: Ao restringir a recompensa ao espaço SSL, o DRL resolve um problema de otimização restrita: ele encontra a distribuição $p$ mais próxima de $p_{base}$ (em divergência KL) tal que o pushforward de $p$ através de $\phi$ coincida com o pushforward de $q$ através de $\phi$. Isso garante que as correções sejam feitas ao longo de eixos semanticamente significativos visíveis ao codificador.

Estágio 2: RL Regularizado por KL via Adjoint Matching

A recompensa estimada $\hat{r}$ é usada para ajustar o modelo base via RL regularizado por KL.

• O objetivo é maximizar a recompensa esperada enquanto minimiza a divergência KL entre as distribuições de trajetória do modelo ajustado e do modelo base.
• Os autores utilizam o Adjoint Matching, um algoritmo de última geração para modelos de fluxo que otimiza o objetivo ao nível da trajetória usando propagação de gradiente através do processo de amostragem.
• Para garantir a estabilidade do treinamento com o SDE sem memória necessário para o tilting imparcial do ponto final, os autores introduzem um Integrador Linear Local. Este integrador lida com a rigidez numérica perto de $t=0$ (onde os coeficientes de drift e difusão divergem) de forma mais eficaz do que as discretizações padrão de Euler-Maruyama com offsets.

Principais Contribuições

1. Análise Teórica das Limitações de Correspondência: O artigo fornece uma análise formal mostrando que as perdas de fluxo/correspondência de escore não garantem o controle sobre lacunas de recompensa (propriedades de interesse) devido a mudanças de distribuição entre treino e teste e ao desalinhamento entre a geometria do espaço de velocidade e a geometria de recompensa do espaço de amostra.
2. Aprendizado por Reforço Guiado por Discriminador (DRL): Um novo framework que utiliza um discriminador treinado em um espaço de representação SSL congelado para gerar um sinal de recompensa. Isso permite que o modelo recupere propriedades dos dados sem rótulos de preferência humana.
3. Integrador Linear Local: Um novo integrador numérico para o SDE sem memória usado no treinamento de RL, que evita a explosão de norma no início do tempo e melhora tanto a qualidade da amostragem quanto a estabilidade do treinamento.
4. Validação Empírica: Demonstração de que o DRL melhora significativamente o alinhamento distributivo e a qualidade da imagem em diversas arquiteturas (SiT, JiT, REPA, RAE) e espaços de características (DINOv2, DINOv3, SigLIP, Inception).

Resultados

Os autores avaliaram o DRL em quatro modelos de fluxo pré-treinados no ImageNet. As principais descobertas incluem:

• Alinhamento Distributivo: O DRL reduz significativamente a Distância de Fréchet (FD). Por exemplo, no SiT, a FD sem orientação no espaço DINOv3 caiu de 88.2 para 19.3. Em todos os modelos e espaços de características, o DRL melhorou consistentemente o alinhamento sem usar CFG (Classifier-Free Guidance).
• Qualidade de Imagem sem Dados de Preferência: O DRL melhorou as pontuações em modelos de recompensa de preferência humana mantidos (held-out) (ImageReward, HPSv2, PickScore, Aesthetics), apesar de nunca ter sido treinado neles. Isso sugere que o DRL recupera com sucesso as propriedades de "realismo" que as perdas de correspondência padrão perdem.
• Base Melhorada para Alinhamento de Preferência: Quando o DRL é usado como ponto de partida para o subsequente RL de preferência (PRL), ele produz uma fronteira de Pareto superior entre recompensa e fidelidade de imagem. O DRL absorve as correções estruturais recuperáveis dos dados, permitindo que o PRL foque em preferências subjetivas genuínas em vez de corrigir artefatos de baixo nível (ex: saturação excessiva, brilho excessivo).
• Lacuna de Destilação: Experimentos mostraram que um modelo estudante treinado via flow matching padrão em amostras de um professor DRL não conseguiu replicar o desempenho do professor, mesmo após ver 50 milhões de amostras. Isso confirma que o próprio objetivo de RL, e não apenas as amostras, é necessário para recuperar essas propriedades.

Significância e Alegações

O artigo alega que o pós-treinamento por RL para modelos baseados em fluxo deve ser visto não meramente como uma ferramenta para otimizar preferências externas, mas como um mecanismo complementar para recuperar a estrutura nos dados de treinamento que é imperfeitamente capturada pelos objetivos de correspondência baseados em padrões.

Ao usar um discriminador em um espaço de representação pré-treinado, o DRL fornece um sinal de recompensa que é:

1. Alinhado com as Propriedades dos Dados: Ele visa a distribuição de dados diretamente, em vez de preferências humanas subjetivas.
2. Geometricamente Sólido: Ele contorna as obstruções da geometria do espaço de velocidade ao otimizar diretamente no panorama de recompensa.
3. Eficiente: Requer significativamente menos computação do que o pré-treinamento e evita o custo de coletar dados de preferência.

Os autores concluem que os objetivos de correspondência e o RL não são métodos com objetivos conflitantes, mas sim duas ferramentas complementares para maximizar a utilidade dos dados de treinamento. O trabalho sugere que a "recompensa" para uma geração de alta qualidade está muitas vezes "nos dados o tempo todo", mas requer um mecanismo on-policy guiado por recompensa para ser totalmente realizado.