DASH: Dual-Branch Score Distillation for Guidance-Calibrated Compact Diffusion Models

Résumé Technique : DASH : Dual-Branch Score Distillation for Guidance-Calibrated Compact Diffusion Models

Énoncé du Problème
Les modèles de diffusion conditionnels par classe utilisant le Classifier-Free Guidance (CFG) s'appuient sur un réseau unique pour produire à la fois les prédictions conditionnelles ($\epsilon_y$) et non conditionnelles ($\epsilon_\emptyset$). Le mécanisme de guidage amplifie l'écart $\Delta = \epsilon_y - \epsilon_\emptyset$ pour orienter la génération. Cependant, la compression de ces modèles via la réduction de paramètres introduit une défaillance structurelle critique dans les objectifs de distillation existants.

La distillation standard supervise uniquement la prédiction composée guidée ($\tilde{\epsilon}_S \approx \tilde{\epsilon}_T$) ou la branche conditionnelle seule. Cela crée un problème d'optimisation sous-déterminé : plusieurs décompositions des branches de l'étudiant peuvent satisfaire la fonction de perte, incluant une solution dégénérée où l'écart de guidage s'effondre ($\Delta_S \to 0$). Dans cet état, les prédictions conditionnelles et non conditionnelles convergent, rendant le poids de guidage $w$ inefficace lors de l'inférence, même si la perte de formation composée est minimisée. De plus, les méthodes existantes ne parviennent pas à transférer le curriculum d'importance par pas de temps (importance weighting) appris par l'enseignant, forçant l'étudiant à réapprendre ce calendrier dans des budgets de distillation limités, ce qui corrompt souvent la trajectoire d'entraînement optimale.

Méthodologie : Cadre DASH
Le document propose DASH (Dual-Branch Score Distillation), un cadre conçu pour résoudre la sous-détermination et préserver le calibrage du guidage grâce à deux mécanismes principaux :

1. Supervision Indépendante à Double Branche :
   Au lieu de faire correspondre la sortie composée, DASH supervise indépendamment les deux branches de score en utilisant des termes de perte distincts :

   • Perte d'Imitation ($L_{im}$) : Fait correspondre la branche conditionnelle ($\epsilon^y_S \approx \epsilon^y_T$).
   • Perte Non Conditionnelle ($L_{un}$) : Fait correspondre la branche non conditionnelle ($\epsilon^\emptyset_S \approx \epsilon^\emptyset_T$).
   • Perte d'Ancrage ($L_{an}$) : Un régularisateur léger ancrant la branche conditionnelle au bruit de vérité terrain ($\epsilon$) pour éviter la dérive lors du début de l'entraînement.

   En imposant des contraintes sur les deux branches simultanément, le cadre spécifie de manière unique les sorties cibles pour chaque échantillon, éliminant l'espace de solutions réalisables qui permet l'effondrement de l'écart de guidage. Les auteurs définissent un ratio d'écart de guidage $\rho = E[\|\Delta_S\|] / E[\|\Delta_T\|]$ pour surveiller le calibrage, où $\rho \approx 1$ indique un succès de préservation.

2. Transfert TIRT (Timestep-Importance Rebalanced Training Transfer) :
   Le modèle enseignant utilise le TIRT, une extension du pondération Min-SNR qui apprend des facteurs d'échelle par pas de temps ($\lambda_t$) pour concentrer l'entraînement sur les régions de milieu de trajectoire difficiles en contenu. DASH introduit le TIRT Transfer, qui copie directement le curriculum d'importance par pas de temps convergé de l'enseignant ($\hat{\omega}_T$) dans l'étudiant en tant que prior gelé. Cela empêche l'étudiant de corrompre le curriculum avec des gradients de distillation qui récompensent la facilité de correspondance de prédiction plutôt que la difficulté de contenu.

Contributions Clés

• Résolution de la Sous-Détermination Structurelle : Le papier identifie et résout l'effondrement dégénéré de l'écart de guidage dans les modèles de diffusion compressés en paramètres. Il démontre que la supervision non conditionnelle n'est pas seulement bénéfique mais essentielle, représentant plus de 60 % du gain total de distillation.
• Transfert de Curriculum : L'introduction du TIRT Transfer, qui initialise l'étudiant avec le poids par pas de temps appris de l'enseignant comme un prior gelé, éliminant le besoin pour l'étudiant de réapprendre le curriculum.
• Cadre à Double Branche : Un pipeline de distillation unifié qui contraint indépendamment les branches conditionnelles et non conditionnelles, validé comme étant empiriquement essentiel pour maintenir la fidélité du guidage sous compression.

Résultats Expérimentaux
Les expériences ont été menées sur CIFAR-10 et CIFAR-100 en utilisant un modèle enseignant (35,8 M de paramètres) et un étudiant compressé (6,1 M de paramètres, compression de 5,9×).

• Performance : L'étudiant DASH atteint un FID de 8,87 (CIFAR-10) et 10,47 (CIFAR-100) à 50 étapes d'échantillonnage DDIM. Cela se situe à moins de 4 points FID de l'enseignant (FID 5,47 / 6,80) et surpasse nettement l'entraînement à partir de zéro (FID 20,47 / 26,14) et les bases à contrainte unique.
• Calibrage du Guidage : DASH maintient un ratio d'écart de guidage $\rho$ de 0,91 (CIFAR-10) et 0,89 (CIFAR-100), avec un cosinus directionnel $\cos(\Delta)$ proche de 0,94. En revanche, les bases de distillation composée souffrent d'un effondrement partiel ($\rho \approx 0,68$), et la supervision uniquement conditionnelle conduit à un effondrement quasi total ($\rho \approx 0,09$), dégradant la performance en dessous de l'entraînement à partir de zéro.
• Études d'Ablation :
  • Supprimer $L_{un}$ (perte non conditionnelle) provoque une chute de 7,77 points de FID sur CIFAR-10 et fait s'effondrer $\rho$ à 0,61, confirmant son rôle de composant dominant.
  • Supprimer le TIRT Transfer (en le remplaçant par un Min-SNR fixe) entraîne une chute de 2,60 points de FID, validant la nécessité de transférer le curriculum appris.
  • Le ratio de performance inter-jeu de données (1,18× pour DASH contre 1,51× pour les bases) confirme une généralisation robuste à travers les complexités de labels.

Signification et Revendications
Le papier affirme que DASH établit un nouveau standard pour la diffusion à paramètres efficaces en abordant une faille structurelle dans les objectifs de distillation existants. Les auteurs affirment que :

1. La supervision non conditionnelle est non négociable : Elle ne peut être substituée par des objectifs composites ou conditionnels ; sans elle, le CFG devient inefficace sous compression.
2. Le transfert de curriculum est distinct de la distillation de caractéristiques : Transférer le calendrier d'entraînement (TIRT) comme un prior gelé est une contribution novatrice qui empêche la corruption du poids adaptatif aux données de l'enseignant.
3. Efficacité sans Compromis : Le cadre permet une réduction significative des paramètres (5,9×) tout en préservant le mécanisme de guidage calibré, permettant une génération de haute qualité dans des environnements à ressources limitées sans augmenter les étapes d'inférence.

Ce travail se positionne à l'intersection de la compression de paramètres et du calibrage au niveau des branches, offrant une solution orthogonale aux méthodes de réduction d'étapes et composable avec elles.