Resumo Técnico: LeapAlign

Declaração do Problema

O artigo aborda o desafio de alinhar modelos de correspondência de fluxo (flow matching) com preferências humanas por meio de pós-treinamento. Embora métodos de gradiente direto ofereçam uma alternativa promissora às abordagens baseadas em gradiente de política (como GRPO), ao aproveitar a natureza diferenciável da amostragem de correspondência de fluxo, eles enfrentam obstáculos significativos quando aplicados a trajetórias longas de geração. Especificamente, a retropropagação de gradientes de recompensa através de toda a sequência de etapas de amostragem de EDO (Equações Diferenciais Ordinárias) incorre em custos de memória proibitivos e riscos de explosão de gradiente. Consequentemente, os métodos existentes de gradiente direto normalmente restringem as atualizações a um único passo de tempo próximo à imagem final ou descartam gradientes aninhados para evitar instabilidade. Essa limitação impede a otimização dos passos iniciais de geração, que são críticos para determinar o layout global e a composição estrutural da imagem final.

Metodologia: LeapAlign

Os autores propõem o LeapAlign, um método de ajuste fino projetado para permitir a propagação direta e estável de gradientes para qualquer etapa de geração, incluindo passos de tempo iniciais, sem incorrer em sobrecarga excessiva de memória. A metodologia central envolve três componentes principais:

1. Construção de Trajetória de Salto em Duas Etapas

Em vez de retropropagar através da trajetória completa de geração, o LeapAlign constrói uma trajetória de salto encurtada consistindo apenas de duas etapas derivadas de uma trajetória completa padrão.

• Processo: Durante cada iteração, uma trajetória completa é amostrada de ruído ($x_1$) até uma imagem limpa ($x_0$). Dois passos de tempo, $k$ e $j$ (onde $k > j$), são selecionados aleatoriamente.
• Saltos: O método utiliza a propriedade de previsão de salto de uma etapa dos modelos de fluxo retificado para estimar latentes futuros em um único passo:
  • Salto 1: Prevê o latent $x_j$ a partir de $x_k$ usando o campo de velocidade em $x_k$: $\hat{x}_{j|k} = x_k - (k-j)v_\theta(x_k)$.
  • Salto 2: Prevê o latent final $x_0$ a partir de $x_j$ usando o campo de velocidade em $x_j$: $\hat{x}_{0|j} = x_j - j v_\theta(x_j)$.
• Conectores Latentes: Para manter a diferenciabilidade ao alinhar estados previstos com estados reais amostrados da trajetória completa, "conectores latentes" são introduzidos usando operações stop_gradient (por exemplo, $x_j = \hat{x}_{j|k} + \text{stop\_gradient}(x_j - \hat{x}_{j|k})$).
• Benefício: Isso reduz o caminho de retropropagação para duas etapas, mantendo os custos de memória constantes independentemente do comprimento original da trajetória, enquanto permite a seleção aleatória de $(k, j)$ para atualizar qualquer parte do processo de geração.

2. Desconto de Gradiente

Para abordar a explosão de gradiente enquanto preserva sinais de aprendizado úteis, o LeapAlign introduz um mecanismo de desconto de gradiente.

• O Problema: A retropropagação através de duas etapas introduz um termo de "gradiente aninhado" ($\frac{\partial v_\theta(x_j)}{\partial x_j} \frac{\partial v_\theta(x_k)}{\partial \theta}$) que captura interações entre passos de tempo. Métodos anteriores como DRTune eliminam esse termo inteiramente ao parar gradientes na entrada do modelo, o que descarta informações valiosas de dependência.
• A Solução: O LeapAlign escala o termo de gradiente aninhado por um fator $\alpha \in [0, 1]$ em vez de removê-lo. Isso é alcançado modificando a entrada para a previsão de velocidade no segundo salto: $\hat{x}_{0|j} = x_j - j v_\theta(\alpha x_j + (1-\alpha)\text{stop\_gradient}(x_j))$.
• Efeito: Isso reduz a magnitude de grandes termos de gradiente para estabilizar o treinamento, mantendo a estrutura completa do gradiente, incluindo o gradiente aninhado.

3. Ponderação por Similaridade de Trajetória

Para garantir que as trajetórias de salto encurtadas forneçam supervisão confiável, o método atribui pesos maiores a saltos que se assemelham estreitamente ao caminho de geração verdadeiro.

• Mecanismo: A perda é ponderada pelo inverso da diferença absoluta média entre os estados previstos ($\hat{x}$) e os estados reais ($x$) nos pontos de conexão ($x_j$ e $x_0$).
• Fórmula: $w_{sim} = \frac{1}{\max(d_j, \tau) + \max(d_0, \tau)}$, onde $d$ representa a distância e $\tau$ é um limite mínimo para evitar ênfase excessiva em pares quase idênticos.
• Objetivo: Isso garante que saltos que se desviam significativamente das dinâmicas verdadeiras contribuam menos para a atualização do gradiente, melhorando a estabilidade e a fidelidade da supervisão.

4. Objetivo de Ajuste Fino

O método emprega uma perda do tipo "hinge" para maximizar recompensas enquanto previne a exploração de recompensas (reward hacking): $L_{raw} = \max(0, \lambda - r(x_0))$. Crucialmente, ao contrário de métodos que estimam recompensas a partir de previsões de salto, o LeapAlign avalia a recompensa na imagem final real ($x_0$) gerada pela trajetória completa, garantindo que o modelo de recompensa avalie a verdadeira qualidade visual e semântica.

Contribuições Principais

1. Estrutura LeapAlign: Um método inovador de pós-treinamento que constrói trajetórias de salto de duas etapas para permitir a retropropagação direta de gradientes para etapas iniciais de geração com custo de memória constante.
2. Técnicas de Estabilidade de Gradiente: A introdução de desconto de gradiente para escalar (em vez de remover) gradientes aninhados, e ponderação por similaridade de trajetória para priorizar saltos consistentes, abordando os problemas de instabilidade e perda de sinal em métodos anteriores de gradiente direto.
3. Avaliação Abrangente: Demonstração de que o LeapAlign supera os métodos de gradiente de política de última geração (baseados em GRPO) e métodos de gradiente direto em múltiplas métricas, incluindo pontuações de preferência humana (HPSv2.1, HPSv3, PickScore) e benchmarks de alinhamento composicional (GenEval).

Resultados

Os autores ajustaram finamente o modelo Flux.1-Dev usando LeapAlign e o avaliaram contra baselines incluindo DanceGRPO, MixGRPO, ReFL, DRaFT-LV e DRTune.

• Alinhamento de Preferência Geral: O LeapAlign alcançou as maiores pontuações em métricas dentro do domínio (HPSv2.1) e fora do domínio (HPSv3, PickScore, UnifiedReward). Notavelmente, superou o MixGRPO (que usa três modelos de recompensa) enquanto foi treinado em um único modelo de recompensa (HPSv2.1).
• Alinhamento Composicional: No benchmark GenEval, o LeapAlign mostrou ganhos significativos, particularmente nas tarefas "dois objetos", "cores", "posição" e "ligação de atributos". A melhoria nessas categorias destaca o sucesso do método em otimizar etapas iniciais de geração que ditam o layout global.
• Análise Qualitativa: Comparações visuais demonstraram que o LeapAlign modifica substancialmente estruturas globais de imagem para alinhar melhor com prompts de texto, enquanto métodos restritos a atualizações de etapas tardias (como ReFL) frequentemente mantinham o layout do modelo pré-treinado.
• Estudos de Ablação: Experimentos confirmaram que o design de salto de duas etapas oferece o melhor compromisso entre desempenho e memória. O desconto de gradiente com $\alpha=0.3$ produziu resultados ótimos, e a seleção aleatória de passos de tempo em todo o intervalo $[0, 1]$ foi superior a intervalos fixos, sublinhando a importância das atualizações de etapas iniciais.

Significado e Alegações

O artigo afirma que o LeapAlign resolve com sucesso o compromisso entre eficiência de memória e cobertura de gradiente no pós-treinamento de correspondência de fluxo. Ao permitir atualizações em etapas iniciais de geração sem custos computacionais proibitivos, ele permite a otimização da estrutura global da imagem, o que é crucial para tarefas composicionais complexas. Os autores posicionam o LeapAlign como uma alternativa estável e eficaz tanto aos métodos estocásticos de gradiente de política quanto às abordagens de gradiente direto instáveis, alcançando alinhamento imagem-texto superior e qualidade de geração. O trabalho sugere que reter e escalar gradientes aninhados, em vez de descartá-los, é um fator chave para melhorar o desempenho do alinhamento por gradiente direto. Trabalhos futuros são notados para estender esses princípios à geração de vídeo.