DiffusionOPD: Единый взгляд на дистилляцию по политике в диффузионных моделях

Аннотация

Обучение с подкреплением стало мощным инструментом для улучшения диффузионных моделей типа «текст-изображение», однако существующие методы в значительной степени ограничены оптимизацией отдельной задачи. Расширение обучения с подкреплением на множество задач представляет собой сложную задачу: совместная оптимизация страдает от межзадачной интерференции и дисбаланса, а каскадное обучение с подкреплением громоздко и подвержено катастрофическому забыванию. Мы предлагаем DiffusionOPD — новую парадигму многозадачного обучения для диффузионных моделей, основанную на онлайн-дистилляции политик (Online Policy Distillation, OPD). DiffusionOPD сначала независимо обучает учителей для отдельных задач, а затем дистиллирует их способности в единого студента вдоль его собственных траекторий развёртывания. Это разделяет исследование отдельных задач и многозадачную интеграцию, избегая бремени оптимизации, связанного с решением всех задач с нуля совместно. Теоретически мы обобщаем фреймворк OPD с дискретных токенов на марковские процессы с непрерывным состоянием, выводя аналитическое выражение для пошаговой цели KL, которая объединяет как стохастическое уточнение с помощью СДУ, так и детерминированное уточнение с помощью ОДУ через согласование средних. Мы формально и эмпирически демонстрируем, что этот аналитический градиент обеспечивает меньшую дисперсию и лучшую обобщаемость по сравнению с традиционными градиентами политик в стиле PPO. Обширные эксперименты показывают, что DiffusionOPD неизменно превосходит как базовые методы многозадачного обучения с подкреплением с несколькими наградами, так и каскадного обучения с подкреплением по эффективности обучения и итоговой производительности, достигая при этом самых современных результатов на всех оценённых эталонных тестах.

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.