A Few-Step Generative Model on Cumulative Flow Maps

Résumé technique : Un modèle génératif en quelques étapes sur les cartes de flux cumulatives

Énoncé du problème
Les modèles génératifs, y compris les modèles de diffusion et l'appariement de flux (flow matching), sont généralement formulés comme l'apprentissage de cartes de flux instantanées (ou de dynamiques instantanées) qui prédisent des mises à jour d'état locales conditionnées par l'état courant $x(t)$. Pour générer des données, ces modèles nécessitent une intégration numérique itérative sur de nombreux petits pas de temps afin d'approcher la carte de flux cumulative — le transport à temps fini qui déplace des échantillons d'une distribution source simple vers une distribution de données complexe. Ce processus d'échantillonnage multi-étapes engendre des coûts de calcul significatifs.

Bien que des efforts récents aient exploré la génération en quelques étapes et en une seule étape, les approches existantes rencontrent des limitations :

1. Les méthodes basées sur la distillation nécessitent souvent des pipelines d'entraînement enseignant-élève complexes.
2. Les modèles de cohérence et le Flux Moyen (qui apprend des vitesses moyennes) sont principalement conçus pour des formulations spécifiques (par exemple, l'appariement de flux de type $u$-prédiction) et ne se généralisent pas naturellement à d'autres cadres largement utilisés en infographie, tels que DDIM, EDM ou l'appariement de flux de type $x_1$-prédiction.
3. L'apprentissage direct de cartes cumulatives à longue portée est difficile car les contreparties conditionnelles requises pour les objectifs d'entraînement standards n'existent pas pour les champs cumulatifs, créant ainsi une barrière fondamentale à l'entraînement.

Méthodologie : Cartes de flux cumulatives (CFM)
Les auteurs proposent les Cartes de flux cumulatives (CFM), un cadre unifié qui relie les mises à jour instantanées locales au transport à temps fini via une paramétrisation explicite de champ cumulatif.

• Abstraction unifiée : La méthode formalise la carte de flux instantanée $\psi_{t \to t+h}$ en utilisant une fonction abstraite $F[m_t(x), x, t, t+h]$, où $m_t(x)$ est le champ instantané (par exemple, la vitesse). Cette abstraction unifie des paramétrisations diverses (u-FM, $x_1$-FM, DDIM, EDM).
• Paramétrisation cumulative : Les CFM étendent cela au transport à longue portée en définissant une carte de flux cumulative $\psi_{t \to r}$ comme la limite de la composition de cartes instantanées. Le modèle apprend un champ de paramétrisation cumulative $m_{t \to r}(x)$ tel que $\psi_{t \to r}(x) = F[m_{t \to r}(x), x, t, r]$. Cela permet au modèle de faire avancer l'échantillonnage du temps $t$ vers un temps futur arbitraire $r$ en une seule ou quelques étapes.
• Entraînement via équations de champ : Un défi clé est qu'un champ cumulatif conditionnel $m_{t \to r}(x | X_1)$ n'existe pas de manière auto-cohérente. Pour surmonter cela, les auteurs dérivent une reformulation basée sur les équations de champ (Théorème 3). Ils expriment le champ cumulatif en termes du champ instantané et de ses dérivées :
  $$m_{t \to r}(x) = G(t, r)m_t(x) + H(t, r)E[\dots]$$
  Cette formulation permet de construire une fonction de perte de substitution (Éq. 7) qui utilise le champ instantané conditionnel $m_t(x | X_1)$ (qui est apprenable) pour superviser le champ cumulatif. La perte inclut des termes de dérivée ($\partial_t m$ et $\partial_x m$) qui sont calculés via des produits Jacobien-vecteur (JVP) ou des approximations discrètes.
• Implémentation : L'approche nécessite des modifications architecturales minimales. Elle enrichit l'incorporation temporelle avec un incorporateur supplémentaire pour le temps cible $r$ et utilise une incorporation moyenne $(emb_t + emb_r)/2$. Aucune distillation ni augmentation de capacité n'est requise.

Contributions clés

1. Abstraction carte de flux instantanée–cumulative : L'article formalise les cartes de flux cumulatives comme des transports à temps fini obtenus par composition de cartes instantanées, unifiant la génération multi-étapes et en quelques étapes sous un seul cadre mathématique.
2. Généralisation au-delà de la prédiction $u$ : Les CFM généralisent le concept de Flux Moyen à un large éventail de formulations génératives, incluant l'appariement de flux $u$ et $x_1$, EDM et DDIM. Cela permet la génération en quelques étapes dans des contextes où les méthodes précédentes (comme le Flux Moyen) sont inapplicables, tels que la modélisation de distributions géométriques et la génération d'images dans l'espace des pixels.
3. Entraînement agnostique au modèle : En dérivant un objectif basé sur les équations de champ, la méthode permet l'apprentissage de cartes de flux cumulatives sans modifications architecturales ni distillation, réduisant substantiellement les étapes d'échantillonnage dans les modèles existants.

Résultats expérimentaux
Les auteurs évaluent les CFM sur cinq tâches d'infographie diverses, démontrant des accélérations significatives (10×–200×) tout en préservant ou en améliorant la qualité de génération :

• Génération d'images (CelebA-HQ) : En utilisant CFM-DDIM, le modèle réalise une génération en 1 étape et en 4 étapes avec des scores FID comparables aux bases de référence à 128 étapes, surpassant la distillation de cohérence et les méthodes d'entraînement.
• Modélisation de distributions géométriques : Sur la tâche GeoDist, CFM-EDM réalise des accélérations de 6× à 10× sans dégradation de la distance de Chamfer par rapport aux bases de référence à 60 étapes. Notamment, les méthodes de prédiction $u$ échouent à supporter la génération en quelques étapes dans ce domaine, tandis que CFM-EDM réussit.
• Prédiction conjointe de positions (PDT) : Appliqué à l'ensemble de données RigNet, CFM-DDIM réduit l'inférence de 1000 étapes à 5 étapes (accélération de 200×) tout en maintenant une précision de prédiction des articulations (CD-J2J, IoU, Précision, Rappel) comparable au PDT original à 1000 étapes.
• Génération de croquis conditionnée par l'image : Sur l'ensemble de données ControlSketch, CFM génère des croquis vectorisés en 1 ou 4 étapes (accélération de 50×) avec une fidélité (MS-SSIM, DreamSim) correspondant à la base de référence SwiftSketch à 50 étapes.
• Génération de SDF 3D : Pour la reconstruction de SDF conditionnels clairsemés (64 points), CFM-$x_1$-FM réalise des accélérations de 6 à 16× par rapport à la Diffusion Fonctionnelle avec une qualité de reconstruction comparable.

Signification et affirmations
L'article affirme que les CFM fournissent un cadre mathématique général particulièrement adapté aux applications d'infographie impliquant des représentations de données diverses (images, nuages de points, champs implicites) et des formulations génératives. Sa signification principale réside dans :

• Découplage de l'efficacité et de l'architecture : Il réalise une accélération drastique de l'inférence (jusqu'à 200×) uniquement en modifiant l'objectif d'entraînement et les incorporations temporelles, sans nécessiter de distillation ni de révisions architecturales.
• Cadre unificateur : Il comble le fossé entre les dynamiques instantanées et le transport à longue portée, offrant une solution fondée sur des principes pour la génération en quelques étapes à travers des modèles où les approches précédentes de « vitesse moyenne » ont échoué.
• Praticité : La méthode est agnostique au modèle et peut être appliquée à des modèles multi-étapes pré-entraînés pour les accélérer, ou utilisée pour entraîner de nouveaux modèles en quelques étapes à partir de zéro.

Les auteurs maintiennent que, bien que les CFM soient applicables à un large éventail de modèles, leur évaluation actuelle se limite à cinq applications et quatre formulations représentatives, laissant l'extension à des ensembles de données à plus grande échelle (par exemple, ImageNet) comme travail futur.