Résumé Technique : DiffusionOPD

Énoncé du Problème

L'apprentissage par renforcement (RL) a prouvé son efficacité pour l'optimisation des modèles de texte vers image basés sur la diffusion contre des signaux de récompense à objectif unique (par exemple, l'esthétique ou la fidélité au texte). Cependant, les méthodes existantes peinent dans les environnements multi-tâches où un modèle unique doit satisfaire simultanément plusieurs objectifs, potentiellement contradictoires. Les approches actuelles font face à deux limitations principales :

1. Optimisation Conjoints : L'entraînement de toutes les tâches simultanément conduit souvent à des interférences inter-tâches et à un déséquilibre de difficulté des tâches, où les tâches plus faciles dominent la dynamique d'apprentissage, supprimant les signaux provenant de tâches plus complexes.
2. RL en Cascade : L'optimisation des tâches de manière séquentielle évite les conflits de gradients directs mais est lourde, nécessitant plusieurs étapes d'entraînement et une planification minutieuse. De plus, elle est sujette à l'oubli catastrophique, où l'adaptation aux tâches ultérieures dégrade les performances sur celles apprises précédemment.

L'article soutient que le RL multi-tâches devrait être découplé en deux processus distincts : l'exploration en ligne à politique unique par tâche et l'intégration des capacités multi-tâches.

Méthodologie : DiffusionOPD

Les auteurs proposent DiffusionOPD, un paradigme d'entraînement multi-tâches unifié basé sur la Distillation de Politique en Ligne (OPD). Le cadre fonctionne en deux étapes :

1. Étape d'Entraînement : Spécialisation Indépendante des Enseignants

Au lieu d'une optimisation conjointe, la méthode entraîne d'abord un ensemble de modèles enseignants spécifiques à chaque tâche de manière indépendante. Chaque enseignant est optimisé pour une seule tâche (par exemple, la reconnaissance optique de caractères (OCR), l'esthétique, l'alignement compositionnel) en utilisant des algorithmes de RL standards et prêts à l'emploi. Cela garantit que chaque enseignant atteint un plafond de performance élevé sur son objectif spécifique sans interférence des autres tâches.

2. Étape d'Intégration : Distillation en Ligne

Un modèle étudiant unifié est ensuite entraîné pour distiller les capacités de ces enseignants spécialisés. Crucialement, l'étudiant apprend le long de ses propres trajectoires de déploiement (rollout) (en ligne).

• Processus : L'étudiant génère une trajectoire de débruitage. Pour chaque étape de cette trajectoire, l'enseignant spécifique à la tâche correspondante fournit une supervision.
• Objectif : L'étudiant est entraîné à correspondre au comportement de transition de l'enseignant à chaque étape de débruitage.

Déduction Théorique : Du Discret au Continu

La contribution théorique centrale consiste à élever le cadre OPD des transitions de tokens discrets (LLM) vers des processus de Markov à état continu (Modèles de Diffusion).

• Formulation de Chaîne de Markov : Le processus de débruitage par diffusion est modélisé comme une chaîne de Markov à temps discret induite par une équation différentielle stochastique (SDE) à temps inverse.
• Transitions Gaussiennes : L'étudiant et l'enseignant définissent tous deux des noyaux de transition gaussiens à un pas. Puisqu'ils partagent la même covariance (déterminée par le planificateur et le niveau de bruit) mais diffèrent par leurs moyennes, la divergence de Kullback-Leibler (KL) inverse admet une expression analytique fermée :
  $$\text{KL}(p_S \parallel p_T) = \frac{\|\mu_S - \mu_T\|^2}{2\sigma^2}$$
  où $\mu_S$ et $\mu_T$ sont respectivement les moyennes de transition de l'étudiant et de l'enseignant.
• Optimisation : L'étudiant est optimisé en minimisant cet objectif KL à forme fermée par rétropropagation directe.

Comparaison avec les Gradients de Style PPO

L'article oppose cette approche à forme fermée à une formulation de gradient de politique de style PPO (traitant l'enseignant comme un modèle de récompense de processus).

• Réduction de la Variance : Bien que les gradients de style PPO soient sans biais en espérance, ils introduisent un terme de fonction de score proportionnel au bruit gaussien, conduisant à une variance de gradient plus élevée.
• Compatibilité Déterministe : L'objectif KL à forme fermée s'étend naturellement aux échantillonneurs ODE déterministes (où la transition est une correspondance directe en $L_2$), alors que les objectifs de style PPO reposent sur des densités de politique stochastiques.
• Conclusion : Le KL à forme fermée fournit un principe d'optimisation à variance plus faible et plus général, applicable aux échantillonneurs SDE et ODE.

Contributions Clés

1. Cadre DiffusionOPD : Un nouveau paradigme qui découple l'exploration par tâche unique de l'intégration multi-tâches, utilisant des enseignants spécifiques à chaque tâche pour superviser un étudiant unifié le long de ses propres trajectoires.
2. Unification Théorique : Dérivation d'un objectif KL par étape à forme fermée et fondé sur des principes pour les modèles de diffusion, unifiant le raffinement stochastique SDE et déterministe ODE, offrant une optimisation à variance plus faible que les gradients de style PPO.
3. Validation Empirique : Des expériences extensives démontrant que DiffusionOPD surpasse à la fois les bases de RL multi-tâches (optimisation conjointe) et les bases de RL en cascade en termes d'efficacité d'entraînement et de performance finale.

Résultats Expérimentaux

Les auteurs ont évalué DiffusionOPD sur divers benchmarks incluant GenEval (génération compositionnelle), OCR (rendu de texte) et Esthétique (qualité visuelle), en utilisant SD3.5-Medium comme modèle de base.

• Performance : DiffusionOPD a obtenu des résultats State-of-the-Art (SOTA) sur tous les benchmarks évalués, surpassant Multi-Task GRPO-Guard, Multi-Task NFT et Cascade NFT.
  • Score Moyen : DiffusionOPD (0,929) contre Cascade NFT (0,851) contre Multi-Task NFT (0,715).
• Efficacité : La méthode a démontré une convergence significativement plus rapide et un plafond de performance plus élevé par rapport aux bases de RL multi-tâches.
• Temps d'Entraînement : DiffusionOPD a nécessité moins de temps total d'horloge pour atteindre les scores cibles par rapport aux méthodes de RL conjointes, tout en évitant la surcharge séquentielle des méthodes en cascade.
• Études d'Ablation :
  • Formulation de la Perte : L'objectif KL à forme fermée a systématiquement surpassé les gradients de politique de style PPO tant en vitesse de convergence qu'en récompense finale.
  • Niveau de Bruit : Des niveaux de bruit plus faibles (s'approchant de l'échantillonnage ODE) ont conduit à une convergence plus rapide et à de meilleures performances, l'échantillonneur ODE étant jusqu'à cinq fois plus efficace que l'échantillonneur SDE avec un bruit élevé.

Importance et Revendications

L'article affirme que DiffusionOPD répond aux limitations fondamentales du RL multi-tâches actuel pour les modèles de diffusion en découplant l'exploration de l'intégration. En tirant parti des propriétés analytiques des transitions de diffusion, il fournit un principe d'entraînement unifié et à faible variance applicable à différents types d'échantillonneurs (SDE et ODE).

Les auteurs positionnent DiffusionOPD non pas simplement comme une méthode de distillation spécifique, mais comme un cadre général pour la distillation en ligne dans les modèles de diffusion. Ils affirment que cette approche résout efficacement le compromis entre le conflit d'optimisation (dans l'entraînement conjoint) et l'oubli catastrophique/l'inefficacité (dans l'entraînement en cascade), établissant une nouvelle fondation pour la modélisation par diffusion multi-tâches et alignée sur les préférences.