Résumé Technique : Suivre la Moyenne : Appariement de Flux Guidé par Référence

Énoncé du Problème

Le contrôle des sorties des modèles génératifs préentraînés (par exemple, l'imposition d'attributs, de styles ou de structures spécifiques) reste un défi majeur. Les approches existantes se répartissent généralement en trois catégories, chacune présentant des limitations distinctes :

1. Affinage et Adaptateurs : Ils nécessitent la modification des paramètres du modèle pour chaque nouvelle cible, ce qui est coûteux en calcul et peu flexible face aux changements dynamiques.
2. Méthodes de Guidage : Elles laissent le générateur figé mais reposent sur des classificateurs auxiliaires, des modèles de récompense ou des signaux de récompense, introduisant une surcharge computationnelle supplémentaire et une dépendance à des réseaux externes.
3. Méthodes Basées sur la Recherche : Elles évitent le réentraînement mais engendrent des coûts d'inférence élevés par un échantillonnage répété, un filtrage ou une optimisation par invite.

Aucune de ces approches ne parvient simultanément à éviter l'entraînement supplémentaire, les réseaux auxiliaires et la recherche au moment du test tout en maintenant une génération de haute qualité. Cet article propose une formulation alternative : Flux Guidés par Référence, qui contrôle la génération en adaptant le modèle via des exemples de données plutôt que par des mises à jour de paramètres ou des signaux de récompense externes.

Méthodologie

L'idée centrale de l'article est que, dans l'Appariement de Flux (Flow Matching - FM), le champ de vitesse $u_t(x)$ est uniquement régi par la moyenne conditionnelle de l'extrémité $\mu_t(x) = \mathbb{E}[x_1 | x_t = x]$. Pour un interpolant linéaire, la vitesse est donnée par :
$$u_t(x) = \frac{\mu_t(x) - x}{1 - t}$$
Par conséquent, décaler la moyenne de l'extrémité décale l'ensemble de la trajectoire du flux. Les auteurs proposent que cette moyenne puisse être décalée en conditionnant sur un ensemble de référence $R$ d'images, orientant efficacement le modèle vers la distribution induite par $R$ sans modifier les poids du modèle.

L'article met en œuvre ce principe sous deux formes :

1. Guidage par Moyenne de Référence (RMG)

Le RMG est une méthode sans entraînement pour contrôler des modèles préentraînés figés.

• Mécanisme : Il calcule une correction sous forme close du champ de vitesse basée sur la différence entre la moyenne de l'extrémité de l'ensemble de référence ($\hat{\mu}^\rho_t$) et la moyenne de l'extrémité du modèle préentraîné ($\mu^\theta_t$).
• Formule : La vitesse guidée est définie comme suit :
  $$u^\pi_t(x) \approx u^\theta_t(x) + \beta_t \frac{\hat{\mu}^\rho_t(x) - \mu^\theta_t(x)}{1 - t}$$
  où $\beta_t$ est un calendrier de guidage et $\hat{\mu}^\rho_t(x)$ est la moyenne pondérée empirique de l'ensemble de référence (calculée via un softmax sur les images de référence).
• Propriétés : Il ne nécessite aucune computation de gradient, aucun réseau auxiliaire et aucune évaluation supplémentaire du modèle au-delà du modèle de base. L'ensemble de référence peut être échangé au moment de l'inférence pour modifier la couleur, l'identité, le style ou la structure.

2. Guidage Semi-Paramétrique (SPG)

Le SPG est une architecture apprenable qui amortit le concept de guidage par moyenne de référence.

• Architecture : Le modèle $\mu^\theta_t(x_t, R)$ conditionne explicitement sur un ensemble de référence $R$ pendant l'entraînement. Il utilise un mécanisme d'attention croisée pour calculer une « ancre » $\bar{x}$ (analogue à la moyenne sous forme close) et un raffineur résiduel appris $f^\theta$ pour corriger les corrélations parasites (par exemple, des arrière-plans partagés) que l'ancre pourrait involontairement copier.
• Entraînement : Le modèle est entraîné avec un objectif de type « leave-one-out » pour empêcher l'ancre de s'attarder sur sa propre extrémité, combiné à une perte résiduelle pour apprendre au raffineur à capturer des détails au-delà de l'ancre.
• Propriétés : Le SPG correspond à la qualité de génération non conditionnelle des modèles standards (par exemple, DiT-B/4) tout en permettant un contrôle au moment de l'inférence par substitution d'ensemble de référence.

Contributions Clés

1. Décalage de Moyenne comme Guidage : Les auteurs définissent le guidage comme une interpolation géométrique entre le chemin de probabilité de l'ensemble d'entraînement et celui d'un ensemble de référence. Ils prouvent que le terme de guidage résultant dépend uniquement de la différence entre les moyennes d'extrémité des deux distributions.
2. Guidage par Moyenne de Référence (RMG) : Une méthode sans entraînement qui calcule le décalage de vitesse sous forme close en utilisant une banque de référence. Elle permet le contrôle de modèles figés (démontré sur FLUX.2-klein) sans gradients, classificateurs ni modèles de récompense.
3. Guidage Semi-Paramétrique (SPG) : Une architecture contrôlable qui apprend à approximer l'ancre de moyenne de référence et une correction résiduelle. Elle préserve la qualité de génération non conditionnelle tout en permettant l'échange de l'ensemble de référence au moment de l'inférence, supprimant efficacement les corrélations parasites comme les arrière-plans partagés.

Résultats

L'article valide l'approche par une analyse mécaniste et des expériences extensives :

• Validation Mécaniste : Sur la distribution à deux lunes et sur MNIST, les auteurs démontrent que modifier la composition de l'ensemble de référence décale directement le champ de flux et l'attracteur final, isolant le rôle causal de la moyenne a posteriori.
• Contrôle Sans Entraînement (RMG) : Appliqué à un modèle FLUX.2-klein (4B) figé, le RMG oriente avec succès la génération pour la couleur (éléphants roses/bleus), l'identité (girafes/zebras), le style (Van Gogh/sketches) et la structure (formes de serrure, poses de mains).
  • Benchmarks GenEval : Le RMG surpasse les méthodes basées sur la recherche (Best-of-4, SMC) et les méthodes basées sur le gradient (ReNO) pour des tâches de texte-à-image compositionnelles (par exemple, position, génération de deux objets) tout en utilisant une seule trajectoire d'échantillonnage et aucun modèle auxiliaire. Il obtient un score moyen de 91,17 contre 80,10 pour la base de référence, avec un temps d'inférence significativement plus faible et aucun appel à un modèle externe.
• Guidage Semi-Paramétrique (SPG) : Entraîné sur AFHQv2, le SPG correspond à la qualité non conditionnelle d'une base de référence DiT-B/4 (FID 23,26 contre 23,11).
  • Contrôle : L'échange de l'ensemble de référence (par exemple, uniquement des chats vs uniquement des chiens) décale systématiquement la distribution de sortie. Les proportions de classes générées suivent étroitement la composition de l'ensemble de référence.
  • Suppression d'Artéfacts : Contrairement au RMG, qui peut transférer des corrélations parasites (par exemple, un arrière-plan blanc provenant de toutes les images de référence), le raffineur résiduel du SPG préserve avec succès le guidage au niveau de l'objet tout en supprimant de tels artéfacts.

Importance et Revendications

L'article soutient que les modèles génératifs peuvent être adaptés par les données plutôt que par des mises à jour de paramètres. En présentant le contrôle comme un problème de décalage des moyennes d'extrémité, les auteurs offrent une alternative simple et fondée sur des principes à l'affinage, au guidage auxiliaire ou à la recherche.

• Pour les Modèles Figés : Le RMG démontre qu'un modèle préentraîné peut être orienté efficacement en modifiant l'ensemble de référence qu'il suit implicitement, offrant un paradigme de « guider par des exemples, non par des récompenses ».
• Pour les Modèles Apprenables : Le SPG montre que ce mécanisme peut être amorti dans une architecture apprenable qui maintient une haute qualité de génération tout en permettant un contrôle flexible au moment de l'inférence.
• Direction Plus Large : Ce travail suggère une voie vers des modèles génératifs qui s'adaptent dynamiquement via des références de données, réduisant le besoin de réentraînement ou d'ingénierie complexe de récompenses pour une génération contrôlable.

Les auteurs reconnaissent les limites, notant que le RMG hérite de la qualité de son ensemble de référence (des références bruyantes introduisent des artéfacts) et que le calcul des moyennes sur de grands ensembles peut être coûteux, bien que l'échantillonnage subséquent offre des remèdes. Ils soulignent la nécessité d'une curation responsable des ensembles de référence pour prévenir les mauvais usages.