ParetoSlider: Diffusion Models Post-Training for Continuous Reward Control

Título: ParetoSlider: Pós-treinamento de Modelos de Difusão para Controle Contínuo de Recompensas

1. O Problema

O alinhamento de modelos generativos (como difusão e flow-matching) com preferências humanas é frequentemente realizado através de Aprendizado por Reforço (RL). No entanto, a maioria dos métodos atuais baseia-se em uma única recompensa escalar ou em uma escalarização precoce (early scalarization), onde múltiplos objetivos são combinados em uma soma ponderada fixa antes do treinamento.

Isso cria dois problemas principais:

1. Compromisso Fixo: O modelo é treinado para um único ponto de equilíbrio (trade-off) entre objetivos conflitantes (ex: aderência ao prompt vs. preservação da imagem original).
2. Falta de Controle em Inferência: Uma vez treinado, o modelo não permite que o usuário ajuste dinamicamente o equilíbrio entre os objetivos durante a geração. Para obter um novo equilíbrio, seria necessário treinar e armazenar múltiplos checkpoints (um para cada combinação de pesos), o que é ineficiente e não escalável.

2. Metodologia: ParetoSlider

O ParetoSlider é um framework de Aprendizado por Reforço Multi-Objetivo (MORL) projetado para treinar um único modelo capaz de aproximar toda a Frente de Pareto (o conjunto de soluções ótimas onde nenhum objetivo pode ser melhorado sem piorar outro).

Principais Componentes Técnicos:

• Condicionamento por Vetor de Preferência ($\omega$):

  • Em vez de treinar modelos separados, o modelo é condicionado explicitamente por um vetor de preferência $\omega$ durante o treinamento.
  • $\omega$ é um vetor de pesos que reside no simplex de probabilidade (soma dos pesos = 1, pesos $\ge$ 0), especificando o trade-off desejado entre os $M$ objetivos de recompensa.
  • Isso permite que o modelo aprenda a mapear qualquer preferência do usuário para o ponto ótimo correspondente na frente de Pareto.

• Escalarização Tardia (Late Scalarization):

  • Um desafio crítico é que diferentes funções de recompensa podem ter escalas numéricas e variâncias muito diferentes, o que pode fazer com que uma recompensa "domine" o sinal de gradiente, ignorando a preferência do usuário ($\omega$).
  • O ParetoSlider evita a escalarização precoce. Em vez disso, ele calcula as vantagens (advantages) e as perdas de treinamento independentemente para cada recompensa.
  • A agregação ponderada pelos pesos $\omega$ ocorre apenas no último passo, ao combinar as perdas individuais antes da atualização do gradiente. Isso garante que a otimização respeite estritamente a preferência definida por $\omega$, independentemente da magnitude bruta das recompensas.

• Arquitetura e Condicionamento:

  • O método é construído sobre o DiffusionNFT (um framework de RL eficiente para modelos de flow-matching).
  • O vetor $\omega$ é injetado no modelo através de módulos de condicionamento leves (LoRA e adaptadores de modulação) que variam conforme a arquitetura base:
    • SD3.5 (Texto-para-Imagem): Injeção via embedding de tempo e modulação residual compartilhada nos blocos do transformador.
    • FluxKontext (Edição de Imagem): Modulação dos parâmetros AdaLN do fluxo de contexto (texto).
    • LTX-2 (Texto-para-Vídeo): Mecanismo de resíduo compartilhado nos blocos do fluxo de vídeo.

• Processo de Treinamento:

  1. Geração de um grupo de amostras para um prompt e um vetor de preferência $\omega$ amostrado.
  2. Avaliação por múltiplos modelos de recompensa.
  3. Normalização independente das vantagens por canal de recompensa.
  4. Cálculo de perdas individuais de DiffusionNFT para cada recompensa.
  5. Agregação das perdas usando $\omega$ e atualização do modelo.

3. Contribuições Chave

1. Controle Contínuo em Inferência: Pela primeira vez em modelos de difusão, permite-se um controle deslizante suave entre objetivos conflitantes sem a necessidade de re-treinamento ou interpolação de checkpoints.
2. Aproximação da Frente de Pareto: O modelo aprende a cobrir todo o espectro de trade-offs ótimos em um único conjunto de parâmetros.
3. Estratégia de Escalarização Tardia: Demonstra que normalizar as vantagens por recompensa antes da agregação ponderada é crucial para evitar o "sequestro" do gradiente por recompensas de grande magnitude e garantir fidelidade à preferência do usuário.
4. Generalização Multi-Modal: O método foi validado com sucesso em três tarefas distintas: Síntese de Imagem (T2I), Edição de Imagem (I2I) e Geração de Vídeo (T2V).

4. Resultados

Os autores avaliaram o ParetoSlider em três backbones de flow-matching de última geração: Stable Diffusion 3.5, FluxKontext e LTX-2.

• Comparação com Baselines:
  • O modelo único do ParetoSlider igualou ou superou o desempenho de múltiplos modelos treinados separadamente com pesos fixos (Fixed-Weights).
  • Superou métodos de interpolação de checkpoints (como Rewarded Soups) e abordagens sem treinamento (como Prompt Rewriting ou Classifier-Free Guidance), que oferecem apenas pontos de operação discretos ou controlam o trade-off de forma grosseira.
• Qualidade Visual:
  • T2I: Transições suaves entre realismo fotográfico e esboço (sketch), ou entre estilos como anime e aquarela.
  • I2I: Controle fino entre a preservação da identidade da imagem original e a aderência estrita às instruções de edição (ex: transformar um retrato em um personagem de fantasia sem perder os traços faciais).
  • T2V: Navegação contínua entre realismo e animação 3D (estilo Disney/Pixar).
• Métricas Quantitativas:
  • O método apresentou o maior Hypervolume (HV), indicando uma melhor cobertura e convergência para a frente de Pareto em comparação com todas as baselines.
  • A ablação mostrou que a Late Scalarization e as arquiteturas de condicionamento específicas (como Shared Residual) são superiores a alternativas como Early Scalarization ou condicionamento via tokens.

5. Significância

O ParetoSlider representa uma mudança de paradigma no alinhamento de modelos generativos visuais. Ao permitir que um único modelo capture a complexidade de múltiplos objetivos conflitantes e ofereça controle granular em tempo de inferência, ele resolve o dilema entre flexibilidade e eficiência.

Isso abre caminho para interfaces de usuário mais intuitivas, onde usuários não especialistas podem ajustar dinamicamente o "estilo" ou o "grau de edição" de uma geração (ex: "mais realista", "mais fiel ao original") sem precisar de conhecimento técnico sobre pesos de treinamento ou múltiplos modelos. O trabalho estabelece um novo padrão escalável para o alinhamento multi-objetivo em IA generativa.