Résumé technique : Yeti – Un tokeniseur de structure protéique compact pour la reconstruction et la génération multi-modale

1. Énoncé du problème

L'intégration des séquences protéiques, des structures 3D et des annotations fonctionnelles dans une représentation multi-modale unifiée est cruciale pour faire progresser la biologie computationnelle. Cependant, les approches actuelles font face à trois défis principaux :

1. Hétérogénéité des données : Les données biologiques couvrent des modalités diverses (contraintes biophysiques, texte, imagerie, données volumétriques) aux propriétés statistiques distinctes.
2. Incomplétude et bruit : Les jeux de données réels souffrent souvent de modalités manquantes, de mesures bruitées et d'une couverture inégale dans l'espace structure-fonction.
3. Extensibilité et couplage : La plupart des modèles opèrent à une seule échelle ou modalité, manquant de mécanismes pour transférer l'information entre les niveaux (de l'atomique au cellulaire) ou pour gérer des processus couplés.

Plus précisément, alors que les modèles génératifs récents traitent les protéines comme des objets 3D utilisant des espaces de coordonnées continues (par exemple, la diffusion SE(3)), il existe une alternative convaincante : encoder les structures atomiques 3D sous forme de tokens discrets. Cela permet un entraînement évolutif et agnostique à la modalité, similaire aux modèles de langage. Cependant, les tokeniseurs de structure protéique existants (par exemple, ESM3, DPLM-2, Kanzi) privilégient souvent la précision de reconstruction aux capacités génératives. Il existe un décalage où une haute fidélité de reconstruction ne garantit ni une diversité de tokens de haute qualité, ni un support pour la co-génération inconditionnelle de séquence et de structure. De plus, de nombreux modèles de l'état de l'art reposent sur des nombres massifs de paramètres ou sur des paradigmes d'entraînement en deux étapes où les encodeurs sont figés, limitant potentiellement l'optimisation de bout en bout.

2. Méthodologie

2.1 Architecture : Yeti

Yeti (Yielding Encoded Tokens for Intermodality) est un tokeniseur de structure protéique compact et en une seule étape, conçu pour mapper les coordonnées atomiques 3D continues vers des tokens discrets.

• Entrée : Structures protéiques représentées par des coordonnées $C_\alpha$ centrées sur la moyenne ($x \in \mathbb{R}^{L \times 3}$).
• Encodeur : Un encodeur basé sur Transformer utilisant une attention multi-têtes avec un encodage positionnel rotatif (RoPE) qui mappe les entrées vers des embeddings latents continus $Z \in \mathbb{R}^{L \times D}$.
• Quantification : Les embeddings sont projetés dans un espace de quantification et traités par un Quantificateur sans recherche (LFQ). L'espace LFQ est un produit cartésien de sous-réseaux binaires ($C_i = \{-1, +1\}$). La quantification est effectuée via une opération de signe, mappant les vecteurs vers des indices discrets sans nécessiter de table de recherche, générant une taille de vocabulaire $K = 2^D$.
• Décodeur : Une architecture StripedHyena (un modèle hybride de traitement du signal) agit comme décodeur. Contrairement aux décodeurs traditionnels qui reposent sur des composants invariants SE(3) comme l'attention par points invariants (IPA), le décodeur StripedHyena apprend un champ de vitesse conditionnel pour reconstruire les structures à partir de coordonnées bruitées.

2.2 Objectif d'entraînement : Flow Matching

Yeti est entraîné de bout en bout en utilisant un objectif de Flow Matching, qui interpole entre une distribution de bruit gaussien ($p_0$) et la distribution cible de structure protéique ($p_{data}$).

• Fonction de perte : Le modèle minimise la différence entre le champ de vitesse conditionnel prédit $v_\theta(x_t, t, c)$ et le champ vectoriel cible ($x_1 - x_0$), où $x_1$ est la structure propre et $x_0$ est un bruit gaussien.
• Régularisation par entropie : Pour éviter l'effondrement du codebook et assurer une forte utilisation, une pénalité d'entropie est ajoutée à la fonction de perte : $L = L_{FM} + \lambda L_{entropy}$.
• Inférence : La Guidance sans classificateur (CFG) est employée lors de l'inférence. En masquant aléatoirement les embeddings de tokens pendant l'entraînement, le modèle apprend à la fois les champs de vitesse conditionnels et inconditionnels, permettant un décodage guidé pour affiner la précision de la structure.

2.3 Co-génération multi-modale

Pour valider l'utilité générative des tokens Yeti, les auteurs ont entraîné un Modèle de Diffusion Masqué (MDM) à partir de zéro. Ce modèle génère conjointement des séquences d'acides aminés et des tokens de structure Yeti sans aucune initialisation pré-entraînée (contrairement à DPLM-2 qui démarre à partir d'un modèle de séquence, ou à ESM3 qui utilise des pipelines en cascade).

3. Contributions clés

1. Tokeniseur compact et de bout en bout : Introduction de Yeti, un tokeniseur en une seule étape entraîné de bout en bout avec flow matching, évitant le paradigme en deux étapes « figer l'encodeur » courant dans les travaux antérieurs.
2. Haute utilisation du codebook : Yeti atteint une utilisation supérieure du codebook et une diversité de tokens par rapport à des modèles plus grands (ESM3, DPLM-2) et à d'autres méthodes de quantification (FSQ, VQ-VAE), malgré un nombre de paramètres nettement inférieur.
3. Efficacité des paramètres : Yeti atteint une précision de reconstruction compétitive avec seulement 62,5 millions de paramètres, soit environ 10 fois moins que ESM3 (648 M) et 2 fois moins que DPLM-2 (118 M).
4. Preuve de concept de co-génération : Démonstration qu'un modèle multi-modal compact entraîné à partir de zéro sur des tokens Yeti peut effectuer une co-génération inconditionnelle de séquences et de structures protéiques, obtenant des résultats comparables à des modèles 10 fois plus grands.
5. Analyse de la dynamique de repliement : Identification d'une émergence hiérarchique de la structure secondaire au cours de la trajectoire de décodage par flow matching, caractérisée par une compaction globale précoce suivie d'un phénomène de « commitment tardif » où la topologie globale se consolide dans les étapes finales.

4. Résultats

4.1 Analyse du codebook

• Utilisation : Sur un codebook de 8 192 entrées, Yeti atteint une perplexité et une entropie élevées, surpassant nettement ESM3 et Kanzi. Sur le jeu de données CATH, Yeti utilise environ 66 % de sa capacité (Perplexité = 5414), avec une diversité intra-structure ($\bar{U}$) de 98,8 %, saturant presque la métrique et indiquant des affectations de tokens hautement distinctes pour les résidus.
• Comparaison : Yeti surpasse DPLM-2 en utilisation du codebook bien que les deux utilisent LFQ, ce qui est attribué au paradigme d'entraînement de bout en bout par opposition à l'approche d'encodeur figé de DPLM-2.

4.2 Précision de reconstruction

• Métriques : Yeti atteint un score TM de 0,96 sur CAMEO, 0,97 sur CASP14, et 0,95 sur CASP15 et CATH.
• Efficacité : Bien qu'ESM3 montre un RMSD légèrement inférieur (en raison de son échelle massive), les scores TM de Yeti sont comparables ou supérieurs, indiquant qu'il récupère efficacement la topologie globale correcte. Le décodeur basé sur le flow de Yeti produit des échantillons avec une variance dans les positions atomiques locales (pénalisant le RMSD) mais une topologie globale cohérente, une caractéristique bénéfique pour la génération multi-conformationnelle.

4.3 Co-génération inconditionnelle

• Performance : Un modèle de 224 millions de paramètres entraîné à partir de zéro sur 2,6 millions de protéines (utilisant des tokens Yeti) a atteint un score TM de cohérence interne (scTM) de 0,70 et un pLDDT de 76,12.
• Comparaison : Cette performance est comparable à DPLM-2 (scTM 0,87) et nettement supérieure à ESM3-Open (scTM 0,46), bien que DPLM-2 repose sur un modèle de séquence pré-entraîné de 650 millions de paramètres et que ESM3 utilise un pipeline en cascade.
• Designabilité : ProteinMPNN a récupéré 30,6 % des séquences générées originales à partir des structures générées, confirmant la cohérence géométrique et la designabilité des paires co-générées.

4.4 Dynamique de la trajectoire Flow

L'analyse de la trajectoire de décodage a révélé une émergence non monotone des caractéristiques :

• Phase précoce ($t < 0,6$) : Compaction globale rapide (diminution du rayon de giration).
• Phase tardive ($t > 0,9$) : Coalescence rapide des éléments de structure secondaire ($\alpha$-hélices et $\beta$-feuillets) et consolidation de la topologie globale.
  Cela suggère que le tokeniseur découple efficacement la densité grossière des détails structuraux haute fréquence.

5. Signification et revendications

L'article revendique que Yeti sert de fondation compacte et expressive pour l'entraînement de modèles protéiques multi-modaux. Sa signification principale réside dans la démonstration que :

1. La tokenisation discrète est viable pour la génération : Une co-génération inconditionnelle de haute qualité de séquence et de structure est possible sans pré-entraînement massif ni pipelines en cascade, à condition que le tokeniseur soit expressif et que l'objectif d'entraînement (flow matching) soit aligné sur les objectifs génératifs.
2. Efficacité plutôt que mise à l'échelle : Un modèle plus petit soigneusement conçu (62,5 millions de paramètres) peut surpasser ou égaler des modèles beaucoup plus grands (600 M+) en diversité de tokens et en reconstruction, remettant en question l'idée que l'échelle est le seul chemin vers la performance dans la modélisation des protéines.
3. Représentation unifiée : La capacité à entraîner un modèle à partir de zéro sur des données conjointes séquence-structure valide le potentiel d'une représentation moléculaire universelle qui capture le couplage intrinsèque entre les séquences d'acides aminés et les structures 3D.

Les auteurs positionnent Yeti non pas comme une solution finale pour toutes les tâches de conception de protéines, mais comme une preuve de concept établissant la viabilité de leur stratégie de tokenisation. Ils reconnaissent des limites, telles que la dépendance aux structures à conformation unique et la nécessité d'une mise à l'échelle supplémentaire pour égaler les performances de modèles spécialisés et fortement optimisés comme La-Proteina. Cependant, ils affirment que Yeti fournit une colonne vertébrale évolutive, efficace et performante pour les futurs modèles protéiques multi-modaux capables de co-générer des séquences et des structures hautement plausibles.