Toward Better Geometric Representations for Molecule Generative Models

Résumé technique : LENSEs – Amélioration Latente pour les Encodages Structuraux Non-Lisses

Énoncé du Problème
La génération conditionnée par des représentations géométriques (GeoRCG) s'est imposée comme un paradigme efficace pour la génération de molécules 3D, découplant la modélisation des représentations moléculaires de la génération des structures 3D. Dans ce cadre, un encodeur moléculaire préentraîné (par exemple, UniMol, Frad) définit un espace latent, un Modèle de Diffusion de Représentation (RDM) échantillonne un vecteur au niveau de la molécule depuis cet espace, et un générateur reconstruit la structure 3D conditionnée par ce vecteur.

Cependant, l'efficacité de GeoRCG est fondamentalement limitée par la qualité de l'espace latent fourni par les encodeurs préentraînés. Les auteurs identifient deux goulots d'étranglement critiques :

1. Variétés de Représentation Non-Lisses : Les encodeurs préentraînés, optimisés principalement pour la prédiction de propriétés via le débruitage de coordonnées, produisent des représentations avec des constantes de Lipschitz élevées. De petites perturbations dans la conformation moléculaire entraînent des déplacements disproportionnés dans l'espace latent. Cette non-lissitude rend la distribution difficile à modéliser pour le RDM, nécessitant souvent des mécanismes compensatoires tels que l'injection de bruit ajustée manuellement.
2. Sémantiques Multi-niveaux Sous-exploitées : Les encodeurs préentraînés présentent des caractéristiques hiérarchiques où les couches peu profondes capturent des motifs structuraux locaux (par exemple, des groupes fonctionnels) et les couches profondes encodent des sémantiques globales. Les pipelines GeoRCG existants utilisent typiquement uniquement la représentation de la dernière couche, jetant l'information sémantique riche préservée dans les couches antérieures. Par conséquent, le générateur doit réapprendre ces sémantiques à partir de zéro, augmentant la difficulté d'optimisation.

Méthodologie : LENSEs
Pour résoudre ces problèmes, les auteurs proposent LENSEs (Amélioration Latente pour les Encodages Structuraux Non-Lisses), un cadre qui affine activement l'espace latent de conditionnement durant l'entraînement génératif. Plutôt que de dépendre passivement d'un encodeur figé, LENSEs introduit trois mécanismes complémentaires :

1. Tête de Représentation avec Pooling de Couches Apprenable :
   Au lieu d'utiliser une sortie déterministe unique de l'encodeur, LENSEs agrège des caractéristiques multi-niveaux. Il emploie un mécanisme de pooling de couches apprenable pour pondérer les représentations provenant de différentes profondeurs de l'encodeur ($H^{(l)}$). Ces caractéristiques poolées sont mappées aux paramètres d'une posteriori gaussienne diagonale ($\mu, \sigma$) via des têtes de projection. Le code latent $z$ est échantillonné en utilisant l'astuce de reparamétrisation. Cette tête est entraînée conjointement avec le générateur pour extraire des représentations à la fois riches sémantiquement et favorables à la génération, régularisées par une perte de divergence KL pour assurer un espace latent compact.

2. Perte de Perception Moléculaire ($L_{PERC}$) :
   Inspirée des pertes de perception en vision par ordinateur, cette perte aligne la représentation sémantique de haut niveau de la molécule débruitée (générée) avec celle de la molécule propre de référence (ground-truth) dans l'espace de caractéristiques de l'encodeur figé. En faisant passer à la fois les molécules générées et propres à travers l'encodeur et en minimisant la distance entre leurs représentations au niveau du graphe poolées, le générateur est encouragé à préserver des motifs structuraux significatifs (par exemple, des groupes fonctionnels) plutôt que de simplement minimiser les erreurs de coordonnées de bas niveau. Une opération de stop-gradient est appliquée à la représentation de référence.

3. Perte d'Alignement de Représentation (REPA) ($L_{REPA}$) :
   Motivée par l'observation que les états cachés du générateur s'alignent naturellement avec les caractéristiques de l'encodeur, cette perte aligne explicitement les caractéristiques au niveau des nœuds du générateur avec celles de l'encodeur préentraîné. Un projecteur apprenable mappe les caractéristiques du générateur dans l'espace de caractéristiques de l'encodeur, et la similarité cosinus est maximisée entre les nœuds correspondants. Cela réduit l'écart sémantique entre le préentraînement et la génération, accélérant l'entraînement et améliorant l'acquisition de connaissances chimiquement significatives.

L'objectif d'entraînement total combine la perte de génération de base ($L_{gen}$) avec la perte KL, la perte de perception et la perte REPA.

Résultats Clés
Les auteurs évaluent LENSEs sur les ensembles de données GEOM-DRUG et QM9, le comparant à des références de l'état de l'art incluant GeoRCG, CanonFlow et SemlaFlow.

• Qualité de Génération : Sur l'ensemble de données difficile GEOM-DRUG, LENSEs atteint des performances de l'état de l'art. Il obtient 97,28 % de validité et 98,51 % de stabilité moléculaire, surpassant le meilleur précédent GeoRCG (95,3 % de validité, 97,20 % de stabilité) et CanonFlow (95,9 % de validité). Il atteint également les métriques d'énergie et de contrainte les plus basses, indiquant une plausibilité physique supérieure.
• Lissitude de la Représentation : L'analyse de la constante de Lipschitz révèle que l'espace latent affiné est significativement plus lisse que l'espace de l'encodeur brut. Sur GEOM-DRUG, la constante de Lipschitz empirique est réduite d'un facteur 4,6, confirmant que la variété de représentation est plus propice à la modélisation par diffusion.
• Richeesse de la Représentation : Le rang effectif de la représentation augmente de 65,32 (encodeur brut) à 88,42 (LENSEs), indiquant une réduction de l'effondrement dimensionnel et un espace latent plus riche et mieux structuré.
• Utilité en Aval : Lorsqu'il est utilisé pour des tâches de régression en aval sur les propriétés QM9 avec un encodeur et une tête de représentation figés, LENSEs surpasse l'encodeur brut sur toutes les propriétés, démontrant que le processus d'entraînement génératif lui-même agit comme un paradigme de préentraînement efficace pour affiner les représentations moléculaires.
• Efficacité d'Échantillonnage : LENSEs démontre une efficacité d'échantillonnage améliorée, atteignant une validité et une stabilité élevées avec moins d'étapes de génération par rapport à la référence.

Signification et Revendications
L'article revendique que la génération de haute qualité de molécules 3D dépend non seulement de l'architecture du générateur, mais de manière critique de la qualité de la représentation de conditionnement. LENSEs démontre que l'espace latent des encodeurs préentraînés peut être systématiquement amélioré pour être plus lisse, plus informatif sémantiquement et mieux aligné avec les objectifs génératifs.

Les auteurs posent que leur travail établit une relation bidirectionnelle entre la qualité de la génération et la qualité de la représentation. Plus précisément, ils suggèrent que l'entraînement génératif avec des objectifs d'alignement peut servir de paradigme de préentraînement potentiel pour les encodeurs moléculaires, affinant les représentations pour qu'elles soient plus prédictives et utiles pour les tâches en aval. Le travail met en évidence que les représentations moléculaires utiles ne sont pas intrinsèquement favorables à la génération et nécessitent un affinage explicite pour débloquer leur plein potentiel dans la modélisation générative.