DiffusionOPD : Une perspective unifiée de la distillation on-policy dans les modèles de diffusion

Résumé

L'apprentissage par renforcement est devenu un outil puissant pour améliorer les modèles texte-image basés sur la diffusion, mais les méthodes existantes sont largement limitées à l'optimisation mono-tâche. Étendre l'apprentissage par renforcement à plusieurs tâches est difficile : l'optimisation conjointe souffre d'interférence entre tâches et de déséquilibre, tandis que l'approche en cascade (cascade RL) est lourde et sujette à l'oubli catastrophique. Nous proposons DiffusionOPD, un nouveau paradigme d'entraînement multi-tâche pour les modèles de diffusion basé sur la distillation de politique en ligne (Online Policy Distillation, OPD). DiffusionOPD entraîne d'abord des enseignants spécifiques à chaque tâche de manière indépendante, puis distille leurs capacités dans un étudiant unifié le long des trajectoires de déploiement (rollout) propres à l'étudiant. Cela découple l'exploration mono-tâche de l'intégration multi-tâche et évite la charge d'optimisation liée à la résolution conjointe de toutes les tâches à partir de zéro. Théoriquement, nous étendons le cadre OPD des tokens discrets aux processus de Markov à états continus, en dérivant un objectif KL par étape sous forme fermée qui unifie le raffinement par EDS stochastique et par EDO déterministe via l'appariement des moyennes. Nous démontrons formellement et empiriquement que ce gradient analytique offre une variance plus faible et une meilleure généralité par rapport aux gradients de politique de type PPO classiques. Des expériences approfondies montrent que DiffusionOPD surpasse systématiquement les références (baselines) de RL multi-récompenses et de RL en cascade en termes d'efficacité d'entraînement et de performance finale, tout en atteignant des résultats de pointe (état de l'art) sur tous les benchmarks évalués.

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.