Les modèles de diffusion sont des générateurs intrinsèquement à un seul pas

Les modèles de diffusion (DMs) ont connu un succès remarquable dans la génération d’images et d’autres domaines. En effectuant un échantillonnage fin le long de la trajectoire définie par un solveur SDE/ODE basé sur un modèle de score bien entraîné, les DMs parviennent à produire des résultats de haute qualité exceptionnels. Toutefois, cet échantillonnage précis nécessite généralement un grand nombre d’étapes et s’avère très exigeant en ressources computationnelles. Pour pallier ce problème, des méthodes de distillation basées sur l’instance ont été proposées afin de distiller un générateur à une seule étape à partir d’un DM, en faisant en sorte qu’un modèle étudiant plus simple imite un modèle enseignant plus complexe. Pourtant, notre recherche met en évidence une limitation intrinsèque à ces approches : le modèle enseignant, doté d’un plus grand nombre d’étapes et de paramètres, occupe des minima locaux différents de ceux du modèle étudiant, ce qui entraîne des performances sous-optimales lorsque ce dernier tente de reproduire le comportement du modèle enseignant. Pour contourner ce problème, nous introduisons une nouvelle méthode de distillation distributionnelle, reposant sur une perte distributionnelle exclusive. Cette approche dépasse les résultats de l’état de l’art (SOTA) tout en nécessitant un nombre significativement réduit d’images d’entraînement. Par ailleurs, nous démontrons que les couches des DMs sont activées de manière différentielle aux différentes étapes temporelles, ce qui révèle une capacité intrinsèque à générer des images en une seule étape. En gelant la majorité des couches convolutionnelles d’un DM pendant la distillation distributionnelle, nous exploitons cette capacité naturelle, entraînant une amélioration supplémentaire des performances. Notre méthode atteint des résultats SOTA sur CIFAR-10 (FID 1,54), AFHQv2 64x64 (FID 1,23), FFHQ 64x64 (FID 0,85) et ImageNet 64x64 (FID 1,16) avec une grande efficacité. La plupart de ces résultats sont obtenus avec seulement 5 millions d’images d’entraînement en moins de 6 heures sur 8 GPU A100.