How to Guide Your Flow: Few-Step Alignment via Flow Map Reward Guidance

1. Énoncé du problème

Dans la modélisation générative (spécifiquement les modèles basés sur les flots et les modèles de diffusion), il existe un besoin critique de générer des échantillons satisfaisant des critères spécifiques définis par l'utilisateur, tels que la qualité esthétique, les préférences humaines ou les contraintes physiques (problèmes inverses). Ce processus est connu sous le nom de guidage.

Les méthodes de guidage actuelles font face à deux limitations principales :

• Intratabilité computationnelle : L'approche théoriquement rigoureuse consiste à échantillonner à partir d'une distribution « inclinée par la récompense » ($\tilde{\rho} \propto e^{r(x)}\rho(x)$). L'échantillonnage exact à partir de cette distribution nécessite généralement des méthodes de Monte Carlo séquentielles (SMC) avec de nombreuses particules et étapes de rééchantillonnage, ce qui est computationnellement prohibitif pour des données de haute dimension comme les images.
• Erreurs d'approximation : Les méthodes à trajectoire unique (comme le Diffusion Posterior Sampling, DPS) évitent le coût des particules mais reposent sur des approximations non contrôlées du score incliné par la récompense. Ces méthodes ont souvent du mal à s'aligner avec les échantillonneurs déterministes modernes (comme les Flow Maps) visant une inférence rapide (peu d'étapes), car elles ont été initialement dérivées de cadres stochastiques.

La question centrale abordée est la suivante : Existe-t-il un cadre principiel pour guider les modèles génératifs basés sur les flots en utilisant très peu d'évaluations de fonctions (NFEs) et une seule trajectoire ?

2. Méthodologie : Flow Map Reward Guidance (FMRG)

Les auteurs proposent un changement de paradigme : au lieu de considérer le guidage comme un échantillonnage à partir d'une distribution inclinée par la récompense (un problème stochastique), ils le formulent comme un problème de contrôle optimal déterministe.

A. Formulation théorique

Le problème est défini comme la minimisation d'une fonction de coût qui arbitre entre l'effort de contrôle et la récompense finale :
$$\min_{u} L(u) = \min_{u} \int_0^1 \frac{\|u_t\|^2}{2\lambda} dt - r(x^u_1)$$
Sous les dynamiques contrôlées : $\dot{x}^u_t = b_t(x^u_t) + u_t$, où $b_t$ est le champ de vitesse pré-entraîné et $u_t$ est le signal de contrôle.

En utilisant le principe du maximum de Pontryagin, les auteurs dérivent la loi de contrôle optimale. Crucialement, ils montrent que dans la limite d'un coût de contrôle faible ($\lambda \to 0$) ou via une approximation « gourmande » (greedy), le signal de guidage optimal dépend explicitement du Flow Map ($X_{t,1}$), qui mappe l'état au temps $t$ vers l'état final au temps $t=1$.

Le signal de guidage résultant est :
$$u_t(x) = \lambda \nabla X_{t,1}(x)^T \nabla r(X_{t,1}(x))$$

B. Composants algorithmiques clés

1. Intégration du Flow Map : Contrairement aux méthodes traditionnelles qui approximent l'état futur en utilisant une seule étape d'Euler (débruiteur), FMRG utilise le Flow Map pré-entraîné pour intégrer exactement les dynamiques de base et prédire le point final $X_{t,1}(x)$.
2. Deux variantes :
   • FMRG-J (Jacobian) : Utilise le Jacobien complet $\nabla X_{t,1}(x)^T$ pour rétropropager le gradient de la récompense. Cela projette le gradient sur l'espace tangent de la variété des données, assurant que la trajectoire reste « sur la variété » et préservant la cohérence sémantique.
   • FMRG-E (Euclidien) : Omet le calcul du Jacobien pour économiser la mémoire, n'utilisant que $\nabla r(X_{t,1}(x))$. C'est plus rapide mais peut dériver hors de la variété, potentiellement causant des artefacts (piratage de la récompense).
3. Décomposition d'opérateurs : L'algorithme alterne entre l'intégration du flot de base (en utilisant le flow map) et l'application d'étapes de gradient pour diriger vers une récompense élevée.
4. Arrêt anticipé : Pour prévenir l'« effondrement de mode » (où un guidage gourmand réduit trop agressivement la variance), les auteurs introduisent un arrêt anticipé. Le guidage est appliqué uniquement jusqu'à un temps $t_{stop}$, après quoi le flow map non contrôlé complète la génération. Cela permet de retrouver la diversité du modèle de base tout en maintenant une récompense élevée.

C. Hiérarchie théorique

Le papier établit une hiérarchie d'approximations. Il prouve que les méthodes existantes comme DPS, FlowDPS et FlowChef sont essentiellement des approximations grossières de FMRG où le Flow Map exact est remplacé par une approximation d'Euler en un seul pas (le débruiteur). FMRG se situe au sommet de cette hiérarchie en tant qu'approximation la plus précise et principielle.

3. Contributions clés

1. Cadre principiel : Reformule le guidage comme un problème de contrôle optimal déterministe, fournissant une fondation théorique pour un guidage à trajectoire unique et en peu d'étapes, natif aux flow maps.
2. Algorithme centré sur le Flow Map : Introduit FMRG, un cadre sans entraînement qui exploite le flow map à la fois pour l'intégration et le guidage. Il unifie et englobe les méthodes existantes à trajectoire unique (DPS, FlowChef, etc.) comme cas particuliers.
3. Insight géométrique : Démontre que le Jacobien du flow map projette naturellement les gradients de récompense sur l'espace tangent de la variété des données, empêchant la génération d'artefacts courants dans l'optimisation hors variété.
4. Efficacité : Atteint des performances de l'état de l'art avec aussi peu que 3 NFEs (évaluations de fonctions de réseau), offrant des accélérations de 10x à 70x par rapport aux méthodes précédentes de l'état de l'art.

4. Résultats expérimentaux

Les auteurs ont évalué FMRG sur FLUX.1-Dev (un modèle à grande échelle basé sur les flots) à travers diverses tâches :

• Problèmes inverses dans l'espace latent (Super-résolution, Défloutage, Inpainting) :
  • FMRG (les deux variantes J et E) a largement surpassé les références (DPS, FlowChef, FlowDPS) en termes de PSNR, SSIM et de métriques perceptuelles (LPIPS, FID).
  • FMRG-E a obtenu des résultats compétitifs avec 3 NFEs, égalant des références qui nécessitaient 20 à 100 NFEs.
• Transfert de style :
  • FMRG-J a capturé les styles cibles plus fidèlement tout en préservant le contenu sémantique par rapport aux méthodes basées sur le débruiteur, qui produisaient souvent des artefacts ou échouaient à intégrer le style.
• Génération guidée par récompense (Préférences humaines et VLM) :
  • Sur le benchmark GenEval (texte-à-image compositionnel), FMRG-J a obtenu un score de 0,80 à 100 NFEs, égalant les performances de FMTT (une méthode beaucoup plus lourde) à 1400 NFEs.
  • Cela représente une réduction de 70x du calcul pour une qualité équivalente.
  • FMRG a guidé avec succès des générations utilisant des récompenses complexes de modèles vision-langage (VLM) pour suivre des invites compositionnelles détaillées que le modèle de base n'arrivait pas à suivre.

5. Signification et impact

• Pont entre théorie et pratique : Le papier résout la tension entre le « tilt de récompense » théorique (stochastique) et le besoin pratique d'un échantillonnage rapide et déterministe. Il montre que le contrôle optimal déterministe est le cadre correct pour les modèles basés sur les flots.
• Révolution de l'efficacité : En exploitant le flow map, FMRG élimine le besoin d'échantillonnage multi-particules coûteux ou de fine-tuning itératif. Cela rend la génération alignée de haute qualité réalisable sur du matériel grand public avec une latence minimale.
• Généralisabilité : Le cadre est agnostique au modèle concernant la fonction de récompense (supportant les pertes $\ell_2$, les préférences de réseaux de neurones et les VLM) et unifie un large éventail de méthodes heuristiques existantes sous une seule bannière mathématique.
• Direction future : Il ouvre la voie à un alignement efficace au moment de l'inférence dans d'autres modèles génératifs déterministes et suggère que le contrôle optimal « gourmand » est une alternative viable au contrôle stochastique complet pour de nombreuses applications pratiques.

En résumé, FMRG fournit une méthode mathématiquement rigoureuse, hautement efficace et sans entraînement pour orienter les modèles génératifs modernes basés sur les flots vers les objectifs de l'utilisateur, obtenant des résultats de l'état de l'art avec une fraction du coût computationnel des approches précédentes.