Résumé technique : Apprentissage par renforcement guidé par un discriminateur pour le Flow Matching

Énoncé du problème

Les modèles basés sur les flux et le score-matching sont devenus le paradigme dominant de la modélisation générative en raison de leurs objectifs d'entraînement par régression sans simulation. Cependant, ces modèles sont généralement entraînés selon un pipeline multi-étapes : d'abord, un modèle de base est ajusté à une distribution de données $q$ via le flow ou le score matching (FSM) ; ensuite, l'apprentissage par renforcement (RL) est appliqué pour orienter le modèle vers des régions désirables définies par une récompense $r(x)$.

L'article identifie un décalage structurel dans ce paradigme. Bien que les pertes de FSM minimisent l'erreur de régression $\ell_2$ sur les champs de vitesse ou de score sous les marginales d'entraînement, elles échouent souvent à récupérer les propriétés intrinsèques de la distribution de données — telles que le réalisme visuel, la cohérence de la structure des objets ou le mouvement physiquement correct — qui sont pourtant théoriquement présentes dans les données. Par conséquent, les praticiens ont recours au RL basé sur les préférences pour récupérer ces propriétés. Cette approche est problématique pour deux raisons :

1. Échec structurel : Il est déroutant que les modèles ne puissent pas apprendre ces propriétés directement à partir des données via des pertes de correspondance (matching losses), mais nécessitent le RL pour les récupérer.
2. Coût et confusion : L'entraînement des récompenses nécessite généralement des données de préférence humaine coûteuses, ce qui confond le réalisme objectif des données avec les inclinations subjectives des annotateurs.

Les auteurs soutiennent que l'échec des objectifs de correspondance à capturer ces propriétés provient d'un « décalage train-test » (où les erreurs s'accumulent le long de la distribution de déploiement du modèle plutôt que le long des marginales d'entraînement) et d'un « décalage géométrique » (où de petites erreurs dans l'espace de vitesse peuvent entraîner de grandes erreurs dans l'espace de qualité d'échantillonnage/récompense).

Méthodologie : Apprentissage par renforcement guidé par un discriminateur (DRL)

L'article propose le Discriminator-Guided RL (DRL), une méthode pour corriger les modèles de flow matching sans dépendre de données de préférence humaine. L'idée centrale est que la récompense optimale pour aligner un modèle de base $p_{base}$ avec une distribution cible $q$ est le rapport de densité logarithmique $\log(q(x)/p_{base}(x))$. Comme ce rapport est intraçable, le DRL l'estime à l'aide d'un discriminateur.

Le pipeline DRL se compose de deux étapes :

Étape 1 : Estimation de la récompense dans l'espace de représentation

Au lieu d'entraîner un discriminateur directement dans l'espace des pixels (ce qui est statistiquement difficile et sémantiquement peu fiable), le DRL entraîne un discriminateur dans un espace de représentation d'apprentissage auto-supervisé (SSL) pré-entraîné.

• Un encodeur gelé $\phi$ (ex: DINOv2) projette les images dans un espace de caractéristiques $Z$.
• Un discriminateur $D$ est entraîné pour distinguer les échantillons de données réelles $x \sim q$ des échantillons du modèle de base $x_{fake} \sim p_{base}$ en utilisant leurs représentations encodées $\phi(x)$.
• La fonction de récompense est définie comme le logit de ce discriminateur : $\hat{r}(x) = \log(D(\phi(x)) / (1 - D(\phi(x))))$.
• Justification théorique : En restreignant la récompense à l'espace SSL, le DRL résout un problème d'optimisation contraint : il trouve la distribution $p$ la plus proche de $p_{base}$ (en divergence KL) telle que la mesure directe de $p$ à travers $\phi$ corresponde à la mesure directe de $q$ à travers $\phi$. Cela garantit que les corrections sont effectuées le long d'axes sémantiquement significatifs visibles pour l'encodeur.

Étape 2 : RL régularisé par KL via l'Adjoint Matching

La récompense estimée $\hat{r}$ est utilisée pour affiner le modèle de base via un RL régularisé par la divergence KL.

• L'objectif est de maximiser la récompense attendue tout en minimisant la divergence KL entre les distributions de trajectoire du modèle affiné et du modèle de base.
• Les auteurs utilisent l'Adjoint Matching, un algorithme de pointe pour les modèles de flux qui optimise l'objectif au niveau de la trajectoire par propagation de gradient à travers le processus d'échantillonnage.
• Pour assurer la stabilité de l'entraînement avec l'SDE sans mémoire requise pour le basculement (tilting) non biaisé de l'extrémité, les auteurs introduisent un Intégrateur Linéaire Local. Cet intégrateur gère l'aspect raide (stiffness) numérique près de $t=0$ (où les coefficients de dérive et de diffusion divergent) plus efficacement que les discrétisations Euler-Maruyama standards avec décalages (offsets).

Contributions clés

1. Analyse théorique des limites du matching : L'article fournit une analyse formelle montrant que les pertes de flow/score matching ne garantissent pas le contrôle des écarts de récompense (propriétés d'intérêt) en raison des décalages de distribution train-test et du désalignement entre la géométrie de l'espace de vitesse et la géométrie de la récompense de l'espace d'échantillonnage.
2. Apprentissage par renforcement guidé par un discriminateur (DRL) : Un nouveau cadre qui utilise un discriminateur entraîné dans un espace de représentation SSL gelé pour générer un signal de récompense. Cela permet au modèle de récupérer les propriétés des données sans étiquettes de préférence humaine.
3. Intégrateur Linéaire Local : Un nouvel intégrateur numérique pour l'SDE sans mémoire utilisé dans l'entraînement par RL, qui empêche l'explosion de la norme aux temps précoces et améliore la qualité de l'échantillonnage ainsi que la stabilité de l'entraînement.
4. Validation empirique : Démonstration que le DRL améliore significativement l'alignement de la distribution et la qualité d'image à travers diverses architectures (SiT, JiT, REPA, RAE) et espaces de caractéristiques (DINOv2, DINOv3, SigLIP, Inception).

Résultats

Les auteurs ont évalué le DRL sur quatre modèles de flux pré-entraînés sur ImageNet. Principales conclusions :

• Alignement de la distribution : Le DRL réduit considérablement la distance de Fréchet (FD). Par exemple, sur SiT, la FD sans guidage dans l'espace DINOv3 est passée de 88,2 à 19,3. Sur tous les modèles et espaces de caractéristiques, le DRL améliore systématiquement l'alignement sans utiliser de CFG (Classifier-Free Guidance).
• Qualité d'image sans données de préférence : Le DRL a amélioré les scores sur des modèles de récompense basés sur les préférences humaines tenus à l'écart (ImageReward, HPSv2, PickScore, Aesthetics) sans jamais avoir été entraîné sur eux. Cela suggère que le DRL récupère avec succès les propriétés de "réalisme" que les pertes de correspondance standard manquent.
• Base améliorée pour l'alignement de préférence : Lorsque le DRL est utilisé comme point de départ pour un RL ultérieur basé sur les préférences (PRL), il produit une frontière de Pareto supérieure entre la récompense et la fidélité de l'image. Le DRL absorbe les corrections structurelles récupérables à partir des données, permettant au PRL de se concentrer sur les véritables préférences subjectives plutôt que sur la correction d'artefacts de bas niveau (ex: sursaturation, luminosité excessive).
• Écart de distillation : Les expériences ont montré qu'un modèle étudiant entraîné via un flow matching standard sur des échantillons provenant d'un enseignant DRL ne pouvait pas reproduire les performances de l'enseignant, même après avoir vu 50 millions d'échantillons. Cela confirme que l'objectif de RL lui-même, et non seulement les échantillons, est nécessaire pour récupérer ces propriétés.

Signification et affirmations

L'article affirme que le post-entraînement par RL pour les modèles de flux doit être perçu non pas simplement comme un outil d'optimisation des préférences externes, mais comme un mécanisme complémentaire pour récupérer la structure des données d'entraînement qui est imparfaitement capturée par les objectifs de correspondance standard.

En utilisant un discriminateur dans un espace de représentation pré-entraîné, le DRL fournit un signal de récompense qui est :

1. Aligné avec les propriétés des données : Il cible directement la distribution des données plutôt que les préférences humaines subjectives.
2. Géométriquement sain : Il contourne les obstructions de la géométrie de l'espace de vitesse en optimisant directement sur le paysage de récompense.
3. Efficace : Il nécessite nettement moins de calcul que le pré-entraînement et évite le coût de la collecte de données de préférence.

Les auteurs concluent que les objectifs de correspondance et le RL ne sont pas des méthodes aux objectifs conflictuels, mais plutôt deux outils complémentaires pour maximiser l'utilité des données d'entraînement. Ce travail suggère que la "récompense" pour une génération de haute qualité est souvent "déjà présente dans les données", mais nécessite un mécanisme sur-politique (on-policy) guidé par la récompense pour être pleinement réalisée.