Résumé Technique : Apprentissage Multi-Objectif pour les Modèles de Diffusion dans un Cadre Semi-Supervisé

1. Énoncé du Problème

Les modèles de diffusion sont de plus en plus déployés en tant que générateurs conditionnels pour diverses tâches, telles que la synthèse texte-image à travers des domaines d'invite hétérogènes ou le contrôle robotique à travers plusieurs environnements. Ces scénarios constituent naturellement un problème d'Apprentissage Multi-Objectif (MOL), où un modèle unique doit atteindre une performance Pareto-optimale à travers plusieurs distributions cibles.

Un défi central dans ce contexte est le coût statistique de l'entraînement d'un modèle « généraliste ». Pour gérer plusieurs objectifs, la classe de modèles généralistes ($\mathcal{F}$) nécessite généralement une capacité nettement supérieure à celle des classes de modèles « spécialistes » ($\mathcal{H}_k$) requises pour les tâches individuelles. Dans des cadres entièrement supervisés, la complexité en échantillons pour l'entraînement du généraliste évolue avec sa propre complexité ($C_\mathcal{F}$), ce qui est souvent prohibitivement élevé.

De plus, les déploiements réels font souvent face à un régime de données asymétrique :

• Les données appariées $(x, y)$ (par exemple, des paires texte-image curatées ou des démonstrations d'experts robotiques) sont rares et coûteuses à acquérir.
• Les données de condition non appariées $y$ (par exemple, des invites textuelles brutes ou des observations d'état) sont abondantes.

L'article aborde la question suivante : Pouvons-nous apprendre un modèle de diffusion Pareto-optimal de manière efficace en échantillons en exploitant des données faiblement supervisées (conditions non étiquetées) ?

2. Méthodologie

Les auteurs proposent un cadre d'entraînement semi-supervisé en deux étapes qui concilie le besoin d'un généraliste haute capacité avec la rareté des données appariées.

Étape 1 : Entraînement des Spécialistes

Pour chacun des $K$ objectifs (distributions cibles $P_k$), une fonction de score spécialiste légère $h_k \in \mathcal{H}_k$ est entraînée en utilisant uniquement les données appariées limitées disponibles $\{(x_i^k, y_i^k)\}_{i=1}^n$.

• Les classes de spécialistes $\mathcal{H}_k$ sont supposées avoir une complexité significativement inférieure ($C_{\mathcal{H}_k} \ll C_\mathcal{F}$) à celle de la classe généraliste.
• L'entraînement minimise l'objectif d'appariement de score de débruitage sur les données appariées.

Étape 2 : Entraînement du Généraliste via Échantillonnage Pseudo

Un grand modèle généraliste $f \in \mathcal{F}$ est entraîné pour approximer l'ensemble Pareto-optimal à travers tous les objectifs.

• Génération de Pseudo-Échantillons : Les spécialistes entraînés $\{h_k\}$ sont utilisés pour générer des données pseudo-appariées. Pour chaque condition non étiquetée abondante $y_i^k \sim P_Y^k$, le spécialiste génère une pseudo-cible $\tilde{x}_i^k \sim P_{h_k}(\cdot|y_i^k)$.
• Objectif : Le généraliste est entraîné sur une combinaison des données étiquetées originales et des données pseudo-appariées générées. L'objectif d'entraînement minimise une scalarisation $S$ des pertes d'appariement de score à travers les $K$ objectifs :
  $$\hat{f}_S = \arg\min_{f \in \mathcal{F}} S(\tilde{L}_1(f), \dots, \tilde{L}_K(f))$$
  où $\tilde{L}_k$ est la perte empirique calculée sur les pseudo-échantillons (avec un terme de correction pour tenir compte du biais du spécialiste).

Extension à la Prise de Décision Séquentielle

Le cadre est étendu aux politiques de diffusion dans l'apprentissage par imitation. Ici, la « condition » est l'état $s$, et la « cible » est l'action $a$. Les auteurs abordent le décalage de distribution inhérent aux déroulements en politique : puisque les distributions d'état expertes sont difficiles à échantillonner directement, la deuxième étape utilise des trajectoires générées par les politiques spécialistes apprises sous les dynamiques de l'environnement pour collecter des pseudo-données pour la politique généraliste.

3. Contributions Théoriques Clés

Bornes de Complexité en Échantillons

L'article établit des bornes de généralisation non asymptotiques pour l'approche semi-supervisée proposée.

• Résultat Principal (Théorème 3.2) : Le nombre de données appariées (étiquetées) requis pour atteindre une erreur cible $\epsilon$ dépend uniquement de la complexité des classes de spécialistes ($C_{\mathcal{H}_k}$), et non de la classe généraliste ($C_\mathcal{F}$).
  • Échantillons étiquetés requis : $\tilde{O}\left(\frac{K C_\mathcal{H}}{\epsilon^2}\right)$.
  • Conditions non étiquetées requises : $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$.
• Comparaison : En revanche, un entraînement entièrement supervisé classique du généraliste nécessiterait $\tilde{O}\left(\frac{K C_\mathcal{F}}{\epsilon^2}\right)$ échantillons étiquetés. Cela démontre une amélioration significative de l'efficacité en échantillons lorsque $C_\mathcal{F} \gg C_\mathcal{H}$.
• Scalarisations Linéaires (Théorème 3.5) : Pour les scalarisations linéaires (somme pondérée) et les classes de fonctions convexes, les auteurs dérivent des taux plus précis, affinant davantage la borne sur l'erreur d'approximation.

Estimation de Distribution et Garanties de Politique

• Estimation de Distribution : L'article comble le fossé entre l'erreur d'appariement de score et l'estimation de distribution (distance de variation totale). Il prouve que le généraliste appris atteint une performance Pareto-optimale dans l'approximation des distributions conditionnelles cibles.
• Politiques de Diffusion (Théorème 3.6) : Les auteurs fournissent la première garantie théorique sur l'écart de sub-optimalité des politiques de diffusion dans l'apprentissage par imitation dans des cadres semi-supervisés. Ils montrent que le cadre gère avec succès le décalage de distribution causé par les déroulements en politique tout en conservant les avantages d'efficacité en échantillons du pipeline spécialiste-vers-généraliste.

4. Résultats Expérimentaux

Les auteurs valident leur théorie sur deux domaines distincts :

1. Manipulation Robotique (StackCube & PickCube) :

   • Configuration : Entraînement de politiques de diffusion pour des tâches sous randomisation de domaine variable (scène, éclairage, pose de caméra).
   • Résultat : La méthode MOL semi-supervisée surpasse nettement une base de référence uniquement étiquetée, en particulier dans des configurations Hors Distribution (OOD) (par exemple, +16 % de taux de réussite sur StackCube Niveau 2, +14 % sur le Niveau 3). Cela suggère que les pseudo-déroulements fournissent une supervision cruciale pour gérer les décalages de distribution.

2. Restauration d'Image (Inpainting CelebA-HQ) :

   • Configuration : Inpainting de visages avec trois types de masques différents (lignes alternées, région du visage, trait large).
   • Résultat : L'approche semi-supervisée a amélioré le PSNR, le SSIM et le LPIPS pour tous les types de masques par rapport à la base de référence uniquement étiquetée, démontrant l'efficacité de la méthode dans les tâches d'images génératives.

5. Signification et Revendications

L'article revendique fournir un cadre statistique principiel pour l'apprentissage multi-objectif dans les modèles de diffusion sous régimes semi-supervisés. Sa signification principale réside dans :

• Liaison entre Théorie et Pratique : Il formalise le pipeline de distillation « spécialiste-vers-généraliste », offrant une justification théorique de la raison pour laquelle cette approche est plus efficace en échantillons que l'apprentissage multi-tâche direct lorsque les données étiquetées sont rares.
• Efficacité en Échantillons : Il prouve que le coût statistique de l'apprentissage d'un généraliste complexe peut être découplé de sa propre capacité, reposant plutôt sur la complexité de spécialistes légers entraînés sur des données limitées.
• Gestion du Décalage de Distribution : Il étend ces garanties à la prise de décision séquentielle (apprentissage par imitation), abordant un défi pratique critique en robotique.

Les auteurs adoptent une position modeste, reconnaissant que leurs garanties reposent sur une hypothèse de réalisabilité (que les vraies fonctions de score se trouvent au sein des classes de spécialistes) et qu'ils n'abordent pas actuellement la généralisation Hors Distribution (OOD) au-delà des niveaux de randomisation de domaine spécifiques testés. Ils considèrent ce travail comme une étape fondamentale vers la compréhension des propriétés statistiques du MOL semi-supervisé pour les modèles génératifs.