DisCo-Diff : Amélioration des modèles de diffusion continus grâce à des latents discrets

Les modèles de diffusion (DMs) ont révolutionné l'apprentissage génératif. Ils exploitent un processus de diffusion pour encoder les données dans une distribution gaussienne simple. Toutefois, encoder une distribution de données complexe, potentiellement multimodale, dans une unique distribution gaussienne continue représente, selon nous, un problème d'apprentissage inutilement difficile. Nous proposons les Modèles de diffusion à variables latentes discrètes-continues (DisCo-Diff) afin de simplifier cette tâche en introduisant des variables latentes discrètes complémentaires. Nous enrichissons les DMs avec des latents discrètes apprenables, estimés par un encodeur, et entraînons conjointement le modèle de diffusion et l'encodeur. DisCo-Diff ne repose pas sur des réseaux pré-entraînés, ce qui confère au cadre une applicabilité universelle. Les latents discrets simplifient considérablement l'apprentissage de la transformation bruit-données complexe du DM en réduisant la courbure de l'ÉDO génératif du modèle. Un transformateur autoregressif supplémentaire modélise la distribution des latents discrets, une étape simple car DisCo-Diff n'exige qu'un petit nombre de variables discrètes avec des codebooks de taille réduite. Nous validons DisCo-Diff sur des données synthétiques, plusieurs tâches de synthèse d'images ainsi que sur le docking moléculaire, et constatons que l'introduction de latents discrets améliore de manière cohérente les performances du modèle. Par exemple, DisCo-Diff atteint des scores FID de pointe sur les jeux de données ImageNet-64/128 conditionnés par classe, utilisant un échantillonneur basé sur une ÉDO.