Résumé Technique : Calibration des Modèles Génératifs selon des Contraintes Distributionnelles

Énoncé du Problème

Les modèles génératifs souffrent fréquemment d'un défaut de calibration, où les statistiques de la distribution d'échantillonnage (par exemple, les fréquences de classes, les proportions de structures secondaires ou la représentation des genres) s'écartent systématiquement des valeurs souhaitées. Ce problème se manifeste dans divers domaines : les modèles d'images présentent un effondrement de modes ; les modèles de langage renforcent les biais sociétaux ; et les modèles de protéines/ADN génèrent des structures avec des fréquences de sous-structures atypiques. Ces erreurs découlent souvent de déséquilibres dans les jeux de données, de dynamiques d'entraînement sous-optimales ou d'ajustements a posteriori tels que l'échantillonnage à basse température.

Les auteurs formulent la calibration comme un problème d'optimisation sous contraintes. Étant donné un modèle génératif de base $p_{\theta_{base}}(x)$, une statistique $h(x)$ et une espérance cible $h^*$, l'objectif est de trouver une nouvelle distribution $p_\theta$ satisfaisant $E_{p_\theta}[h(x)] = h^*$ tout en minimisant la divergence de Kullback-Leibler (KL) par rapport au modèle de base :
$$p_{\theta^*} = \arg \min_{p_\theta} D_{KL}(p_\theta \parallel p_{\theta_{base}}) \quad \text{s.t.} \quad E_{p_\theta}[h(x)] = h^*$$
Cette formulation cherche le modèle le plus proche de la distribution de base qui satisfait la contrainte de calibration. Cependant, résoudre ce problème exactement est intraitable car l'objectif et la contrainte impliquent tous deux des espérances difficiles à calculer pour les modèles génératifs de haute dimension.

Méthodologie

Pour pallier l'intratabilité, l'article introduit deux objectifs de substitution pour le fine-tuning, nommés CGM-relax et CGM-reward. Les deux méthodes reposent sur une optimisation stochastique utilisant des estimations de gradient non biaisées dérivées d'échantillons.

1. CGM-relax (Perte de Relaxation)

Cette approche remplace la contrainte rigide par un terme de pénalité. La fonction objectif est :
$$L_{relax}(\theta) = \| E_{p_\theta}[h(x)] - h^* \|^2 + \lambda D_{KL}(p_\theta \parallel p_{\theta_{base}})$$

• Violation de la contrainte ($L_{viol}$) : La norme $\ell_2$ au carré de la différence entre la statistique empirique et la cible. Les auteurs dérivent un estimateur non biaisé pour ce terme en utilisant une formule de variance d'échantillon corrigée du biais.
• Divergence KL ($L_{KL}$) : Mesurée par échantillonnage de Monte Carlo.
• Optimisation : La méthode utilise l'échantillonnage par importance pour calculer des gradients non biaisés. Plus précisément, elle échantillonne à partir du modèle actuel $p_\theta$ (avec arrêt du gradient) et répond les échantillons pour estimer les gradients du terme KL et de la violation de contrainte. Un hyperparamètre $\lambda$ arbitre entre la satisfaction de la contrainte et le maintien de la proximité avec le modèle de base.

2. CGM-reward (Perte de Récompense)

Cette approche relie le problème de calibration au principe de Maximisation de l'Entropie. La solution au problème de maximisation de l'entropie (minimisation de la KL par rapport au modèle de base sous contraintes de moments) est une déformation exponentielle du modèle de base : $p_{\alpha^*}(x) \propto p_{\theta_{base}}(x) \exp(\alpha^{*\top} h(x))$.

• Estimation de $\alpha^*$ : Puisque la vraie distribution de maximisation de l'entropie est intraitable, la méthode estime d'abord les paramètres de récompense optimaux $\alpha^*$ en résolvant un problème d'optimisation duale utilisant $N$ échantillons du modèle de base.
• Fonction de Perte : L'objectif de fine-tuning est la divergence KL entre le modèle actuel $p_\theta$ et la distribution de maximisation de l'entropie estimée $p_{\hat{\alpha}_N}$ :
  $$L_{reward}(\theta) = D_{KL}(p_\theta \parallel p_{\hat{\alpha}_N}) = D_{KL}(p_\theta \parallel p_{\theta_{base}}) + E_{p_\theta}[-r_{\hat{\alpha}_N}(x)] + C$$
  Cela équivaut à un problème de fine-tuning par récompense où la « récompense » est dérivée des contraintes.
• Optimisation : De manière similaire à CGM-relax, elle utilise des estimateurs de gradient par fonction de score avec des baselines leave-one-out pour réduire la variance.

Estimation du Gradient

Les deux méthodes emploient des estimateurs de gradient par fonction de score (style REINFORCE) combinés à l'échantillonnage par importance. Pour gérer la dépendance de l'espérance par rapport à $\theta$, les auteurs utilisent une astuce « stop-gradient » où les échantillons sont tirés d'une version fixe du modèle actuel, et les gradients sont calculés via le rapport de densité (fonction de score). La réduction de variance est obtenue grâce à des baselines leave-one-out.

Contributions Clés

1. Formulation : Présenter la calibration des modèles génératifs comme un problème de minimisation de KL sous contraintes.
2. Algorithmes : Proposer deux algorithmes pratiques de fine-tuning (CGM-relax et CGM-reward) qui approximent la solution du problème contraint en utilisant l'optimisation stochastique.
3. Évolutivité : Démontrer que ces méthodes peuvent gérer des centaines de contraintes simultanées et des modèles allant jusqu'à 9 milliards de paramètres.
4. Lien Théorique : Établir le lien entre calibration, maximisation de l'entropie et fine-tuning par récompense, montrant que CGM-reward vise la solution de maximisation de l'entropie tandis que CGM-relax vise une version régularisée.

Résultats

Les auteurs évaluent les méthodes sur trois applications diverses :

1. Conception de Protéines (Genie2 et ESM3) :

   • Objectif : Calibrer le contenu en structures secondaires (hélices alpha et brins bêta) pour correspondre aux statistiques des protéines naturelles (jeu de données CATH).
   • Contraintes : Jusqu'à 99 contraintes simultanées (CDF bivariées).
   • Résultat : CGM-relax a réduit la distance aux distributions de protéines naturelles de plus de 4x pour Genie2 et de plus de 2x pour ESM3. CGM-reward a obtenu des améliorations plus modestes, en partie dû à la difficulté d'estimer $\alpha^*$ avec des contraintes de haute dimension. Les deux méthodes ont maintenu la plausibilité biophysique, bien qu'avec une légère augmentation des « échecs de conception ».

2. Génération d'Images (TarFlow) :

   • Objectif : Équilibrer les proportions de classes au sein de la catégorie « faune sauvage » du jeu de données AFHQ (par exemple, égaliser lions, tigres, renards, loups).
   • Contraintes : 5 contraintes simultanées.
   • Résultat : CGM-relax a réduit la distance de variation totale des proportions de classes de 0,306 à 0,101. Bien que le réalisme des échantillons (FID) ait légèrement diminué, le modèle calibré a produit des générations plus diversifiées et équilibrées. CGM-reward n'a pas réussi à réduire significativement la violation des contraintes.

3. Modélisation du Langage (Gemma-2-9B-IT) :

   • Objectif : Réduire le biais de genre dans les histoires pour enfants générées pour diverses professions (par exemple, assurer une probabilité égale d'avocats masculins/féminins).
   • Contraintes : 8 contraintes conditionnelles par profession.
   • Résultat : Les deux méthodes ont réduit le déséquilibre de genre. CGM-relax a réduit le déséquilibre de plus de trois fois par rapport à la ligne de base, tandis que CGM-reward a montré une amélioration statistiquement significative mais plus faible. Les modèles fine-tunés ont maintenu la qualité des histoires (aucune dégradation de la log-probabilité).

Expériences Synthétiques :

• Événements Rares : CGM a réussi à surpondérer les modes rares dans des mélanges gaussiens, même lorsque la probabilité de base était aussi faible que $10^{-3}$.
• Hautes Dimensions : CGM-relax a géré efficacement jusqu'à $10^3$ contraintes. En revanche, CGM-reward est devenu inapplicable pour des contraintes $>30$ en raison de la difficulté de résoudre le problème empirique de maximisation de l'entropie (Équation 7).

Importance et Revendications

L'article revendique que CGM fournit un cadre pratique pour calibrer des modèles génératifs de haute dimension afin de respecter des contraintes distributionnelles sans nécessiter de modifications de l'architecture du modèle ni de réentraînement à partir de zéro.

• Robustesse : CGM-relax s'avère plus robuste que CGM-reward, en particulier dans des contextes de contraintes de haute dimension où l'estimation des paramètres de maximisation de l'entropie devient intraitable.
• Compromis : Les méthodes offrent un compromis navigable entre la précision de la calibration et la fidélité au modèle de base (contrôlée par $\lambda$).
• Limites : Les auteurs reconnaissent que des erreurs résiduelles subsistent, en particulier pour les événements très rares. De plus, le cadre repose actuellement sur des modèles ayant des vraisemblances traitables (modèles de diffusion, flows normalisants, modèles autorégressifs), limitant son application directe aux modèles implicites comme les GAN.

Ce travail positionne la calibration comme une direction de recherche distincte et fructueuse, séparée du fine-tuning par récompense standard, en utilisant les contraintes elles-mêmes comme signal moteur pour l'alignement distributionnel.