Resumo Técnico: DiffusionOPD

Declaração do Problema

O aprendizado por reforço (RL) provou ser eficaz para otimizar modelos de texto para imagem baseados em difusão contra sinais de recompensa de objetivo único (por exemplo, estética ou fidelidade ao texto). No entanto, os métodos existentes lutam em cenários de múltiplas tarefas onde um único modelo deve satisfazer múltiplos objetivos, potencialmente conflitantes, simultaneamente. As abordagens atuais enfrentam duas limitações primárias:

1. Otimização Conjunta: O treinamento de todas as tarefas simultaneamente frequentemente leva a interferência entre tarefas e desequilíbrio de dificuldade de tarefa, onde tarefas mais fáceis dominam a dinâmica de aprendizado, suprimindo sinais de tarefas mais desafiadoras.
2. RL em Cascata: A otimização de tarefas sequencialmente evita conflito direto de gradiente, mas é cômoda, exigindo múltiplos estágios de treinamento e agendamento cuidadoso. Além disso, é propensa ao esquecimento catastrófico, onde a adaptação a tarefas posteriores degrada o desempenho nas anteriormente aprendidas.

O artigo argumenta que o RL de múltiplas tarefas deve ser desacoplado em dois processos distintos: exploração on-policy de tarefa única e integração de capacidade de múltiplas tarefas.

Metodologia: DiffusionOPD

Os autores propõem o DiffusionOPD, um paradigma unificado de treinamento de múltiplas tarefas baseado em Destilação de Política Online (OPD). O framework opera em dois estágios:

1. Estágio de Treinamento: Especialização Independente do Professor

Em vez de otimização conjunta, o método primeiro treina um conjunto de modelos professores específicos de tarefa independentemente. Cada professor é otimizado para uma única tarefa (por exemplo, OCR, estética, alinhamento composicional) usando algoritmos RL padrão prontamente disponíveis. Isso garante que cada professor atinja um teto de alto desempenho em seu objetivo específico sem interferência de outras tarefas.

2. Estágio de Integração: Destilação On-Policy

Um modelo aluno unificado é então treinado para destilar as capacidades desses professores especializados. Crucialmente, o aluno aprende ao longo de suas próprias trajetórias de rollout (on-policy).

• Processo: O aluno gera uma trajetória de remoção de ruído. Para cada etapa nesta trajetória, o professor específico da tarefa correspondente fornece supervisão.
• Objetivo: O aluno é treinado para corresponder ao comportamento de transição do professor em cada etapa de remoção de ruído.

Derivação Teórica: De Discreto a Contínuo

A contribuição teórica central é elevar o framework OPD de transições de tokens discretos (LLMs) para processos de Markov de estado contínuo (Modelos de Difusão).

• Formulação de Cadeia de Markov: O processo de remoção de ruído da difusão é modelado como uma cadeia de Markov de tempo discreto induzida por uma Equação Diferencial Estocástica (EDE) de tempo reverso.
• Transições Gaussianas: Tanto o aluno quanto o professor definem kernels de transição Gaussianos de um passo. Como eles compartilham a mesma covariância (determinada pelo agendador e nível de ruído) mas diferem em suas médias, a divergência Kullback-Leibler (KL) reversa admite uma expressão de forma fechada:
  $$\text{KL}(p_S \parallel p_T) = \frac{\|\mu_S - \mu_T\|^2}{2\sigma^2}$$
  onde $\mu_S$ e $\mu_T$ são as médias de transição do aluno e do professor, respectivamente.
• Otimização: O aluno é otimizado minimizando este objetivo KL de forma fechada via retropropagação direta.

Comparação com Gradientes Estilo PPO

O artigo contrasta esta abordagem de forma fechada com uma formulação de gradiente de política estilo PPO (tratando o professor como um modelo de recompensa de processo).

• Redução de Variância: Embora os gradientes estilo PPO sejam não tendenciosos em expectativa, eles introduzem um termo de função de pontuação proporcional ao ruído gaussiano, levando a uma variância de gradiente mais alta.
• Compatibilidade Determinística: O objetivo KL de forma fechada estende-se naturalmente para amostradores ODE determinísticos (onde a transição é uma correspondência direta $L_2$), enquanto objetivos estilo PPO dependem de densidades de política estocásticas.
• Conclusão: O KL de forma fechada fornece um princípio de otimização de menor variância e mais geral, aplicável tanto a amostradores EDE quanto ODE.

Contribuições Principais

1. Framework DiffusionOPD: Um novo paradigma que desacopla a exploração de tarefa única da integração de múltiplas tarefas, utilizando professores específicos de tarefa para supervisionar um aluno unificado ao longo de suas próprias trajetórias.
2. Unificação Teórica: Derivação de um objetivo KL por etapa de forma fechada e principial para modelos de difusão que unifica o refinamento estocástico EDE e determinístico ODE, oferecendo otimização de menor variância do que gradientes estilo PPO.
3. Validação Empírica: Experimentos extensivos demonstrando que o DiffusionOPD supera tanto as linhas de base de RL de múltiplas tarefas (otimização conjunta) quanto as linhas de base de RL em cascata em eficiência de treinamento e desempenho final.

Resultados Experimentais

Os autores avaliaram o DiffusionOPD em diversos benchmarks incluindo GenEval (geração composicional), OCR (renderização de texto) e Estética (qualidade visual), usando o SD3.5-Medium como modelo base.

• Desempenho: O DiffusionOPD alcançou resultados State-of-the-Art (SOTA) em todos os benchmarks avaliados, superando Multi-Task GRPO-Guard, Multi-Task NFT e Cascade NFT.
  • Pontuação Média: DiffusionOPD (0,929) vs. Cascade NFT (0,851) vs. Multi-Task NFT (0,715).
• Eficiência: O método demonstrou convergência significativamente mais rápida e um teto de desempenho mais alto em comparação com as linhas de base de RL de múltiplas tarefas.
• Tempo de Treinamento: O DiffusionOPD exigiu menos tempo total de relógio para atingir pontuações-alvo em comparação com métodos de RL conjunta, evitando ao mesmo tempo a sobrecarga sequencial dos métodos em cascata.
• Estudos de Ablação:
  • Formulação de Perda: O objetivo KL de forma fechada superou consistentemente os gradientes de política estilo PPO tanto na velocidade de convergência quanto na recompensa final.
  • Nível de Ruído: Níveis de ruído mais baixos (aproximando-se da amostragem ODE) levaram a uma convergência mais rápida e desempenho mais alto, com o amostrador ODE sendo até cinco vezes mais eficiente do que o amostrador EDE com alto ruído.

Significado e Alegações

O artigo afirma que o DiffusionOPD aborda as limitações fundamentais do RL de múltiplas tarefas atual para modelos de difusão desacoplando a exploração da integração. Ao aproveitar as propriedades analíticas das transições de difusão, ele fornece um princípio de treinamento unificado e de baixa variância que é aplicável através de diferentes tipos de amostradores (EDE e ODE).

Os autores posicionam o DiffusionOPD não meramente como um método de destilação específico, mas como um framework geral para destilação on-policy em modelos de difusão. Eles afirmam que esta abordagem resolve efetivamente o trade-off entre conflito de otimização (no treinamento conjunto) e esquecimento catastrófico/ineficiência (no treinamento em cascata), estabelecendo uma nova fundação para modelagem de difusão de múltiplas tarefas e alinhada a preferências.