DiffusionOPD: Uma Perspectiva Unificada da Destilação On-Policy em Modelos de Difusão

Resumo

A aprendizagem por reforço emergiu como uma ferramenta poderosa para aprimorar modelos de difusão de texto para imagem, mas os métodos existentes estão amplamente limitados à otimização de tarefa única. Estender a aprendizagem por reforço para múltiplas tarefas é desafiador: a otimização conjunta sofre de interferência entre tarefas e desequilíbrio, enquanto a aprendizagem por reforço em cascata é complicada e propensa ao esquecimento catastrófico. Propomos o DiffusionOPD, um novo paradigma de treinamento multitarefa para modelos de difusão baseado na Destilação de Políticas Online (OPD). O DiffusionOPD primeiro treina professores específicos para cada tarefa de forma independente e, em seguida, destila suas capacidades em um estudante unificado ao longo das trajetórias de rollout do próprio estudante. Isso desacopla a exploração de tarefa única da integração multitarefa e evita o ônus da otimização de resolver todas as tarefas conjuntamente do zero. Teoricamente, elevamos o arcabouço OPD de tokens discretos para processos de Markov de estado contínuo, derivando um objetivo KL por etapa em forma fechada que unifica tanto o refinamento estocástico por EDE quanto o determinístico por EDO por meio de correspondência de médias. Demonstramos formal e empiricamente que esse gradiente analítico fornece menor variância e melhor generalidade em comparação com os gradientes de política convencionais do tipo PPO. Experimentos extensivos mostram que o DiffusionOPD supera consistentemente tanto as linhas de base de aprendizagem por reforço com múltiplas recompensas quanto as de aprendizagem por reforço em cascata em eficiência de treinamento e desempenho final, enquanto alcança resultados de ponta em todos os benchmarks avaliados.

English

Reinforcement learning has emerged as a powerful tool for improving diffusion-based text-to-image models, but existing methods are largely limited to single-task optimization. Extending RL to multiple tasks is challenging: joint optimization suffers from cross-task interference and imbalance, while cascade RL is cumbersome and prone to catastrophic forgetting. We propose DiffusionOPD, a new multi-task training paradigm for diffusion models based on Online Policy Distillation (OPD). DiffusionOPD first trains task-specific teachers independently, then distills their capabilities into a unified student along the student own rollout trajectories. This decouples single-task exploration from multi-task integration and avoids the optimization burden of solving all tasks jointly from scratch. Theoretically, we lift the OPD framework from discrete tokens to continuous-state Markov processes, deriving a closed-form per-step KL objective that unifies both stochastic SDE and deterministic ODE refinement via mean-matching. We formally and empirically demonstrate that this analytic gradient provides lower variance and better generality compared to conventional PPO-style policy gradients. Extensive experiments show that DiffusionOPD consistently surpasses both multi-reward RL and cascade RL baselines in training efficiency and final performance, while achieving state-of-the-art results on all evaluated benchmarks.