Calibrating Generative Models to Feature Distributions with MMD Finetuning

Résumé Technique : Calibrer les modèles génératifs aux distributions de caractéristiques avec le fine-tuning MMD (kCGM)

Énoncé du problème
Les modèles génératifs produisent souvent des échantillons qui sont individuellement plausibles mais qui présentent un désalignement distributionnel concernant des caractéristiques clés d'intérêt. Par exemple, un modèle préentraîné sur un espace chimique large peut générer des molécules qui diffèrent significativement de la distribution de caractéristiques d'une classe thérapeutique spécifique, telle que les antibiotiques. Corriger cela via un fine-tuning direct sur l'ensemble cible est problématique : cela risque de provoquer un surapprentissage (overfitting) sur de petits jeux de données, nécessite l'intégralité des échantillons cibles (qui peuvent être rares), et ne permet pas de spécifier quelles caractéristiques précises doivent être égalées. De plus, les méthodes existantes peinent souvent avec les caractéristiques à valeurs discrètes ou "boîte noire" (par exemple, les empreintes Morgan) qui sont non différentiables. L'article identifie une lacune dans les algorithmes capables de calibrer les modèles génératifs pour correspondre à une distribution de caractéristiques cible en utilisant uniquement une supervision au niveau des caractéristiques, sans nécessiter la différentiabilité de la carte de caractéristiques (feature map) ou de la procédure d'échantillonnage.

Méthodologie : kCGM
Les auteurs introduisent les modèles génératifs de calibration par noyau (kCGM), une méthode de fine-tuning qui minimise la divergence de Maximum Mean Discrepancy (MMD) entre la distribution de caractéristiques générée et une distribution de caractéristiques cible, régularisée par la divergence de Kullback-Leibler (KL) par rapport au modèle préentraîné.

• Fonction d'objectif : La méthode optimise une fonction d'objectif $L_\theta$ composée de deux termes :

  1. Perte MMD : Un estimateur non biaisé par échantillonnage fini du carré de la MMD entre la distribution de caractéristiques induite par le modèle $P$ (où $y = h(x)$ et $x \sim p_\theta$) et la distribution cible empirique $Q$ (provenant des caractéristiques cibles $y^{targ}_{1:N}$).
  2. Régularisation KL : Un terme $\lambda D_{KL}(p_\theta || p_{\theta_{pre}})$ pour garantir que le modèle ajusté reste proche de la distribution préentraînée, évitant ainsi l'oubli catastrophique ou l'effondrement de mode (mode collapse).

• Sélection du noyau : Contrairement aux approches précédentes qui pourraient s'appuyer sur des noyaux de produit scalaire (qui ne correspondent que les moyennes), kCGM utilise des noyaux qui définissent une métrique sur les distributions de probabilité (nulle uniquement pour des distributions identiques).

  • Pour les caractéristiques euclidiennes continues, un noyau de distance d'énergie (energy-distance kernel) est utilisé.
  • Pour les caractéristiques de vecteurs de bits discrets (ex: empreintes Morgan), un noyau de Tanimoto/Jaccard est employé.
  • Ces choix évitent d'introduire des hyperparamètres de bande passante supplémentaires.

• Estimation du gradient : Puisque la carte de caractéristiques $h$ et le processus d'échantillonnage peuvent être non différentiables, kCGM emploie des gradients de fonction de score (de type REINFORCE).

  • Base de référence Leave-One-Out (LOO) : Une contribution technique cruciale est la dérivation d'une nouvelle base de référence LOO pour la perte MMD. Les bases de référence LOO standards moyennent les coefficients d'autres échantillons dans un batch, mais les coefficients MMD sont couplés entre les échantillons. Les auteurs dérivent un estimateur spécifique "MMD-LOO" qui approxime le coefficient moyen comme si un échantillon spécifique n'était pas apparu dans le mini-batch, réduisant la variance sans introduire de biais.
  • Algorithme : Le processus implique l'échantillonnage à partir du modèle actuel (sans suivi de gradient), le calcul des vecteurs de caractéristiques, le calcul des coefficients MMD et KL avec les bases de référence LOO, et la mise à jour des poids du modèle via des gradients de fonction de score.

Contributions clés

1. Formulation : L'article formule la correspondance de distribution de caractéristiques comme un fine-tuning MMD régularisé par KL en utilisant uniquement des échantillons de caractéristiques cibles, éliminant le besoin d'échantillons cibles complets dans l'espace d'origine.
2. Estimateur non biaisé : Les auteurs dérivent un estimateur de gradient de fonction de score non biaisé pour l'objectif MMD, permettant l'utilisation de cartes de caractéristiques non différentiables et d'échantillonneurs non reparamétrisables (ex: modèles autorégressifs, diffusion discrète).
3. Généralisation : kCGM généralise le modèle CGM. Alors que CGM fait correspondre les moyennes de caractéristiques cibles, kCGM fait correspondre la distribution complète via la MMD. L'article note qu'utiliser un noyau de produit scalaire dans kCGM permet de retrouver l'objectif de relaxation de CGM.

Résultats expérimentaux
Les auteurs évaluent kCGM à travers trois domaines : la génération de petites molécules, la génération de structures protéiques et la génération d'ADN régulateur.

• Antibiotiques (Petites Molécules) :

  • Tâche : Calibrer un modèle de graphe autorégressif G2PT pour correspondre à la distribution de caractéristiques de 174 antibiotiques connus.
  • Caractéristiques : Empreintes Morgan (substructures locales), squelettes Murcko (topologie) et descripteurs physico-chimiques.
  • Résultats : kCGM atteint un compromis supérieur entre la correspondance avec la distribution de caractéristiques cible (mesurée par MMD) et la distance par rapport au modèle préentraîné par rapport à CGM et au fine-tuning direct. Crucialement, alors que le fine-tuning direct sur le petit ensemble d'antibiotiques entraînait une proportion élevée de molécules chimiquement invalides, kCGM a amélioré la correspondance des caractéristiques tout en augmentant la validité chimique. kCGM a également démontré un transfert de caractéristiques croisées, améliorant la correspondance des squelettes et des descripteurs même lorsqu'il n'a été entraîné que sur les empreintes.

• Diversité des structures protéiques :

  • Tâche : Fine-tuner Genie 2 (un modèle de diffusion de protéines) pour correspondre à la distribution de la structure secondaire (hélices alpha vs feuillets bêta) de l'ensemble de données CATH.
  • Résultats : Le modèle préentraîné produisait principalement des hélices alpha. kCGM a réduit le désaccord avec la distribution cible d'environ 2 fois par rapport à CGM. kCGM a montré une courbe de compromis fluide avec le poids KL ($\lambda$), contrairement à CGM dont la performance était moins sensible à $\lambda$, rendant l'optimisation difficile.

• Génération d'ADN régulateur :

  • Tâche : Fine-tuner un modèle de diffusion discrète pour générer des séquences d'ADN avec des profils d'activité cellulaires spécifiques réalistes (prédits par AlphaGenome).
  • Résultats : kCGM a réussi à faire correspondre la distribution complète des profils d'activité pour 47 sujets de régulation cis différents, surpassant CGM (qui ne fait correspondre que les moyennes). Le modèle fine-tuné a généré des séquences avec des profils d'activité similaires à de vrais enhancers, tandis que le modèle préentraîné générait des séquences sous-actives.

Signification et affirmations
L'article affirme que kCGM fournit une alternative robuste au fine-tuning direct, particulièrement lorsque les jeux de données cibles sont petits ou que seuls certains aspects d'une distribution doivent être préservés. Sa principale importance réside dans sa capacité à :

• Aligner les modèles génératifs sur des distributions de caractéristiques spécifiées par l'utilisateur en utilisant uniquement une supervision au niveau des caractéristiques.
• Gérer des caractéristiques "boîte noire" et non différentiables (ex: empreintes, sorties de prédicteurs) qui sont courantes dans les domaines scientifiques mais incompatibles avec le fine-tuning standard basé sur le gradient.
• Maintenir ou améliorer la qualité des échantillons (validité) tout en corrigeant le désalignement distributionnel, un mode d'échec courant du fine-tuning direct.
• Être applicable à diverses architectures, incluant les modèles autorégressifs, la diffusion en espace continu et la diffusion discrète.

Les auteurs notent des limites, telles qu'une potentielle inefficacité pour les caractéristiques de très haute dimension due à la malédiction de la dimensionnalité dans les tests d'hypothèse MMD, et la nécessité de probabilités logarithmiques de trajectoire d'échantillonnage tractables (excluant les modèles de flow matching). Ils suggèrent que les travaux futurs pourraient étendre ce cadre à des contextes où seules des contraintes distributionnelles existent sans échantillons complets dans l'espace génératif d'origine.