PG-MAP: Joint MAP Optimization for Inference-Time Alignment of Diffusion and Flow-Matching Models

Resumo Técnico: PG-MAP: Otimização de MAP Conjunta para Alinhamento em Tempo de Inferência de Modelos de Difusão e Flow-Matching

Definição do Problema
Os atuais métodos de alinhamento em tempo de inferência para modelos pré-treinados de texto-para-imagem (Difusão e Flow-Matching) operam tipicamente ao longo de um único eixo de controle. Métodos do lado do condicionamento (ex: edição de atenção, aprendizado de embedding de texto) modificam o embedding de texto $c$, enquanto métodos do lado do latente (ex: perturbações de latente guiadas por recompensa) modificam o estado latente $z_t$. Os frameworks existentes falham em modelar as dependências conjuntas entre variáveis de condicionamento e latentes, tratando as atualizações como aditivas em vez de coordenadas. Além disso, esses métodos são frequentemente estáticos, fixando estratégias de controle no início da trajetória ($z_T$) ou offline, ignorando as necessidades dinâmicas de informação do denoiser conforme a trajetória evolui do regime de alto ruído (layout global) para o de baixo ruído (detalhe local). Essa limitação dificulta a transferência de técnicas de alinhamento entre diferentes transportes generativos (Difusão vs. Flow-Matching).

Metodologia: PG-MAP
Os autores propõem o PG-MAP (Preference-Guided Adaptive MAP), um framework livre de treinamento que formula o alinhamento em tempo de inferência como uma otimização de energia proximal / Gibbs-MAP em nível de trajetória. A inovação central é um acoplamento de consistência de avanço que otimiza conjuntamente o condicionamento $c$ e o estado latente $z_t$ em cada passo de denoising.

• Função Objetivo: O framework define uma função de energia dependente do passo $J_t(c, z_t)$ composta por três componentes:

  1. Resíduo de Consistência de Avanço: Um termo $\ell_t(c, z_t)$ que penaliza o desvio entre o latente atual $z_t$ e a previsão do próprio modelo para o estado do passo anterior, garantindo que a trajetória permaneça consistente com o kernel de avanço do denoiser.
  2. Priors de Ancoragem Gaussiana: Termos $R_c(c)$ e $R_z(z_t)$ que ancoram as variáveis otimizadas aos seus valores não perturbados ($c_0$ e o ponto da trajetória DDIM $z^{ddim}_t$). Crucialmente, a escala do prior latente $\sigma_z(t)$ é adaptativa ao cronograma (schedule-adaptive), acompanhando a escala de ruído marginal do processo de difusão ($\sigma_z(t) = \gamma\sqrt{1-\bar{\alpha}_t}$) para fornecer uma região de confiança invariante de escala que se estreita conforme a trajetória se aproxima do manifold de dados.
  3. Inclinação de Recompensa de Preferência (Preference Reward Tilt): Um termo $\lambda Q(\hat{x}_0, y)$ que guia a otimização em direção à preferência humana (ex: PickScore, HPS) com base na prévia por passo da imagem limpa gerada pelo modelo.

• Conjunto Ativo Adaptativo: O framework emprega um conjunto ativo dependente do passo $A_t \subseteq \{c, z_t\}$ e um portão de recompensa $\lambda_t$.

  • Modelos de Difusão (SD 1.5, SDXL): O framework refina conjuntamente tanto $c$ quanto $z_t$, principalmente na janela de alto ruído, permitindo que o condicionamento e as variáveis latentes informem um ao outro conforme a trajetória se desenrola.
  • Modelos de Flow-Matching (SD3.5): Devido à natureza determinística da ODE e à alta capacidade dos encoders de texto concatenados, o framework reduz-se a uma variante apenas de latente (UG-FM). Aqui, o conjunto ativo colapsa para $\{z_t\}$, restrito à janela do lado dos dados (baixo ruído), pois as atualizações de condicionamento foram consideradas ineficazes e as perturbações do lado do ruído são excessivamente amplificadas pelas dinâmicas da ODE.

• Otimização: O par refinado $(c^\star_t, z^\star_t)$ é obtido através de um pequeno número de passos de ascensão conjunta (tipicamente $K=2$) usando gradientes derivados da função de energia. Este processo é composível com métodos existentes como o Classifier-Free Guidance (CFG), visto que o CFG modifica o campo vetorial enquanto o PG-MAP move o ponto de consulta sob um campo vetorial fixo.

Principais Contribuições

1. Framework MAP Conjunto (c, $z_t$): O primeiro framework de tempo de inferência a acoplar variáveis de condicionamento e latentes através do kernel de avanço do denoiser, abordando simultaneamente falhas de composição (lado c) e falhas de textura (lado z).
2. Objetivo Unificado: O objetivo $J_t$ recupera analiticamente métodos de eixo único anteriores (MAP-c, Reward-z) como casos especiais, enquanto o CFG permanece como um mecanismo externo composível, em vez de um caso especial.
3. Otimização Adaptativa ao Cronograma: Introduz uma região de confiança dependente do tempo e um conjunto ativo que se adapta ao cronograma de ruído, relaxando as restrições em alto ruído e estreitando-as próximo ao fim dos dados.
4. Análise Dependente do Transporte: Demonstra que o conjunto ativo ideal depende do transporte generativo: a Difusão beneficia-se do refinamento conjunto em alto ruído, enquanto o Flow-Matching beneficia-se do refinamento apenas de latente no final dos dados, uma hipótese apoiada por análise de perturbação local e validação empírica.

Resultos

• Backbones de Difusão (SD 1.5, SDXL): O PG-MAP melhora consistentemente as métricas de alinhamento. No SDXL, alcança uma taxa de vitória de 56,4% no PickScore e 56,2% em Estética (Aesthetic) contra um baseline estático. Combinar o PG-MAP com CFG ajustado produz o melhor desempenho geral, atingindo 64,6% de HPS e 56,5% de Estética.
• Backbone de Flow-Matching (SD3.5-medium): A variante apenas de latente (UG-FM) alcança taxas de vitória de 91,9% no PickScore e 75,7% no HPS contra o baseline estático. Experimentos controlados descartam artefatos relacionados ao ruído, confirmando que os ganhos se devem a perturbações latentes estruturadas de baixa frequência.
• Avaliação Humana: No SDXL, o PG-MAP é preferido em relação aos baselines estáticos, CFG ajustado e Universal Guidance com computação equivalente em avaliações pareadas humanas (taxas de vitória de 60–67%).
• Roteamento Oracle: Uma análise oracle revela que diferentes tipos de prompts (ex: vinculação de atributos vs. cenas atmosféricas) beneficiam-se de diferentes variantes de otimização. Um selecionador por prompt poderia, teoricamente, melhorar as métricas em mais 5–14 pontos percentuais, indicando um espaço significativo para futuros mecanismos de roteamento.

Significância e Alegações
O artigo afirma que o PG-MAP representa uma mudança de heurísticas de orientação de eixo único e estáticas para uma otimização dinâmica e consciente da trajetória como um princípio de design padrão para alinhamento em tempo de inferência. Ao formular o alinhamento como uma otimização de energia proximal com acoplamento de consistência de avanço, o framework permite atualizações coordenadas entre modalidades que são compatíveis tanto com transportes de difusão quanto de flow-matching. Os autores enfatizam que os ingredientes "estruturais" são o acoplamento conjunto e o agendamento não estacionário, em vez de simplesmente passos maiores ou sinais de recompensa mais fortes. O trabalho fornece uma base teórica e empírica unificada para entender como o condicionamento e as variáveis latentes interagem durante o processo de geração, oferecendo uma abordagem flexível e livre de treinamento para melhorar o alinhamento de modelos generativos.