🤖 machine learning

COT-FM: Cluster-wise Optimal Transport Flow Matching

COT-FM est un cadre général qui améliore la génération par Flow Matching en réorganisant les trajectoires de transport optimal via un regroupement par clusters, permettant ainsi d'obtenir des champs vectoriels plus linéaires, un échantillonnage plus rapide et une qualité de génération supérieure sans modifier l'architecture du modèle.

Auteurs originaux : Chiensheng Chiang, Kuan-Hsun Tu, Jia-Wei Liao, Cheng-Fu Chou, Tsung-Wei Ke

Publié 2026-03-17

📖 4 min de lecture☕ Lecture pause café

CC BY 4.0

Auteurs originaux : Chiensheng Chiang, Kuan-Hsun Tu, Jia-Wei Liao, Cheng-Fu Chou, Tsung-Wei Ke

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 essayez de guider un groupe de personnes (les données) depuis un point de départ simple, comme une place publique vide (une distribution gaussienne), jusqu'à un lieu de destination complexe et coloré, comme un festival avec plusieurs scènes distinctes (une image ou un objet 3D).

C'est ce que font les modèles de génération d'images actuels, appelés Flow Matching (ou "Appariement de Flux"). Ils apprennent un chemin pour que les gens marchent de la place vers le festival.

Le Problème : La Foule Se Croise et Fait des Boucles

Dans les méthodes actuelles, le problème est que le modèle essaie de guider tout le monde en même temps, sans faire de distinctions.

L'analogie : Imaginez un chef d'orchestre qui crie des instructions à 10 000 personnes en même temps. Certaines doivent aller à la scène de rock, d'autres à la scène de jazz. Comme le chef ne fait pas la différence, il donne des ordres moyens : "Allez un peu vers la gauche, un peu vers la droite".
Le résultat : Au lieu de lignes droites et efficaces, les gens font des zigzags, se croisent, se bousculent et finissent par faire des boucles inutiles. Pour arriver à destination, ils doivent marcher longtemps (beaucoup d'étapes de calcul), et à la fin, ils sont souvent un peu fatigués et mal placés (l'image générée est floue ou déformée).

La Solution : COT-FM (Le Guide de Groupes Intelligents)

Les auteurs de cette paper proposent une nouvelle méthode appelée COT-FM. Au lieu de traiter tout le monde comme une seule masse, ils utilisent une stratégie "diviser pour régner".

Voici comment ça marche, étape par étape :

1. Le Tri par Groupes (Le Clustering)

Au lieu de crier à tout le monde, le modèle regarde la destination et dit : "Tiens, il y a un groupe qui veut aller à la scène de rock, et un autre qui veut aller au stand de nourriture".

En langage simple : Ils séparent les données en petits groupes (clusters) basés sur ce qu'elles ont en commun.

2. La Carte de Retour (L'Optimal Transport)

Pour chaque groupe, le modèle fait une expérience géniale : il imagine comment les gens sont arrivés là. Il "rembobine" le voyage pour voir d'où ils venaient exactement.

L'analogie : C'est comme si le modèle regardait une vidéo à l'envers pour dire : "Ah, les gens de la scène de rock venaient tous de la porte Nord, et ceux du stand de nourriture de la porte Sud".
Au lieu d'avoir un seul point de départ flou pour tout le monde, chaque groupe a maintenant son propre point de départ précis.

3. Des Chemins Droits et Rapides

Maintenant que chaque groupe a son propre point de départ et sa propre destination, le modèle peut tracer des lignes droites pour chacun.

Le résultat : Plus de croisements, plus de boucles. C'est comme passer d'une route de montagne sinueuse à une autoroute droite.
L'avantage : Les gens arrivent beaucoup plus vite (moins d'étapes de calcul) et arrivent exactement à la bonne place (images plus nettes et réalistes).

Pourquoi c'est révolutionnaire ?

C'est comme un "Plug-and-Play" : Vous n'avez pas besoin de reconstruire toute la voiture (le modèle) pour avoir des pneus meilleurs. Vous changez juste la carte de navigation (la trajectoire), et la voiture va plus vite et plus loin.
Moins d'effort, plus de qualité : Avec cette méthode, on peut générer une image de haute qualité en un seul pas (comme un coup de baguette magique) là où les anciennes méthodes avaient besoin de 50 pas laborieux.
Pour les robots aussi : Le papier montre que ça marche même pour les robots qui doivent attraper des objets. Au lieu de faire des mouvements hésitants et lents, le robot sait exactement où mettre sa main dès le premier mouvement.

En Résumé

Imaginez que vous envoyez des colis.

L'ancienne méthode : Vous mettez tous les colis dans un seul camion et vous essayez de les livrer en faisant des détours constants pour éviter les embouteillages. C'est lent et les colis arrivent parfois abîmés.
COT-FM : Vous séparez les colis par quartier, vous envoyez un petit camion pour chaque quartier avec un itinéraire direct. C'est rapide, efficace, et les colis arrivent en parfait état.

C'est exactement ce que fait COT-FM : il organise le chaos en petits groupes bien gérés pour que la génération d'images (ou d'actions de robots) soit plus rapide, plus simple et plus belle.

1. Problématique

Le Flow Matching (FM) est un cadre de modélisation générative qui apprend un champ de vecteurs déterministe pour transporter des échantillons d'une distribution source simple (ex: Gaussienne) vers une distribution cible complexe. Bien que flexible et évolutif, le FM souffre de deux limitations majeures lors de l'inférence :

Trajectoires courbées : Les méthodes d'appariement (couplings) courantes, comme l'appariement aléatoire ou l'appariement optimal par lots (batch-wise OT), génèrent des trajectoires de transport non linéaires. Ces courbures augmentent l'erreur de discrétisation temporelle, obligeant à utiliser un grand nombre d'étapes d'intégration (NFE) pour obtenir des échantillons de haute qualité.
Compromis qualité/vitesse : Les méthodes existantes tentent de "redresser" ces trajectoires soit par distillation (souvent coûteuse), soit en apprenant un champ de vecteurs moyen (comme MeanFlow). Cependant, ces approches ne modifient pas fondamentalement la courbure sous-jacente du champ de vecteurs, limitant ainsi l'amélioration de la qualité de génération, surtout avec un nombre d'étapes très faible (1 ou 2 étapes).

L'objectif est de concevoir un cadre qui redresse les trajectoires de transport tout en préservant la structure des données, permettant une génération rapide et de haute qualité sans modifier l'architecture du modèle.

2. Méthodologie : COT-FM

Les auteurs proposent COT-FM (Cluster-wise Optimal Transport Flow Matching), une approche "plug-and-play" qui repose sur une stratégie de diviser pour régner. Au lieu de mapper l'ensemble de la distribution cible depuis une seule distribution source globale, la méthode partitionne les données et attribue une distribution source dédiée à chaque sous-ensemble.

Le processus se déroule en deux étapes principales, itératives :

A. Identification des distributions sources par cluster

Clustering : Les échantillons de données cibles sont divisés en $K$ clusters (via K-means non supervisé, étiquettes de classes, ou descriptions textuelles).
Estimation de la source : Pour chaque cluster $C_k$ , la méthode utilise un modèle FM pré-entraîné (entraîné initialement avec un couplage aléatoire) pour inverser le processus d'ODE. En intégrant le champ de vecteurs appris à l'envers, on récupère les échantillons sources $\hat{x}_0$ qui ont généré chaque échantillon du cluster.
Approximation Gaussienne : La distribution source pour ce cluster est approximée par une loi Gaussienne $\mathcal{N}(\mu_{0,k}, \Sigma_{0,k})$ calculée à partir des statistiques des échantillons inversés. Cette étape garantit que les chemins entre les clusters sont peu intersectés.

B. Optimisation par Transport Optimal (OT) intra-cluster

Couplage Optimal Local : Au lieu de calculer un OT global (coûteux en $O(n^3)$ ), COT-FM calcule des cartes de transport optimal approximatives (par lots) entre les échantillons sources générés et les échantillons cibles au sein de chaque cluster.
Réduction de la complexité : En réduisant la taille des problèmes d'OT à la taille des clusters, l'approximation par lots devient beaucoup plus précise, produisant des couplages cohérents et des trajectoires plus droites.
Entraînement Alterné : Le modèle FM est mis à jour pour régresser ce nouveau champ de vecteurs basé sur les couplages OT locaux. Ce processus alterne entre le raffinement des distributions sources (via l'inversion ODE) et l'optimisation du modèle FM.

3. Contributions Clés

Cadre Général et Plug-and-Play : COT-FM ne modifie pas l'architecture du modèle FM ni ses mécanismes d'entrée/sortie. Il agit uniquement sur la définition du chemin de probabilité cible, ce qui le rend compatible avec la plupart des modèles FM existants (Rectified Flow, OT-CFM, MeanFlow, etc.).
Redressement des Trajectoires : En partitionnant l'espace des données et en appariant localement les sources et les cibles, la méthode réduit drastiquement la courbure des champs de vecteurs, minimisant ainsi l'erreur de discrétisation.
Génération Haute Qualité à Faible Budget : La méthode permet d'obtenir des résultats de haute qualité avec un nombre d'étapes d'inférence très réduit (1 ou 2 étapes), là où les méthodes de base échouent ou nécessitent des centaines d'étapes.
Extensibilité aux Conditions : Le cadre s'étend naturellement à la génération conditionnelle (images, robotique) où les conditions (étiquettes, textes, observations) servent de clusters naturels.

4. Résultats Expérimentaux

Les auteurs ont évalué COT-FM sur plusieurs benchmarks, montrant des améliorations systématiques par rapport aux méthodes de l'état de l'art :

Données 2D Synthétiques : Sur des distributions comme "5-Gaussians" et "Two Moons", COT-FM atteint les distances de Wasserstein les plus faibles et la courbure la plus basse, surpassant Rectified Flow, OT-CFM et MeanFlow.
Génération d'Images (CIFAR-10) :
- En 10 étapes, le FID passe de 12.6 (Rectified Flow) à 8.23 avec COT-FM.
- En 1 étape (1-step), le FID est réduit de 378.0 à 205.0 pour Rectified Flow, et de 2.92 à 2.60 pour MeanFlow.
Génération Conditionnelle (ImageNet 256x256) : Avec des architectures SiT, COT-FM améliore le FID à tous les niveaux d'étapes (NFE), avec des gains particulièrement marqués pour les faibles nombres d'étapes.
Robotique (LIBERO) : Sur des tâches de manipulation robotique conditionnées par le texte, COT-FM atteint un taux de réussite de 96.1% (Spatial) et 94.5% (Long) avec une seule étape d'évaluation (1 NFE). Cela surpasse la politique FLOWER qui nécessite 4 étapes pour atteindre des performances similaires (97.1% et 93.5%).

5. Signification et Impact

COT-FM démontre que la courbure du flux est le facteur limitant principal pour la génération rapide et de haute qualité dans les modèles Flow Matching. En résolvant le problème de transport optimal à l'échelle locale (par cluster) plutôt que globalement, la méthode contourne les limitations computationnelles de l'OT exact tout en évitant les pièges des couplages aléatoires ou par lots.

Cela ouvre la voie à des modèles génératifs ultra-rapides (en temps réel pour la robotique ou l'interaction humaine) qui ne sacrifient pas la fidélité des échantillons. La capacité de COT-FM à fonctionner comme une couche d'amélioration "plug-and-play" sur des modèles pré-entraînés en fait une contribution pratique majeure pour le déploiement de modèles génératifs dans des applications exigeantes en latence.