Modèles de trajectoire de cohérence : apprentissage de la trajectoire d’ODE de flux de probabilité pour la diffusion

Les modèles de cohérence (Consistency Models, CM) (Song et al., 2023) accélèrent l’échantillonnage des modèles de diffusion fondés sur les scores au détriment de la qualité des échantillons, tout en manquant d’une approche naturelle pour équilibrer qualité et vitesse. Pour surmonter cette limitation, nous proposons le Consistency Trajectory Model (CTM), une généralisation qui englobe à la fois les CM et les modèles fondés sur les scores comme cas particuliers. Le CTM entraîne un seul réseau neuronal capable, en une seule passe avant, de produire des scores (c’est-à-dire les gradients de la densité logarithmique) et permet un parcours sans restriction entre tout instant initial et tout instant final le long de l’équation différentielle ordinaire (ODE) de flux de probabilité dans un processus de diffusion. Grâce au CTM, il devient possible de combiner efficacement l’entraînement adversaire et la perte de correspondance des scores de débruitage, ce qui améliore significativement les performances. Le CTM atteint de nouveaux états de l’art en termes de FID pour l’échantillonnage en une seule étape sur CIFAR-10 (FID 1,73) et sur ImageNet à une résolution de 64×64 (FID 1,92). Le CTM ouvre également la voie à une nouvelle famille de méthodes d’échantillonnage, à la fois déterministes et stochastiques, incluant des « sauts longs » le long des trajectoires de solution de l’ODE. Il améliore de manière cohérente la qualité des échantillons à mesure que le budget computationnel augmente, évitant ainsi la dégradation observée dans les CM. En outre, contrairement aux CM, l’accès au fonctionnement de score fourni par le CTM facilite l’intégration de méthodes établies de génération contrôlée ou conditionnelle issues de la communauté des modèles de diffusion. Ce même accès permet également le calcul de la vraisemblance. Le code est disponible à l’adresse suivante : https://github.com/sony/ctm.