Résumé Technique : Expansion du Domaine de Diffusion (DDE)

Énoncé du Problème

Les modèles de diffusion pré-entraînés ont obtenu des résultats de pointe dans la génération d'images, d'audio et d'autres modalités de données. Cependant, réentraîner ces modèles à partir de zéro pour gérer des domaines de sortie plus vastes (par exemple, des images panoramiques, des pistes audio longues) ou des entrées de conditionnement plus complexes (par exemple, plusieurs emplacements d'objets) est coûteux en termes de calcul. Les approches existantes pour étendre les domaines génératifs reposent souvent sur des mécanismes de coordination algorithmiques fixes (par exemple, MultiDiffusion) qui ne nécessitent pas d'entraînement supplémentaire mais sont limités à des tâches spécifiques et manquent de flexibilité. D'autres méthodes, telles que l'entraînement de modules de diffuseurs dynamiques, nécessitent souvent des modalités fixes au moment de l'entraînement ou échouent à se généraliser à des domaines plus vastes que ceux observés durant l'entraînement.

Le défi central abordé par cet article est de savoir comment étendre efficacement le domaine génératif des modèles de diffusion pré-entraînés pour générer des objets plus grands et gérer un conditionnement plus complexe sans réentraîner les modèles de base, tout en garantissant que la solution se généralise à des instances plus grandes que celles rencontrées durant l'entraînement du coordinateur.

Méthodologie : Expansion du Domaine de Diffusion (DDE)

Les auteurs proposent l'Expansion du Domaine de Diffusion (DDE), un cadre qui emploie un réseau « coordinateur » compact et entraînable pour orchestrer les sorties de modèles de diffusion pré-entraînés.

Formulation Générale

Étant donné un modèle de diffusion conditionnel pré-entraîné $p(x|y)$ avec un réseau de débruitage $D(x, y, t)$, l'objectif est de générer un objet étendu $X^{[L]}$ conditionné par une entrée étendue $Y^{[L]}$, où $L$ représente l'échelle (par exemple, le nombre de patches ou d'étiquettes de conditionnement).

1. Décomposition : La paire étendue $(X^{[L]}, Y^{[L]})$ est décomposée en $L$ parties plus petites $([x_1, \dots, x_L], [y_1, \dots, y_L])$ à l'aide d'une fonction de mappage $F$. Ces parties relèvent du domaine original du modèle de base.
2. Génération de Base : Chaque partie $x_i$ est traitée par le débruiteur pré-entraîné $D(x_i, y_i, t)$.
3. Coordination : Un réseau coordinateur entraînable $C[\cdot]$ prend en entrée la séquence des sorties débruitées $[D(x_i, y_i, t)]_{i=1}^L$ et les informations de conditionnement $[y_i]_{i=1}^L$. Il apprend à coordonner ces sorties pour produire une estimation cohérente de l'objet étendu complet $X^{[L]}(0)$.
4. Objectif d'Entraînement : Le coordinateur est entraîné en minimisant l'erreur de débruitage du débruiteur composite $D^{[L]}$ sur un ensemble de données de taille $L_{train}$. Crucialement, la méthode est conçue pour se généraliser à des tailles de test $L_{test} \ge L_{train}$.

Architecture du Coordinateur

Les auteurs utilisent une architecture de Vision Transformer (ViT) pour le coordinateur, adaptée de Peebles & Xie (2023).

• Traitement de l'Entrée : Les patches débruités provenant des modèles pré-entraînés sont « patchifiés » et augmentés d'encodages de position globaux par rapport à l'objet plus grand. Les auteurs emploient des Encodages de Position Rotatifs (RoPE) pour gérer efficacement des dimensions spatiales plus grandes.
• Réconciliation : Durant le processus direct, le transformateur traite tous les patches. Pour gérer les patches se chevauchant (utilisés pour assurer la continuité), la sortie finale est construite en moyennant les valeurs dans les régions de chevauchement, inspiré par la méthode d'échantillonnage MultiDiffusion.
• Efficacité : Le coordinateur est significativement plus petit que le modèle de base (par exemple, 66M contre 405M de paramètres pour la génération de musique) et nécessite moins d'époques d'entraînement.

Contributions Clés

1. Cadre Novel : Introduction du DDE, une méthode qui utilise un petit coordinateur entraînable pour étendre les capacités génératives des modèles de diffusion pré-entraînés vers des objets plus grands et un conditionnement plus complexe.
2. Capacité de Généralisation : Démonstration que le coordinateur peut se généraliser pour générer des objets et gérer des entrées de conditionnement plus grandes en échelle que celles présentes dans ses données d'entraînement ($L_{test} \ge L_{train}$).
3. Conception Architecturale : Validation d'une architecture basée sur ViT comme coordinateur approprié, qui apprend les dépendances à longue portée tout en imposant une cohérence locale par le biais de la moyenne des chevauchements.
4. Efficacité : L'approche atteint une génération de haute qualité avec nettement moins de paramètres et de temps d'entraînement par rapport au réentraînement des modèles de base ou à l'utilisation d'architectures de coordinateurs plus grandes.

Résultats Expérimentaux

Les auteurs ont évalué le DDE sur trois domaines :

1. Génération de Pistes Audio Longues (Slakh2100)

• Tâche : Générer des pistes musicales d'une longueur de $10l$ (120 secondes) en utilisant un modèle de base entraîné sur une longueur $l$ (12 secondes).
• Lignes de Base : Concaténation et MultiDiffusion.
• Résultats : Le DDE a obtenu une Distance Audio Fréchet (FAD) de 2,112 pour les pistes 4l et 2,142 pour les pistes 10l. Cela surpasse significativement MultiDiffusion (4,796) et la Concaténation (4,596) pour la tâche 10l. Le coordinateur (66M de paramètres) était beaucoup plus petit que le modèle de base (405M).

2. Génération d'Images Conditionnelle (CLEVR)

• Tâche : Générer des images conditionnées par plusieurs emplacements d'objets ($L$ étiquettes), là où le modèle de base ne gère qu'un seul emplacement.
• Lignes de Base : RRR (Reduce, Reuse, Recycle) et MultiDiffusion.
• Résultats : Le DDE a démontré une précision supérieure dans la satisfaction de multiples contraintes de conditionnement. Pour 5 conditionnements, le DDE a atteint 44,5 % de précision (échantillonneur Heun), contre 33,6 % pour RRR et 36,3 % pour MultiDiffusion. Le coordinateur s'est généralisé de l'entraînement sur $L_{train}=2$ au test sur $L_{test}=5$.

3. Génération d'Images Satellitaires

• Tâche : Générer de grandes images satellites conditionnées par des cartes schématiques.
• Lignes de Base : MultiDiffusion.
• Résultats : Le DDE a obtenu un score FID inférieur à celui de MultiDiffusion pour les tailles en distribution ($N=96$, FID 31,75 contre 37,81) et hors distribution ($N=128$, FID 27,37 contre 35,01). L'inspection visuelle a montré que le DDE produisait des images plus cohérentes, sans les coutures visibles observées dans les échantillons MultiDiffusion.

Importance et Revendications

L'article revendique que le DDE offre une solution pratique et efficace pour mettre à l'échelle les modèles de diffusion sans le coût prohibitif du réentraînement. En apprenant à coordonner des processus pré-entraînés, la méthode permet :

• Expansion du Domaine : Génération d'objets plus grands (audio plus long, images plus grandes) et gestion de conditions plus complexes (plusieurs étiquettes) au-delà des capacités du modèle original.
• Généralisation : La capacité d'extrapoler vers des échelles plus grandes que celles rencontrées durant l'entraînement du coordinateur.
• Qualité et Cohérence : Le coordinateur basé sur ViT apprend efficacement les dépendances à longue portée, résolvant les conflits entre les étapes de débruitage locales pour produire des sorties globalement cohérentes.

Les auteurs notent que la limitation principale est l'exigence de données supplémentaires et de temps pour entraîner le modèle coordinateur, bien que ce coût soit substantiellement inférieur à l'entraînement d'un modèle de diffusion complet à partir de zéro. Un travail futur est proposé pour appliquer le DDE à la génération de séquences de protéines et à la création de vidéos à partir de courts clips.