Resumo Técnico: Modelo de Valor Costurado para Alinhamento de Difusão (StitchVM)

1. Declaração do Problema

Modelos generativos baseados em difusão e fluxo destacam-se na captura de distribuições de dados, mas frequentemente exigem alinhamento com recompensas específicas da tarefa (por exemplo, fidelidade ao prompt, preferência estética, preferência humana) para serem praticamente úteis. O desafio central nesse alinhamento é que as recompensas são tipicamente definidas para imagens de saída limpas ($x_0$), enquanto os procedimentos de alinhamento (tanto no tempo de inferência quanto no de treinamento) exigem a estimativa da recompensa esperada de amostras limpas induzidas por latentes intermediários ruidosos ($z_t$). Essa estimativa é conhecida como função de valor.

Os métodos existentes para aproximar essa função de valor sofrem de trade-offs significativos:

• Aproximação de Tweedie: Estima a média posterior da amostra limpa e computa a recompensa nesse proxy. Embora eficiente, introduz viés substancial em regimes de alto ruído e requer avaliações adicionais e dispendiosas do denoiser e do decodificador VAE.
• Aproximação de Monte Carlo (MC): Executa múltiplas trajetórias de desruidização a partir de $z_t$ e faz a média das recompensas. Embora mais precisa, incorre em custos computacionais proibitivos devido à necessidade de múltiplas execuções completas.
• Modelos de Valor Aprendidos: Treinar diretamente uma função de valor para latentes ruidosos é um problema em aberto. Tentativas existentes frequentemente dependem de treinamento em pequena escala ou conjuntos de dados de preferência restritos, resultando em precisão e generalização inferiores em comparação com modelos de recompensa no espaço de pixels em escala de fundação. Além disso, treinar um novo modelo de valor para cada novo backbone de difusão ou modelo de recompensa é computacionalmente insustentável.

2. Metodologia: StitchVM

Os autores propõem o StitchVM (Stitched Value Model), um framework que transfere as capacidades de modelos de recompensa no espaço de pixels pré-treinados e de alta qualidade para o regime de latentes ruidosos com custo computacional mínimo. A abordagem alavanza o model stitching (costura de modelos), uma técnica que recombina camadas de diferentes redes neurais.

Arquitetura Central

O StitchVM constrói um modelo híbrido ao anexar um backbone de difusão truncado e congelado (a "cabeça") a um modelo de recompensa pré-treinado fatiado (a "cauda") por meio de uma camada de costura leve.

• Cabeça: Um backbone de difusão (por exemplo, SD 3.5, FLUX) truncado na camada $i$. Este componente processa nativamente latentes ruidosos e extrai características conscientes do ruído.
• Cauda: Um modelo de recompensa pré-treinado (por exemplo, CLIP, HPSv2, Aesthetic Predictor) começando na camada $j$. Este componente retém a capacidade robusta e específica da tarefa de recompensa aprendida durante o pré-treinamento em grande escala.
• Camada de Costura ($s_\psi$): Um módulo leve que alinha o espaço de características da cabeça de difusão com a cauda de recompensa.

Procedimento de Treinamento

O treinamento do StitchVM ocorre em duas etapas:

1. Seleção de Interface (Forma Fechada): Os autores buscam os índices de costura ótimos $(i, j)$ onde as características de difusão (de latentes ruidosos) e as características de recompensa (de imagens limpas) são mais compatíveis. Isso é alcançado ajustando um mapeamento linear $W$ entre características pareadas extraídas de um pequeno conjunto de sondagem de imagens limpas e seus latentes ruidosos correspondentes. O par que minimiza a perda de correspondência de características é selecionado.
2. Ajuste Fino Leve: Uma vez fixada a interface, a camada de costura e a cauda de recompensa truncada são ajustadas finamente usando imagens limpas não rotuladas. O objetivo é regredir a função de valor $V(z_t)$ para a pontuação de recompensa original $r(z_0)$ da imagem limpa. Este objetivo off-policy permite que o modelo aprenda a recompensa esperada de amostras limpas condicionadas a latentes ruidosos sem exigir execuções completas.

Aplicação ao Alinhamento

O StitchVM serve como um estimador direto de função de valor, substituindo aproximações tanto no alinhamento no tempo de inferência quanto no de treinamento:

• Tempo de Inferência:
  • DPS (Diffusion Posterior Sampling): Substitui o gradiente baseado em Tweedie pelo gradiente direto do StitchVM ($\nabla_{z_t} V(z_t)$), eliminando a necessidade de retropropagação através do denoiser e do VAE.
  • FK Steering: Habilita a "escalonamento de proposta", onde cada partícula gera múltiplas propostas locais. Estas são pontuadas de forma barata pelo StitchVM para selecionar o melhor candidato, um processo que seria muito custoso com aproximações de Tweedie.
• Tempo de Treinamento:
  • Ajuste Fino Direto de Recompensa (por exemplo, AlignProp, DRaFT): Permite que o treinamento pare em latentes ruidosos intermediários $z_\tau$ e use a previsão do StitchVM como alvo de supervisão, evitando execuções completas de trajetória e longas cadeias de retropropagação.
  • DiffusionNFT: Substitui recompensas finais limpas por estimativas de valor intermediárias, permitindo RL online eficiente no processo forward.

3. Contribuições Principais

• Um Novo Paradigma para Modelagem de Valor: O StitchVM demonstra que modelos de recompensa poderosos no espaço de pixels podem ser transferidos efetivamente para o regime de latentes ruidosos sem retreinamento em grande escala, simplesmente costurando-os a um backbone de difusão.
• Eficiência e Precisão: O método alcança desempenho de última geração na estimativa de valor enquanto é excepcionalmente leve. Por exemplo, costurar um CLIP ViT-L em um modelo de valor SD 3.5 Medium leva apenas $\approx 10$ horas de GPU.
• Ganhos de Desempenho no Alinhamento:
  • Inferência: O DPS torna-se 3,2$\times$ mais rápido com redução de 50% no pico de memória da GPU. O FK steering com StitchVM supera significativamente o FK steering padrão e a amostragem Best-of-$N$, particularmente ao permitir seleção eficiente de propostas locais.
  • Treinamento: O DiffusionNFT torna-se 2,3$\times$ mais rápido, e o ajuste fino direto de recompensa (DRaFT) torna-se 1,3$\times$ mais rápido, alcançando ao mesmo tempo maior qualidade de geração (por exemplo, +30% GenEval) devido à supervisão efetiva em etapas de alto ruído.
• Generalização: A abordagem funciona em diversos backbones de difusão (SD 3.5 Medium/Large, FLUX) e modelos de recompensa (CLIP, DFN-CLIP, HPSv2, Aesthetic Predictor).

4. Resultados Experimentais

O artigo avalia o StitchVM em múltiplos benchmarks:

• Qualidade da Função de Valor: Em latentes ruidosos com níveis de ruído variados ($\sigma \in [0.1, 0.9]$), o StitchVM corresponde de perto ao desempenho dos modelos de recompensa limpos originais em recuperação zero-shot (MSCOCO, Flickr30K), previsão de preferência (HPDv2, ImageReward) e pontuação estética (AVA). Supera significativamente baselines como costura VAE (que colapsa em alto ruído) e modelos treinados do zero em dados ruidosos (por exemplo, NoisyCLIP, DiNa-LRM).
• Alinhamento de Inferência: Em experimentos de DPS e FK steering no DrawBench, os métodos aprimorados com StitchVM melhoram consistentemente métricas como ImageReward, HPSv2 e GenEval enquanto reduzem drasticamente o custo computacional. Notavelmente, aumentar o número de propostas locais ($M$) no FK steering usando StitchVM prova ser mais eficiente do que aumentar o número de partículas ($N$).
• Alinhamento de Treinamento: As curvas de treinamento mostram que as variantes do StitchVM atingem pontuações finais mais altas com menos horas de GPU em comparação com DRaFT e DiffusionNFT padrão. Resultados qualitativos indicam que os modelos aumentados com StitchVM convergem mais rápido e produzem imagens mais nítidas e de maior qualidade mais cedo no treinamento.

5. Significado e Alegações

O artigo afirma que o StitchVM fornece um modelo genérico para combinar backbones de difusão latente com modelos feedforward no espaço de pixels sem sacrificar o extenso pré-treinamento de qualquer um deles. Ao elevar modelos de recompensa poderosos para o espaço latente, ele habilita um novo estilo de alinhamento de difusão onde a função de valor correta para latentes ruidosos é construída uma vez e amortizada sobre muitas amostras e iterações.

Os autores posicionam o StitchVM não como uma substituição para o pré-treinamento de modelos de fundação, mas como um mecanismo prático e de baixo custo para adaptar esses modelos a tarefas específicas. Eles enfatizam que essa abordagem torna o alinhamento baseado em modelo de valor acessível e eficiente, potencialmente acelerando a pesquisa em modelagem generativa mais segura e confiável. O trabalho sugere que a "costura" de componentes pré-treinados é uma alternativa viável aos custos proibitivos de treinar funções de valor do zero ou depender de aproximações enviesadas.