Modèles de cohérence tronqués

Les modèles de cohérence ont récemment été introduits afin d’accélérer l’échantillonnage à partir des modèles de diffusion en prédisant directement la solution (c’est-à-dire les données) du système d’équations différentielles stochastiques de flux de probabilité (PF ODE) à partir d’un bruit initial. Toutefois, l’entraînement de ces modèles nécessite d’apprendre à mapper tous les points intermédiaires le long des trajectoires du PF ODE vers leurs points finaux correspondants. Cette tâche s’avère bien plus exigeante que l’objectif ultime de génération en une étape, qui ne concerne que la transformation du bruit en données selon le PF ODE. Nous observons empiriquement que ce paradigme d’entraînement limite les performances de génération en une étape des modèles de cohérence. Pour résoudre ce problème, nous généralisons l’entraînement de cohérence à un intervalle temporel tronqué, ce qui permet au modèle d’ignorer les tâches de débruitage aux instants initiaux et de concentrer ses capacités sur la génération. Nous proposons une nouvelle paramétrisation de la fonction de cohérence ainsi qu’un procédé d’entraînement en deux étapes, qui empêche l’entraînement sur un intervalle tronqué de se réduire à une solution triviale. Des expériences menées sur les jeux de données CIFAR-10 et ImageNet $64\times64$ montrent que notre méthode atteint des scores FID en une et deux étapes supérieurs à ceux des modèles de cohérence les plus avancés actuellement disponibles, tels que iCT-deep, tout en utilisant des réseaux d’une taille inférieure de plus de 2 fois. Page du projet : https://truncated-cm.github.io/