🤖 machine learning

Diffusion Fine-tuning with Rewarded Moment Matching Distillation

Cet article introduit la Distillation par Appariement de Moments Récompensés (RMMD), un nouveau cadre qui unifie la distillation de modèles de diffusion et l'apprentissage par renforcement pour atteindre des compromis vitesse-qualité supérieurs, démontrés par des améliorations significatives tant sur ImageNet que sur le modèle de prévision météorologique haute dimension GenCast.

Auteurs originaux : Alexis Jacq, Guillaume Couairon, Valentin De Bortoli, Quentin Berthet, Arnaud Doucet, Romuald Elie

Publié 2026-06-30

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Alexis Jacq, Guillaume Couairon, Valentin De Bortoli, Quentin Berthet, Arnaud Doucet, Romuald Elie

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous ayez un chef étoilé (le Modèle Enseignant) capable de cuisiner un plat complexe et parfait, mais qu'il lui faut 59 heures pour le faire. Il est lent, mais la nourriture est délicieuse et précise. Vous voulez un sous-chef (le Modèle Étudiant) qui puisse cuisiner le même plat en seulement quelques minutes, mais vous voulez aussi qu'il ajuste légèrement la recette pour la rendre plus épicée ou plus saine (optimisation pour une Récompense).

Le problème est que si vous dites simplement au sous-chef de « rendre le plat plus épicé » alors qu'il est encore en train d'apprendre les bases, il risque de rater complètement le plat ou d'oublier comment cuisiner correctement.

Ce document présente une nouvelle méthode d'entraînement appelée Rewarded Moment Matching Distillation (RMMD). Considérez cela comme un camp d'entraînement en deux phases qui résout ce problème.

Phase 1 : La phase d'« Observation » (Distillation)

D'abord, le sous-chef passe du temps à observer le chef étoilé. Il ne se contente pas de regarder le plat final ; il observe chaque étape du processus de cuisson.

L'analogie : Imaginez que le chef épluche un oignon couche par couche. L'étudiant apprend à prédire l'aspect de l'oignon à chaque étape de l'épluchage, et pas seulement le résultat final.
Le résultat : L'étudiant apprend à cuisiner une version presque parfaite du plat en seulement 8 étapes (au lieu de 59), tout en préservant le « naturel » et la qualité de l'original. C'est ce qu'on appelle le Moment Matching.

Phase 2 : La phase de « Test de Goût » (Affinage par Récompense)

Maintenant que l'étudiant est un cuisinier habile, vous voulez qu'il ajuste la saveur (par exemple, rendre le plat plus rouge, ou dans le cas du document, rendre les prévisions météorologiques plus précises).

Le problème des anciennes méthodes : Les méthodes précédentes tentaient d'apprendre à l'étudiant à changer la saveur en observant une version « re-bruité » d'un plat qu'il n'avait pas réellement cuisiné lui-même (off-policy). C'est comme dire à un chef d'améliorer une recette en se basant sur la photo d'un plat préparé par quelqu'un d'autre. L'étudiant s'embrouille car les ingrédients ne correspondent pas à ceux qu'il utilise habituellement.
La solution RMMD : Le document utilise une approche On-Policy. L'étudiant cuisine un plat, puis l'enseignant ajoute un peu de « bruit » (comme saupoudrer des épices aléatoires) et l'étudiant tente de le corriger à nouveau.
Le tour de magie : Pendant que l'étudiant essaie de corriger le plat pour le rendre « plus épicé » (maximiser la récompense), le système lui murmure également : « Hé, n'oublie pas de bien éplucher l'oignon ! ». Il utilise les leçons d'« Observation » de la Phase 1 comme un filet de sécurité. Cela empêche l'étudiant de devenir fou et de ruiner le plat juste pour obtenir la récompense.

Pourquoi est-ce spécial ?

C'est un équilibre : Le document montre que vous pouvez rendre le plat plus rapide et meilleur en même la fois. Les anciennes méthodes vous forçaient généralement à choisir : soit garder la qualité mais rester lent, soit devenir rapide mais gâcher le goût. Le RMMD trouve le point d'équilibre idéal.
Cela fonctionne sur la science réelle : Les auteurs n'ont pas seulement testé cela sur des photos de chats et de chiens. Ils l'ont appliqué à GenCast, un modèle de prévision météorologique extrêmement complexe.
- L'Enseignant : Prend 59 étapes pour prédire la météo des 12 prochaines heures.
- L'Étudiant RMMD : Prend seulement 8 étapes (ce qui le rend 7,5 fois plus rapide).
- Le Résultat : L'étudiant rapide n'a pas seulement été plus rapide ; il a en fait prédit la météo mieux que le maître lent pour 93 % des variables (comme la température et le vent). Il a également corrigé un problème courant où les modèles météorologiques sont trop confiants (sous-dispersés), rendant les prévisions plus fiables.

L'essentiel à retenir

Le document affirme que le RMMD est une manière plus intelligente d'entraîner les modèles d'IA. Il leur apprend à être rapides sans oublier d'être précis, et il leur apprend à suivre de nouveaux objectifs (comme « être plus rouge » ou « être plus précis ») sans briser les règles fondamentales de la génération de données.

En bref : c'est comme entraîner un pilote de course qui peut rouler à 200 mph (rapide) mais qui sait exactement comment rester sur la piste (précis) et peut naviguer sur un nouvel itinéraire (récompense) sans s'écraser.

Résumé Technique : Fine-tuning de Diffusion avec la Distillation par Appariement de Moments Récompensée (RMMD)

Énoncé du Problème

Les modèles de diffusion sont devenus la norme pour la synthèse d'images haute fidélité et les prévisions scientifiques en raison de leurs objectifs d'entraînement stables. Cependant, le déploiement pratique est entravé par le coût de calcul élevé de l'inférence, qui nécessite des dizaines ou des centaines d'étapes de débruitage. Bien que les techniques de distillation permettent de compresser ces modèles en générateurs à peu d'étapes, et que le fine-tuning par apprentissage par renforcement (RL) les aligne avec des fonctions de récompense spécifiques (par exemple, les préférences humaines ou les métriques scientifiques), la combinaison de ces deux phases reste difficile.

Les approches existantes pour fusionner la distillation et l'optimisation de la récompense font face à des limitations significatives :

Fragilité du l'entraînement conjoint : Fusionner la distillation et le fine-tuning de récompense en une seule phase provoque souvent une dérive des échantillons générés hors de la distribution d'entrée du modèle enseignant, invalidant ainsi le signal de distillation.
Problèmes de mémoire et de biais : Le fine-tuning d'un modèle déjà distillé via une rétropropagation multi-étapes (par exemple, DRaFT) est gourmand en mémoire. Tronquer la rétropropagation (par exemple, ReFL) introduit un biais en évaluant les récompenses sur des latents bruités.
Contraintes structurelles : Des méthodes comme HyperNoise évitent la rétropropagation en chaîne en perturbant le bruit initial, mais sont structurellement confinées aux changements d'image à basse fréquence, ce qui limite leur capacité à optimiser des récompenses structurelles complexes à haute fréquence.

Méthodologie : Distillation par Appariement de Moments Récompensée (RMMD)

Les auteurs proposent RMMD, un cadre en deux phases qui connecte de manière fondamentale la distillation et le fine-tuning de récompense tout en préservant la « naturalité » de haute fidélité de la distillation avancée.

Phase 1 : Distillation par Appariement de Moments (MMD)

Le processus commence par la distillation d'un modèle enseignant de base en utilisant la Distillation par Appariement de Moments (MMD). Cette méthode fait correspondre les distributions de débruitage intermédiaires le long de la trajectoire d'échantillonnage, produisant un modèle étudiant multi-étapes qui suit étroitement les marginales de l'enseignant. Cet étudiant est ensuite gelé pour servir de référence distributionnelle stable.

Phase 2 : Fine-tuning de Récompense On-Policy

Dans la seconde phase, le modèle étudiant est affiné pour maximiser une fonction de récompense tout en maintenant la fidélité distributionnelle. Les innovations clés incluent :

Échantillonnage On-Policy : Au lieu d'utiliser des points de données off-policy, la méthode échantillonne des états bruités $x^{pol}_t$ à partir de la politique actuelle de l'étudiant (en faisant fonctionner l'étudiant pendant $K_{student}$ étapes, en re-bruitant la sortie, et en arrêtant le gradient). Cela garantit que l'objectif de récompense approxime directement la récompense attendue de l'étudiant et empêche la dérive distributionnelle d'invalider le signal d'entraînement.
Régularisation Recyclée : La contribution centrale est la réutilisation de la perte d'appariement de moments comme régularisateur pendant le fine-tuning. Plutôt qu'une pénalité heuristique, cette perte pénalise explicitement l'écart entre les distributions de débruitage intermédiaires de l'étudiant et celles du modèle de référence gelé. Cela fournit un levier interprétable (le poids de régularisation) pour arbitrer entre la maximisation de la récompense et la fidélité distributionnelle.
Gradient à étape unique : La méthode capture toute la gamme de fréquences sans déroulement multi-étapes en effectuant une seule étape de débruitage sur des échantillons corrompus on-policy et en évaluant la récompense.
Régularisation L2 : Un terme L2 supplémentaire pénalise les écarts significatifs par rapport aux prédictions du modèle initialement distillé par MMD pour stabiliser davantage l'entraînement.

L'objectif total combine le gradient de récompense on-policy avec le régularisateur d'appariement de moments et le terme L2 :
$\nabla_\theta \mathcal{L}_{RMMD} = \mathbb{E} \left[ \left(\frac{\partial \hat{x}_0}{\partial \theta}\right)^\top \left( m_{student} - m_{teacher} - \lambda \nabla_{\hat{x}_0} R(\hat{x}_0) \right) \right] + \lambda_{reg} \mathcal{L}_{L2}$

Contributions Clés

Procédure RMMD : Une nouvelle méthode de distillation de diffusion qui maximise simultanément une fonction de récompense, en utilisant une version on-policy de la perte d'appariement de moments pour la régularisation.
Supériorité Empirique : Démontre que RMMD surpasse systématiquement les baselines à une étape (DI++) et les concurrents multi-étapes (DRaFT, HyperNoise) sur les fronts de Pareto FID-Récompense à travers diverses fonctions de récompense, allant des métriques au niveau du pixel à l'alignement sémantique (CLIP).
Application Scientifique : Application réussie de RMMD à GenCast, un modèle de prévision météorologique basé sur la diffusion de pointe, optimisant le Score de Probabilité Classique (CRPS).

Résultats

Expériences ImageNet

Avantage Multi-étapes : L'utilisation d'un étudiant distillé par MMD à 8 étapes comme point de départ a nettement surpassé les baselines à 1 étape (DI++). L'étudiant RMMD à 8 étapes a obtenu des compromis supérieurs sur les fronts de Pareto FID-Récompense.
Comparaison : RMMD a généré de meilleurs fronts de Pareto que DRaFT et HyperNoise à travers diverses résolutions (64x64 et 512x512) et types de récompenses (alignement CLIP, Inception Score, lissage).
Analyse Qualitative : Les visualisations ont montré que si DRaFT-1 préservait le contenu mais introduisait des artefacts adverses, et que DRaFT-2 souffrait de détournement de récompense (changements de distribution drastiques), RMMD a réussi à intégrer des modifications subtiles pour augmenter les récompenses sans trop dévier de la distribution de données originale.

Prévisions Météorologiques GenCast

Accélération : Le modèle distillé a obtenu une accélération de 7,5× (8 étapes contre 59 évaluations de fonction) par rapport à l'enseignant.
Précision : Le modèle RMMD optimisé avec le CRPS a surpassé le modèle enseignant sur 93 % des variables météorologiques cibles (contre 75 % pour le meilleur baseline MMD seul).
Calibration : RMMD a traité le problème de sous-dispersion courant dans la distillation MMD. Bien que la variante on-policy présente une dispersion légèrement moins bonne que la variante off-policy, elle offre les meilleures améliorations globales du CRPS.
Généralisation : L'optimisation du CRPS à un horizon de 12 heures a entraîné des améliorations constantes pour les prévisions allant jusqu'à 7 jours, suggérant que la méthode améliore la modélisation distributionnelle sous-jacente plutôt que de simplement faire du surapprentissage sur la récompense immédiate.

Signification et Revendications

L'article affirme que RMMD fournit une solution principalisée au problème de « dérive » inhérent à la distillation et au fine-tuning conjoints. En recyclant la perte d'appariement de moments comme régularisateur et en utilisant l'échantillonnage on-policy, la méthode permet l'optimisation d'objectifs complexes et différentiables (comme le CRPS dans les prévisions météorologiques) sans sacrifier la qualité générative ou la stabilité du modèle distillé.

Les auteurs soulignent que cette approche s'adapte à des domaines scientifiques complexes et de haute dimension. L'application réussie à GenCast démontre que les modèles de diffusion peuvent être à la fois considérablement accélérés et améliorés dans leurs prédictions probabilistes, ouvrant la voie à des modèles de diffusion plus efficaces et plus précis dans le calcul scientifique où existent des objectifs différentiables. Le travail suggère que le maintien d'une référence distributionnelle stable via l'appariement de moments est crucial pour un fine-tuning de récompense efficace dans les modèles de diffusion.