Resumo Técnico: Correspondência Adjoint Eficiente para Ajuste Fino de Modelos de Difusão

Declaração do Problema

O ajuste fino baseado em recompensa tornou-se a abordagem padrão para alinhar modelos de difusão e fluxo pré-treinados com preferências humanas na geração de texto para imagem (T2I). Entre estes, a Correspondência Adjoint (AM) oferece uma formulação principista ao formular o ajuste fino por recompensa como um problema de Controle Ótimo Estocástico (SOC). A AM visa transformar um modelo pré-treinado em uma distribuição alvo $p^\star_1(x) \propto e^{\beta r(x)}p_{data}(x)$, onde $r(x)$ é uma função de recompensa.

No entanto, a AM sofre de ineficiências computacionais significativas que impedem sua aplicação prática:

1. Gargalo de Simulação Forward: A AM requer a simulação de trajetórias generativas completas usando um solver de equação diferencial estocástica (SDE) sob dinâmicas "sem memória". Isso necessita de um grande número de avaliações de função (NFEs) e exige o armazenamento de toda a trajetória na memória.
2. Gargalo de Simulação Backward: Para calcular o sinal de controle, a AM requer uma simulação backward de Equação Diferencial Ordinária (ODE) do estado adjoint ao longo da trajetória amostrada. Isso envolve o cálculo de Produtos Vetor-Jacobiano (JVPs) através da rede pré-treinada em cada passo de integração, o que é computacionalmente caro.

Os autores observam que esses gargalos estão intrinsecamente ligados ao desvio base não trivial herdado do modelo pré-treinado nas formulações padrão de AM.

Metodologia: Correspondência Adjoint Eficiente (EAM)

O artigo propõe a Correspondência Adjoint Eficiente (EAM), que reformula o problema de SOC para eliminar ambos os gargalos computacionais, preservando a garantia teórica de alcançar a distribuição alvo inclinada pela recompensa.

1. Redesenho das Dinâmicas Base

A EAM introduz duas condições estruturais para as dinâmicas base:

• Linearidade: O desvio base $b(x,t)$ é redesenhado para ser linear em $x$ (isto é, $b(x,t) = D(t)x$).
• Sem Memória: A dinâmica base permanece sem memória, permitindo que estados intermediários sejam amostrados diretamente a partir do ponto final.

Ao fixar o coeficiente de difusão $\sigma(t)$ para corresponder à configuração padrão, os autores derivam uma família específica de desvios lineares admissíveis (Eq. 15) que satisfazem essas condições. Esse desvio linear permite uma solução de forma fechada para o estado adjoint, eliminando a necessidade de simulação backward de ODE e avaliações de JVP.

2. Correção de Custo Terminal

Alterar o desvio base altera a distribuição terminal do sistema não controlado. Para garantir que a distribuição controlada ótima ainda corresponda à distribuição alvo desejada inclinada pela recompensa $p^\star_1$, a EAM deriva um custo terminal corrigido $g(x)$. O gradiente desse custo, $\nabla g(x)$, inclui termos para a pontuação pré-treinada, a nova pontuação terminal base e o gradiente da recompensa (Eq. 16).

3. Algoritmo de Treinamento

O loop de treinamento da EAM (Algoritmo 1) procede da seguinte forma:

1. Geração de Ponto Final: Em vez de um rollout completo de SDE, a imagem de ponto final $X_1$ é gerada usando um solver determinístico de ODE de poucos passos (por exemplo, 10 NFEs).
2. Construção de Estado Intermediário: Estados intermediários $X_t$ são amostrados diretamente do kernel de ruído original $q_t(\cdot | X_1)$ condicionado ao ponto final gerado, contornando a necessidade de simulação forward de SDE.
3. Adjoint de Forma Fechada: O estado adjoint é computado via uma escala de forma fechada do gradiente do ponto final $\nabla g(X_1)$, removendo completamente a simulação backward.
4. Parametrização de Controle: O controle $u(x,t)$ é parametrizado via um modelo de velocidade treinável (por exemplo, LoRA) inicializado a partir do modelo pré-treinado, permitindo que o ajuste fino comece a partir das dinâmicas pré-treinadas.

Contribuições Principais

• Proposta de Algoritmo: A introdução da EAM, um algoritmo de ajuste fino baseado em gradiente de recompensa que alcança os mesmos objetivos teóricos que a AM, mas com custo computacional significativamente reduzido.
• Caracterização Teórica: Uma caracterização de desvios base lineares que satisfazem os requisitos para construção eficiente de trajetórias (sem memória) e computação de adjoint de forma fechada.
• Validação Empírica: Demonstração de que a EAM converge até 4$\times$ mais rápido que a AM em benchmarks padrão de T2I, enquanto iguala ou supera seu desempenho em múltiplas métricas.

Resultados

Os autores avaliaram a EAM no Stable Diffusion 3.5-Medium (SD3.5-M) usando o dataset Pick-a-Pic, comparando-a com a AM e a linha de base SD3.5-M.

• Eficiência: A EAM reduz o custo de treinamento por iteração em até 4$\times$. Isso é alcançado substituindo simulações de SDE de 40 passos por amostragem de ponto final de ODE de 10 passos e eliminando simulações de adjoint backward.
• Desempenho: Em benchmarks incluindo PickScore, ImageReward, HPSv2.1, CLIPScore e Aesthetics, a EAM consistentemente iguala ou supera a AM.
  • Em configurações de recompensa única (PickScore), a EAM alcançou um PickScore de 23,16 comparado aos 23,00 da AM.
  • Em configurações de múltiplas recompensas (PickScore + HPSv2.1 + Aesthetics), a EAM alcançou um PickScore de 22,93 e uma pontuação de Aesthetics de 6,56, superando a AM (6,21).
• Melhorias Qualitativas: Comparações visuais indicam que a EAM preserva melhor detalhes de alta granularidade e composições complexas em comparação com a AM, que às vezes suaviza detalhes. Isso é atribuído à consistência da EAM entre as trajetórias de treinamento (ruído condicionado ao ponto final) e de inferência.

Significado e Alegações

O artigo alega que as limitações primárias da Correspondência Adjoint não são inerentes à própria formulação de SOC, mas sim artefatos da escolha específica de desvio base. Ao redesenhar o desvio base para ser linear e sem memória, e ajustar correspondentemente o custo terminal, a EAM remove os dois principais gargalos computacionais (simulação forward de SDE e simulação backward de adjoint) sem sacrificar as garantias teóricas de alinhamento.

Os autores posicionam a EAM como uma solução prática e escalável para ajuste fino por recompensa que mantém a eficiência de amostragem de métodos baseados em gradiente, enquanto reduz drasticamente a sobrecarga computacional, tornando o alinhamento de alta qualidade de modelos de difusão mais acessível. O trabalho não alega resolver todos os desafios de otimização (por exemplo, replay de buffer ou agendamentos adaptativos são notados como trabalho futuro), mas estabelece uma nova linha de base para eficiência neste domínio.