Understanding protein function with a multimodal retrieval-augmented foundation model

Ce papier présente PoET-2, un modèle fondamental protéique multimodal et augmenté par la récupération qui, en combinant l'apprentissage contextuel des contraintes évolutives familiales avec une architecture de transformateur hiérarchique, atteint des performances de pointe dans la prédiction des effets de variants et l'apprentissage des relations séquence-fonction, notamment sur de petits jeux de données.

L'essentiel

Imaginez que les protéines sont comme des orchestres géants composés de 20 types d'instruments différents (les acides aminés). Pour qu'une protéine fonctionne bien (comme digérer un aliment ou combattre un virus), ses instruments doivent être dans le bon ordre et jouer la bonne partition. Si un musicien change de note (une mutation), l'orchestre peut jouer une musique magnifique, un bruit discordant, ou s'arrêter complètement.

Le problème ? Il existe un nombre astronomique de partitions possibles. Les scientifiques veulent prédire ce qui se passe si on change une note, mais c'est comme essayer de deviner le résultat d'une symphonie en changeant un seul instrument sans avoir entendu la musique complète.

Voici comment PoET-2 (le sujet de ce papier) change la donne, expliqué simplement :

1. Le Problème : Les anciens "musiciens" avaient la mémoire courte

Les modèles précédents (comme ESM ou ProtT5) étaient comme des étudiants brillants qui avaient lu des millions de partitions, mais qui avaient deux gros défauts :

• Ils avaient peur des changements de longueur : Si on ajoutait ou retirait un instrument (une mutation par insertion ou suppression), ils paniquaient. Ils ne savaient pas prédire le résultat.
• Ils étaient trop "gros" et coûteux : Pour devenir meilleurs, il fallait les rendre énormes (des milliards de paramètres), ce qui les rendait lents et chers à utiliser, sans toujours améliorer leur compréhension de la fonction réelle.

2. La Solution PoET-2 : Le "Chef d'orchestre" avec un assistant de recherche

PoET-2 est un nouveau modèle conçu par OpenProtein.AI. Imaginez-le non pas comme un étudiant qui a tout mémorisé par cœur, mais comme un chef d'orchestre génial qui a deux super-pouvoirs :

A. L'Assistant de Recherche (Rétro-ingénierie par la recherche)

Au lieu d'essayer de tout savoir par lui-même, PoET-2 a un assistant qui fouille instantanément dans une bibliothèque géante de partitions similaires (les protéines de la même "famille").

• L'analogie : Si vous demandez à PoET-2 de prédire l'effet d'une mutation sur une protéine, il ne devine pas. Il regarde : "Tiens, dans cette famille de protéines, quand on a changé cette note-là il y a 10 millions d'années, est-ce que l'orchestre a survécu ?"
• Cela lui permet d'apprendre les règles cachées de l'évolution sans avoir besoin d'être un monstre de taille. Il est petit (182 millions de paramètres) mais très efficace.

B. La Vision 3D (Multimodalité)

PoET-2 ne regarde pas seulement la partition (la séquence de lettres), il regarde aussi la forme de l'orchestre (la structure 3D).

• L'analogie : C'est comme si, pour comprendre si une note est bonne, le chef regardait non seulement la note écrite, mais aussi la position des musiciens sur scène. Si un musicien est trop loin des autres, la note sonnera faux, même si elle est juste sur le papier.
• PoET-2 peut dire : "Si on change cette note, est-ce que la forme de la protéine va s'effondrer ?"

3. Les Résultats Magiques

Grâce à cette approche, PoET-2 fait des choses que les autres ne pouvaient pas faire :

• Il gère les "trous" et les "ajouts" : Il peut prédire ce qui se passe si on insère ou supprime des instruments (les mutations d'insertion/délétion). C'est comme si le chef pouvait réarranger toute la partition si un musicien manquait ou si un nouveau rejoignait le groupe.
• Il comprend les effets en cascade : Parfois, changer une note affecte une autre note loin dans la partition (effets épistatiques). PoET-2 voit ces liens complexes là où les autres modèles voyaient seulement des notes isolées.
• Il apprend vite avec peu de données : Dans le monde réel, on n'a pas toujours des millions d'expériences. PoET-2 est si bon qu'il peut apprendre à prédire le comportement d'une protéine avec très peu d'exemples (comme un chef qui peut improviser une symphonie après avoir écouté seulement 3 répétitions).

En résumé

PoET-2 est comme un chef d'orchestre de génie qui, au lieu d'essayer de mémoriser chaque note de l'univers, utilise un assistant rapide pour consulter l'histoire de la musique (l'évolution) et regarde la scène en 3D pour comprendre comment les musiciens interagissent.

Il est plus petit, plus rapide et plus intelligent que ses prédécesseurs pour prédire comment les changements dans l'ADN (les mutations) affectent la santé humaine, la création de nouveaux médicaments ou la conception de protéines pour l'industrie. C'est une étape majeure pour comprendre le "langage de la vie".

Résumé technique

1. Problématique

Les modèles de langage protéique (PLM) actuels, bien qu'efficaces pour la prédiction de la structure, rencontrent des limites majeures dans la compréhension et la conception des protéines :

• Limites des approches existantes : La plupart des PLM reposent sur des modèles de langage masqué (MLM) qui peinent à prédire les effets des mutations complexes, notamment les insertions/délétions (indels) et les effets épistatiques (interactions entre mutations multiples). Ils sont souvent limités à des substitutions simples.
• Efficacité des données : Dans les campagnes d'ingénierie de protéines, les données expérimentales sont souvent rares (few-shot). Les modèles actuels nécessitent de grandes quantités de données pour généraliser correctement, en particulier pour des positions non vues lors de l'entraînement.
• Coût et généralisation : L'augmentation de la taille des modèles (scaling) améliore la prédiction de structure mais n'apporte pas de gains significatifs, voire négatifs, pour la modélisation de la fitness (aptitude fonctionnelle) et la prédiction de fonction, tout en augmentant considérablement les coûts de calcul.
• Manque d'intégration multimodale : Peu de modèles combinent efficacement l'apprentissage par récupération (retrieval-augmentation) et l'information structurelle pour apprendre des distributions génératives conditionnelles.

2. Méthodologie : PoET-2

PoET-2 est un modèle fondation protéique multimodal et augmenté par la récupération (retrieval-augmented), conçu pour apprendre des distributions génératives sur les séquences protéiques en conditionnant sur des homologues et des structures.

Architecture Principale

Le modèle utilise un encodeur-décodeur Transformer avec une architecture hiérarchique et deux décodeurs distincts :

• Encodeur Hiérarchique Équivariant :
  • Traite un ensemble de protéines (le "contexte") de manière équivariante à l'ordre des protéines. Cela signifie que la permutation des protéines dans le prompt n'affecte pas la sortie, éliminant le besoin de modèles massifs pour capturer les dépendances contextuelles.
  • Utilise une attention à deux étages : d'abord intra-protéine (avec biais d'attention basé sur la structure), puis inter-protéines.
  • Intègre les séquences, les structures (distances Cα, coordonnées du squelette) et la confiance structurelle (pLDDT).
• Double Décodeur :
  1. Décodeur Auto-régressif (CLM) : Entraîné avec un objectif de modélisation de langage causal. Il est utilisé pour la génération de séquences et le calcul de vraisemblance (log-likelihood) pour la prédiction "zero-shot" des effets de mutations (y compris les indels).
  2. Décodeur Bidirectionnel (MLM) : Entraîné avec un objectif de langage masqué. Il produit des embeddings riches et contextuels pour l'apprentissage supervisé et la prédiction de fonction.

Approche Multimodale et Récupération

• Prompting : L'utilisateur fournit un "contexte" (un ensemble d'homologues) et une "requête" (une séquence ou une structure partielle). Le modèle infère les contraintes évolutives du contexte pour générer ou évaluer des variantes.
• Conditionnement Structurel : Le modèle peut être conditionné par des structures atomiques partielles (inverse folding) ou par les structures prédites des homologues dans le contexte.
• Taille du modèle : PoET-2 est compact avec 182 millions de paramètres, ce qui le rend efficace en inférence par rapport aux modèles de plusieurs milliards de paramètres.

3. Contributions Clés

1. Prédiction Zero-Shot Avancée : PoET-2 est le premier modèle à prédire efficacement les effets des indels et des mutations d'ordre supérieur (plusieurs mutations simultanées) en zero-shot, surpassant les méthodes précédentes de plus de 20% sur certains benchmarks.
2. Efficacité des Données (Few-Shot) : Dans les réglages supervisés, les embeddings de PoET-2 permettent d'apprendre des relations séquence-fonction avec des données très limitées, surpassant les méthodes de l'état de l'art (comme Kermut) même avec des ensembles de données beaucoup plus petits.
3. Architecture Équivariante et Multimodale : L'introduction d'une attention hiérarchique équivariante permet d'utiliser le contexte évolutif sans nécessiter des modèles massifs, tout en intégrant nativement la structure 3D.
4. Objectifs d'Entraînement Dual : La combinaison de la génération (CLM) et de la représentation (MLM) dans un même modèle permet une flexibilité unique pour les tâches de conception et d'analyse.

4. Résultats Expérimentaux

Les évaluations ont été réalisées sur le benchmark ProteinGym, couvrant des données de Deep Mutational Scanning (DMS) et des variants cliniques.

• Prédiction Zero-Shot (DMS) :
  • Indels : PoET-2 établit un nouvel état de l'art avec un coefficient de corrélation de Spearman ($\rho$) de 0,566, surpassant PoET-1 (0,517) et VenusREM (0,414).
  • Mutations d'ordre supérieur : Pour les variants avec 3 à 5+ mutations, PoET-2 surpasse significativement les modèles existants (ex: $\Delta\rho \approx 0,09$ par rapport à VenusREM pour 3 mutations).
  • Substitutions : Les performances sont comparables à l'état de l'actuel (VenusREM), avec une légère infériorité sur la corrélation globale mais une supériorité sur la détection des mutations bénéfiques (NDCG).
• Prédiction Clinique :
  • Amélioration significative de la prédiction de la pathogénicité des variants humains (AUROC de 0,952 pour les indels et 0,928 pour les substitutions), battant les records précédents.
• Apprentissage Supervisé (Few-Shot) :
  • Sur le benchmark DMS supervisé, PoET-2 GP (Gaussian Process utilisant les embeddings PoET-2) atteint un $\rho$ moyen de 0,693, surpassant Kermut (0,664).
  • Efficacité des données : Avec seulement 100 points de données d'entraînement, PoET-2 atteint les performances de modèles concurrents entraînés sur des milliers de points.
• Impact de la Structure : L'ajout de la structure dans le prompt améliore la prédiction zero-shot (surtout pour la stabilité), mais n'apporte pas de gain significatif en apprentissage supervisé, suggérant que les embeddings capturent déjà l'information structurelle de manière implicite.

5. Signification et Impact

PoET-2 représente une avancée majeure dans le domaine de l'intelligence artificielle appliquée à la biologie :

• Démocratisation de l'ingénierie de protéines : Avec seulement 182M de paramètres, le modèle est accessible et rapide à déployer, contrairement aux modèles géants.
• Résolution de problèmes complexes : La capacité à gérer les indels et les mutations multiples en zero-shot ouvre de nouvelles voies pour la conception de protéines et la compréhension des maladies génétiques complexes.
• Nouveau paradigme d'apprentissage : L'approche combinant récupération de contexte (retrieval) et multimodalité démontre que l'intégration intelligente de l'information évolutive et structurelle est plus efficace que la simple augmentation de la taille des paramètres.

En conclusion, PoET-2 offre un outil robuste et efficace pour la prédiction des effets de variants et la conception de protéines, comblant le fossé entre la modélisation générative et la compréhension fonctionnelle dans des régimes de données limités.