Amélioration de l'entraînement des flux rectifiés

Les modèles de diffusion ont montré un grand potentiel pour la génération d’images et de vidéos, mais l’échantillonnage à partir des modèles les plus avancés nécessite une intégration numérique coûteuse d’une équation différentielle ordinaire (ODE) générative. Une approche pour résoudre ce problème repose sur les flows rectifiés, qui apprennent itérativement des trajectoires continues d’ODE moins sensibles aux erreurs de troncature. Toutefois, les flows rectifiés nécessitent encore un nombre relativement élevé d’évaluations de fonction (NFE). Dans ce travail, nous proposons des techniques améliorées pour l’entraînement des flows rectifiés, permettant à ces derniers de rivaliser avec les méthodes de distillation de connaissances, même dans le régime à faible NFE. Notre principal constat est que, dans des conditions réalistes, une seule itération de l’algorithme Reflow pour l’entraînement des flows rectifiés suffit à apprendre des trajectoires presque rectilignes ; par conséquent, la pratique courante consistant à utiliser plusieurs itérations de Reflow est inutile. Nous proposons donc des techniques visant à améliorer l’entraînement en une seule boucle, notamment une distribution en forme de « U » des pas de temps et un pré-métrique LPIPS-Huber. Grâce à ces améliorations, nous améliorons le FID du modèle 2-rectified flow précédent de jusqu’à 75 % dans le cas d’un seul échantillonnage (1 NFE) sur CIFAR-10. Sur ImageNet 64×64, notre flow rectifié amélioré dépasse les méthodes de distillation de pointe telles que la distillation de cohérence et la distillation progressive, tant dans les scénarios à un pas qu’à deux pas, et atteint des performances comparables à celles de l’entraînement amélioré par cohérence (iCT) en termes de FID. Le code est disponible à l’adresse suivante : https://github.com/sangyun884/rfpp.