Résumé Technique : Stable Velocity : Une perspective sur la variance de Flow Matching

1. Énoncé du Problème

Bien que le Flow Matching (FM) et le Conditional Flow Matching (CFM) offrent des cadres élégants pour la modélisation générative en apprenant le flux de probabilité sans simuler explicitement les SDE ou PF-ODE progressives, ils souffrent d'une limitation fondamentale : la variance élevée de leurs cibles d'entraînement. Dans le CFM, l'objectif d'entraînement repose sur une estimation de Monte Carlo à échantillon unique de la véritable vitesse marginale, $v_t(x_t | x_0)$, pour approximer la vitesse marginale $v_t(x_t)$.

Les auteurs identifient que cet estimateur à échantillon unique présente une variance significative, particulièrement aux instants où la distribution marginale reste proche de la distribution a priori (régime à haute variance). Cette haute variance conduit à :

1. Une optimisation instable : Des gradients bruités qui ralentissent la convergence.
2. Une inefficacité de l'inférence : L'incapacité à exploiter la nature déterministe du flux dans certains régimes, menant à des étapes d'intégration inutilement petites.
3. Une supervision auxiliaire inefficace : Les méthodes existantes d'alignement de représentation (comme REPA) appliquent souvent la supervision de manière uniforme, ne parvenant pas à reconnaître que l'alignement sémantique est mal posé dans les régimes à haute variance (bruyants) mais hautement efficace dans les régimes à faible variance.

2. Méthodologie

Le papier propose le cadre Stable Velocity, qui traite ces problèmes en caractérisant explicitement la variance des cibles de vitesse conditionnelle et en adaptant les stratégies d'entraînement et d'échantillonnage en conséquence.

2.1 Analyse de la Variance

Les auteurs analysent la variance de la cible CFM, $V_{CFM}(t)$, à travers différentes distributions de données (GMM, CIFAR-10, ImageNet). Ils révèlent une structure à deux régimes :

• Régime à Haute Variance ($t \in [\xi, 1]$) : Près de l'a priori (bruit), la distribution postérieure $p(x_0|x_t)$ est multimodale, ce qui provoque des fluctuations significatives de la vitesse conditionnelle.
• Régime à Faible Variance ($t \in [0, \xi]$) : Près de la distribution de données, la distribution postérieure se concentre sur un échantillon de référence unique, faisant que la vitesse conditionnelle $v_t(x_t | x_0)$ coïncide presque avec la véritable vitesse marginale $v_t(x_t)$.

2.2 Entraînement : Stable Velocity Matching (StableVM)

Pour réduire la variance sans introduire de biais, les auteurs proposent le Stable Velocity Matching (StableVM).

• Mécanisme : Au lieu d'utiliser un échantillon de référence unique $x_0$, StableVM construit un chemin conditionnel composite à l'aide d'un lot de $n$ échantillons de référence $\{x_0^i\}_{i=1}^n$.
• Cible : La cible d'entraînement est une moyenne pondérée par importance auto-normalisée des vitesses conditionnelles sur le lot de référence :
  $$\hat{v}_{StableVM}(x_t) = \frac{\sum_{k=1}^n p_t(x_t | x_0^k) v_t(x_t | x_0^k)}{\sum_{j=1}^n p_t(x_t | x_0^j)}$$
• Propriétés : Théoriquement prouvé comme étant non biaisé (partageant le même minimiseur global que le CFM) tout en réduisant strictement la variance de la cible d'entraînement par un facteur de $O(1/n)$.
• Extension Conditionnelle par Classe : Pour gérer la parcité de la génération conditionnelle par classe, une banque de mémoire est utilisée pour maintenir un ensemble diversifié d'échantillons de référence par classe, garantissant que le mélange reste efficace même avec une faible fréquence de classe par lot.

2.3 Entraînement : Variance-Aware Representation Alignment (VA-REPA)

S'appuyant sur l'aperçu du régime de variance, les auteurs proposent VA-REPA pour optimiser la supervision sémantique auxiliaire.

• Aperçu : L'alignement de représentation (par exemple, aligner les caractéristiques latentes avec des encodeurs pré-entraînés) n'est informatif que dans le régime à faible variance où l'information sémantique est préservée. Dans le régime à haute variance, le signal est bruité et mal posé.
• Stratégie : Une fonction de pondération $w(t)$ est appliquée à la perte d'alignement, restreignant la supervision au régime à faible variance ($t < \xi$). Cela empêche le modèle d'apprendre à partir de signaux bruyants et non informatifs.

2.4 Inférence : Stable Velocity Sampling (StableVS)

Les auteurs exploitent le régime à faible variance pour accélérer l'échantillonnage sans ajustement fin (fine-tuning).

• Aperçu : Dans le régime à faible variance, la dynamique du flux est effectivement déterministe, et la trajectoire approxime une ligne droite avec une vitesse constante.
• Mécanisme : StableVS dérive des simplifications en forme close pour les SDE et PF-ODE inverses dans ce régime. Pour les interpolants linéaires, la règle de mise à jour se simplifie en :
  $$x_\tau = x_t + (\tau - t)v_t(x_t)$$
• Implémentation : L'échantillonneur utilise des solveurs standards (ex: Euler, DPM-Solver++) dans le régime à haute variance, mais bascule vers la mise à jour en forme close de StableVS dans le régime à faible variance. Cela permet de grandes étapes d'intégration (moins d'étapes totales) sans dégrader la qualité des échantillons.

3. Contributions Clés

1. Perspective de la Variance Théorique : Le papier fournit une analyse fondée qui révèle une structure à deux régimes dépendante de la dimensionnalité (haute vs faible variance) dans les interpolants stochastiques, régissant à la fois la dynamique d'entraînement et le comportement d'inférence.
2. StableVM : Un objectif d'entraînement non biaisé qui réduit la variance en agrégeant plusieurs échantillons de référence pour stabiliser l'optimisation, surpassant le CFM standard et les techniques existantes de réduction de variance (comme STF) en termes de vitesse de convergence et de qualité d'échantillonnage.
3. VA-REPA : Une adaptation de l'alignement de représentation sensible au régime qui applique sélectivement la supervision auxiliaire dans le régime à faible variance, améliorant l'efficacité et la stabilité de l'entraînement.
4. StableVS : Une stratégie d'accélération sans ajustement fin qui exploite la nature déterministe du régime à faible variance pour obtenir des accélérations de plus de 2x sur l'échantillonnage de modèles de pointe.

4. Résultats Expérimentaux

Le cadre a été validé sur ImageNet 256×256 et divers modèles de grande échelle de texte-vers-image (T2I) et texte-vers-vidéo (T2V) (SD3.5, Flux, Qwen-Image, Wan2.2).

• Efficacité de l'Entraînement :
  • StableVM + VA-REPA a systématiquement surpassé le REPA vanilla et les autres bases (REG, iREPA) à travers les échelles de modèles (SiT-B à SiT-XL) et les durées d'entraînement.
  • Sur ImageNet, la méthode a obtenu des scores FID et IS compétitifs avec nettement moins d'époques d'entraînement par rapport aux bases nécessitant un ajustement complet (ex: REPA-E).
  • Les études d'ablation ont confirmé que la restriction de l'alignement au régime de faible variance produit un meilleur FID que l'application uniforme.
• Vitesse d'Échantillonnage :
  • StableVS a obtenu une accélération de l'inférence >2x (ex: égalant les bases de 30 étapes avec environ 17-20 étapes totales) sur SD3.5, Flux et Wan2.2.
  • Crucialement, cette accélération n'a pas dégradé la qualité des échantillons. Les métriques de référence (PSNR, SSIM, LPIPS) ainsi que les métriques génératives (GenEval, T2V-CompBench) sont restées comparables ou meilleures que celles des solveurs standards avec un nombre réduit d'étapes.
  • Les résultats qualitatifs ont montré que les sorties de StableVS étaient visuellement indiscernables des bases de 30 étapes, alors qu'une réduction naïve des étapes provoquait des artefacts notables.

5. Signification et Revendications

Le papier affirme que la modélisation explicite de la structure de la variance le long de la trajectoire générative fournit une base fondamentale pour concevoir des modèles génératifs plus efficaces.

• Cadre Unifié : Il unifie la réduction de la variance dans l'entraînement et l'accélération sensible au régime dans l'échantillonnage sous une perspective théorique unique.
• Impact Pratique : Les méthodes sont modulaires et compatibles avec les architectures existantes (ex: SiT, DiT), offrant des améliorations immédiates de la stabilité de l'entraînement et de la vitesse d'inférence sans nécesser de changements architecturaux ou d'ajustements fins coûteux.
• Aperçu Théorique : Ce travail remet en question la vision des trajectoires de diffusion/flux comme des processus uniformes, démontant que des régimes distincts (haute vs faible variance) nécessitent des stratégies d'optimisation et d'échantillonnage distinctes.

Les auteurs soulignent que leur approche complète les techniques existantes (comme la réduction de dimensionnalité ou l'amélioration des solveurs) en traitant la source spécifique de variance induite par la stochasticité de l'échantillon de référence dans les objectifs de Flow Matching.