Discriminative Flow Matching Via Local Generative Predictors

1. Problématique

Le paradigme dominant en vision par ordinateur discriminative repose sur des projections statiques. Des architectures comme les ResNets et les Vision Transformers (ViT) fonctionnent comme des mappings fixes qui projettent les caractéristiques d'entrée directement sur des logits ou des coordonnées en une seule étape computationnelle. Bien que efficaces, ces approches manquent de la réfinition itérative et de la robustesse inhérentes à la vision biologique et aux modèles génératifs modernes (comme les modèles de diffusion).

L'adaptation des processus génératifs continus (comme l'appariement de flux ou Flow Matching) aux tâches discriminatives se heurte à deux obstacles majeurs :

1. Goulots d'étranglement computationnels : L'entraînement d'un réseau profond pour approximer un champ vectoriel global nécessite souvent la rétropropagation à travers des solveurs d'équations différentielles ordinaires (ODE), ce qui entraîne une consommation mémoire linéaire par rapport à la profondeur et une instabilité d'optimisation.
2. Limites des méthodes d'apprentissage local existantes : Des alternatives comme Forward-Forward ou NoProp découpent l'entraînement des couches mais peinent à s'étendre aux tâches de prédiction dense (comme la détection d'objets) car elles reposent sur des objectifs scalaires qui ne préservent pas suffisamment les informations spatiales haute dimensionnelle.

2. Méthodologie : Discriminative Flow Matching (DFM)

Les auteurs proposent un cadre novateur qui reformule la classification et la détection d'objets comme un processus de transport conditionnel. Au lieu d'une projection statique, l'inférence est conçue comme un processus de consensus guidé par un champ vectoriel.

A. Formulation du Problème

Le modèle apprend à transporter des échantillons depuis une distribution de bruit simple (Gaussienne, $z_0$) vers une variété cible alignée sur la tâche ($z_1$) via un champ vectoriel appris.

• Pour la classification : La cible $z_1$ est une embedding de classe apprenable.
• Pour la détection d'objets : La cible $z_1$ est un ensemble de requêtes contenant des coordonnées de boîtes englobantes et des classes.
• Chemin de probabilité : Une interpolation linéaire est définie entre le bruit et la cible : $z_t = (1-t)z_0 + tz_1$. Le réseau doit apprendre le champ vectoriel $v_t = z_1 - z_0$.

B. Architecture : Ensemble de Flux Génératifs (Generative Flow Ensemble)

L'innovation centrale réside dans la décomposition du modèle :

1. Backbone Partagé : Un réseau extracteur de caractéristiques (CNN ou ViT) est partagé par tous les blocs.
2. Prédicteurs Locaux Indépendants : Au lieu d'un seul réseau profond, le modèle utilise $T$ prédicteurs de champ vectoriel indépendants ($v_{\theta_1}, ..., v_{\theta_T}$) attachés au backbone.
3. Entraînement Local (Block-wise Local Training) :
   • Chaque bloc est entraîné indépendamment avec sa propre fonction de perte locale (Flow Matching Loss).
   • Pas de rétropropagation globale : Les gradients sont calculés uniquement pour le bloc courant et le backbone, sans traverser les autres blocs. Cela évite la nécessité de stocker toutes les activations intermédiaires d'un réseau profond.
   • Objectif de tâche : Une perte auxiliaire (classification ou détection) est ajoutée pour garantir l'utilité du flux pour la tâche finale.

C. Stratégies d'Inférence

Le framework propose plusieurs modes d'inférence :

• Single-Step : Une seule mise à jour d'Euler ($\hat{z}_1 = z_0 + v_\theta(z_0, 0)$).
• Ensemble : Agrégation des prédictions de plusieurs prédicteurs indépendants pour un consensus robuste.
• Multi-Step (ODE) : Intégration numérique itérative du champ vectoriel (souvent instable si les prédicteurs ne sont pas spécialisés par intervalle de temps).

3. Contributions Clés

1. Extension de l'Appariement de Flux : Application réussie du Flow Matching conditionnel aux tâches discriminatives complexes (classification et détection d'objets), formulant ces tâches comme un transport de bruit vers des variétés cibles.
2. Ensemble de Flux Génératifs : Introduction d'une architecture où plusieurs têtes de prédiction apprennent des champs vectoriels locaux sur un backbone partagé. Cela offre la robustesse d'un ensemble sans le coût mémoire de l'entraînement de multiples réseaux complets.
3. Efficacité Mémoire et Scalabilité : Démonstration que l'entraînement par blocs locaux permet de réduire la consommation mémoire de pic (indépendante du nombre de blocs $T$) tout en maintenant des performances compétitives sur des architectures CNN et ViT.
4. Analyse des Dynamiques d'Inférence : Mise en évidence du fait que de fortes représentations de caractéristiques ne garantissent pas toujours une cohérence du champ vectoriel, et que l'agrégation par ensemble est cruciale pour la stabilité, surtout dans les architectures convolutionnelles.

4. Résultats Expérimentaux

Les expériences ont été menées sur PASCAL VOC (détection), CIFAR-10/100, ImageNet-R et Tiny-ImageNet (classification).

• Détection d'objets (PASCAL VOC) :
  • Avec un backbone ResNet-50, la méthode DFM atteint un mAP d'environ 30%, comparable à la rétropropagation standard (31%).
  • Gain mémoire : La consommation de VRAM est réduite, et l'avantage augmente avec la profondeur du modèle (nombre de couches de décodeur). Par exemple, avec un batch size de 1, DFM utilise ~510 Mo contre ~579 Mo pour la rétropropagation standard.
• Classification (Fine-tuning ViT) :
  • Sur CIFAR-10, DFM surpasse le fine-tuning standard (98,19% vs 96,74%) avec une consommation mémoire inférieure.
  • Sur CIFAR-100 et ImageNet-R, DFM atteint des performances très proches du standard (légèrement inférieures sur les tâches complexes) mais avec une économie de mémoire significative.
• Analyse des Dynamiques :
  • L'inférence par Ensemble s'avère être la stratégie la plus stable.
  • L'intégration ODE multi-étapes peut devenir instable (dérive de trajectoire), en particulier sur les CNN, soulignant l'importance de l'agrégation de consensus.
  • Les CNN bénéficient d'un entraînement séquentiel pour la qualité des caractéristiques, tandis que les ViT montrent une dynamique différente.

5. Signification et Conclusion

Ce travail établit un pont pratique entre la dynamique générative continue et l'apprentissage discriminatif efficace.

• Réduction de la barrière mémoire : En découplant l'optimisation des blocs génératifs, la méthode permet d'entraîner des trajectoires génératives profondes sans épuiser les ressources matérielles, contournant les limitations de la rétropropagation standard.
• Nouveau paradigme d'inférence : Il propose de voir l'inférence discriminative non pas comme une projection instantanée, mais comme un processus de transport itératif ou de consensus, offrant une robustesse accrue face aux ambiguïtés.
• Flexibilité Architecturale : La méthode s'adapte aussi bien aux CNN qu'aux Transformers, suggérant que le Flow Matching discriminatif est une approche générale pour l'apprentissage de représentations robustes et économes en ressources.

En résumé, le Discriminative Flow Matching offre une alternative viable et efficace aux architectures discriminatives traditionnelles, particulièrement dans des scénarios contraints par la mémoire ou nécessitant une grande robustesse d'inférence.