Understanding protein function with a multimodal retrieval-augmented foundation model

Dit artikel introduceert PoET-2, een multimodaal, retrieval-versterkt foundation-model voor eiwitten dat door het integreren van evolutionaire beperkingen en structurele informatie state-of-the-art prestaties bereikt bij het voorspellen van varianteffecten en het leren van sequentiefunctierelaties.

De kern

De "Protein Poet": Een Slimme Vertaler voor het Leven

Stel je voor dat het leven een gigantische bibliotheek is. In deze bibliotheek staan niet boeken, maar eiwitten. Eiwitten zijn de kleine machines in ons lichaam die alles doen: ze bouwen spieren, verteren voedsel, vechten tegen virussen en sturen signalen door je hersenen.

Elk eiwit is geschreven in een heel speciaal alfabet van 20 letters (de aminozuren). Als je deze letters in de juiste volgorde zet, krijg je een werkend eiwit. Zet je één letter verkeerd (een mutatie), dan kan het hele eiwit kapot gaan. Dat is vaak de oorzaak van ziektes.

De onderzoekers van OpenProtein.AI hebben een nieuwe, slimme computer geïntroduceerd genaamd PoET-2. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "Grote Broer" die alles kent (Retrieval-Augmentation)

Stel je voor dat je een moeilijk raadsel moet oplossen, maar je bent niet slim genoeg om het alleen te doen. Wat doe je? Je vraagt het aan je familieleden die al eerder met dit soort raadsels hebben gewerkt.

De meeste oude AI-modellen voor eiwitten probeerden het allemaal zelf te raden door miljoenen voorbeelden uit hun hoofd te leren. Ze werden steeds groter en duurder, maar niet per se slimmer.

PoET-2 doet het anders. Het is als een slimme detective die tijdens het oplossen van een raadsel direct naar de "familie" van het eiwit kijkt.

• De analogie: Als je een nieuw eiwit wilt ontwerpen, vraagt PoET-2 niet alleen aan zijn eigen geheugen, maar hij haalt ook direct de beste voorbeelden van verwante eiwitten op uit een enorme database. Hij leert van de "context" (de familiegeschiedenis) in plaats van alles uit zijn hoofd te moeten hebben. Dit maakt hem veel efficiënter en slimmer zonder dat hij gigantisch groot hoeft te zijn.

2. De Twee-Oren-Methode (Multimodaal)

Eiwitten zijn niet alleen een rijtje letters; ze zijn ook 3D-gebouwen die in elkaar gevouwen zijn.

• Oude modellen keken vaak alleen naar de rijtje letters (de tekst).
• PoET-2 heeft twee zintuigen: hij leest de tekst én hij kijkt naar de 3D-vorm.

De analogie: Stel je voor dat je een origami-kraan wilt maken.

• Een oude AI zou alleen de instructies lezen: "Vouw papier, vouw papier..."
• PoET-2 leest de instructies, maar kijkt ook naar een foto van de gevouwen kraan. Als hij ziet dat een bepaalde vouw niet past bij de vorm, past hij de instructies direct aan. Hierdoor begrijpt hij veel beter hoe het eiwit echt werkt.

3. Twee Manieren van Leren (Dual Decoders)

PoET-2 heeft twee verschillende "hoofden" om te leren:

1. De Schrijver (Generatief): Dit deel kan nieuwe eiwitten "schrijven". Het denkt: "Als ik deze letters in deze volgorde zet, wordt het eiwit sterk." Dit is handig om nieuwe medicijnen te ontwerpen.
2. De Lezer (Begrijpend): Dit deel is een expert in het begrijpen van bestaande teksten. Het kan een eiwit bekijken en zeggen: "Ah, deze mutatie is gevaarlijk, die andere is onschuldig."

Waarom is dit een doorbraak?

Het probleem met de oude modellen:
Stel je voor dat je een eiwit hebt met één foutje. De oude modellen konden dat goed oplossen. Maar wat als je veel foutjes tegelijk hebt, of als er letters verdwijnen of bijkomen (zoals een zin waarin je woorden verwijdert of toevoegt)? De oude modellen raakten dan in paniek en gaven foute antwoorden.

De oplossing van PoET-2:
Omdat PoET-2 leert van de "familie" en de 3D-vorm, kan hij zelfs complexe mutaties voorspellen.

• Indels (Invoegingen/Verwijderingen): Hij kan een eiwit lezen waarbij stukken zijn weggeknipt of erbij zijn geplakt, en zeggen: "Dit zal waarschijnlijk nog steeds werken, of juist niet."
• Klinische toepassingen: Hij kan beter voorspellen of een mutatie bij een mens een ziekte veroorzaakt (pathogeen) of niet.

De Resultaten in het Kort

• Sneller en Kleiner: PoET-2 is niet de grootste AI (hij heeft "maar" 182 miljoen parameters), maar hij is slimmer dan de reuzen die er zijn. Hij is als een slimme student die beter presteert dan een professor die alles uit zijn hoofd moet leren.
• Minder Data nodig: Als je een nieuwe ziekte bestudeert en hebt maar weinig data, kan PoET-2 nog steeds goede voorspellingen doen. Hij is een meester in het "leren van weinig".
• De "Super-Combo": Als je PoET-2 combineert met een andere slimme methode (VenusREM), krijg je een "super-AI" die bijna perfect is in het voorspellen van hoe eiwitten zich gedragen.

Conclusie

PoET-2 is als een meester-vertaler die niet alleen de woorden (de aminozuren) kent, maar ook de cultuur (de evolutie) en de architectuur (de 3D-vorm) begrijpt. Hierdoor kan hij niet alleen vertalen, maar ook nieuwe, betere verhalen (eiwitten) schrijven die ons kunnen helpen bij het genezen van ziektes en het ontwerpen van nieuwe medicijnen.

Het is een enorme stap voorwaarts in het begrijpen van de taal van het leven.

Technische samenvatting

Probleemstelling

Proteïnen zijn essentieel voor bijna alle biologische processen, en het voorspellen van het effect van mutaties op hun functie is cruciaal voor ziektebegrip, medicijnontwikkeling en proteïne-engineering. Hoewel recente vooruitgang in Proteïne Taalmodellen (PLMs) de zero-shot voorspelling van mutatie-effecten heeft verbeterd, blijven er significante uitdagingen bestaan:

1. Beperkingen bij mutatietypes: De meeste bestaande PLMs zijn gebaseerd op Masked Language Modeling (MLM) en kunnen voornamelijk alleen enkelvoudige substituties voorspellen. Ze hebben moeite met inserties en deleties (indels) en kunnen epistatische effecten (interacties tussen meerdere mutaties) niet goed modelleren.
2. Data-efficiëntie: In supervised settings (waarbij beperkte experimentele data beschikbaar is) presteren bestaande modellen vaak niet optimaal, vooral bij het generaliseren naar posities die niet tijdens het trainen zijn gezien.
3. Schalingsproblemen: Het simpelweg vergroten van het aantal parameters in modellen leidt vaak tot betere structuurvoorspelling, maar niet noodzakelijk tot betere functievoorspelling of fitness-modellering, terwijl de rekentkosten exponentieel stijgen.
4. Integratie van modaliteiten: Bestaande modellen gebruiken vaak óf sequentie-informatie óf structurele informatie, maar zelden een geïntegreerde aanpak die beide modaliteiten en evolutionaire context (retrieval) combineert.

Methodologie: PoET-2

PoET-2 is een multimodaal, retrieval-versterkt foundation model dat is ontworpen om generatieve distributies over proteïnesequenties te leren, geconditioneerd op sequenties en structuren van homologen. De architectuur combineert drie kernideeën:

1. Multimodaliteit: Het model redeneert over zowel sequentie als structuur. Het kan conditioneren op sequentie- en/of structurele homologen, inclusief het genereren van sequenties op basis van gedeeltelijk waargenomen backbonestructuren (inverse folding).
2. Retrieval-Augmentatie: In plaats van een enorm model te trainen dat alle kennis in de parameters opslaat, gebruikt PoET-2 een context-conditioning framework. De gebruiker voert een "prompt" in die bestaat uit een set van evolutionair relevante proteïnen (context) en optioneel een deels gespecificeerd proteïne (query). Het model leert de evolutionaire beperkingen van deze familie "in-context".
   • De encoder gebruikt een hiërarchische attention-architectuur die volledig equivariant is voor de volgorde van proteïnen in de prompt. Dit elimineert de noodzaak voor miljarden parameters.
3. Dual Training Objectives: PoET-2 heeft twee decoders met verschillende doelen:
   • Autoregressieve Decoder (CLM): Trained met Causal Language Modeling. Deze is geoptimaliseerd voor generatie en het berekenen van log-likelihoods voor varianten (inclusief indels).
   • Bidirectionele Decoder (MLM): Trained met Masked Language Modeling. Deze produceert krachtige, context-bewuste embeddings voor representatieleren en supervisie.

Architectuurdetails:

• Inputs: Aminozuursequenties, lokale backbonestructuur (afstanden tussen N, Cα, C atomen) en voorspelde structuurvertrouwen (pLDDT).
• Attention Bias: Er wordt een structuur-gebaseerde attention-bias toegepast binnen individuele proteïnen, gebaseerd op de 3D-afstanden (Cα-Cα) in plaats van lineaire sequentieposities.
• Grootte: Het model is compact met slechts 182 miljoen parameters, wat het zeer efficiënt maakt voor inferentie.

Kernbijdragen

• Nieuwe Architectuur: Een hybride encoder-decoder Transformer die sequentie, structuur en evolutionaire context (via retrieval) combineert met equivariantie voor proteïnevolgorde.
• Indel-voorspelling: Voor het eerst kan een PLM op grote schaal indels en hogere-orde mutaties (meer dan één mutatie) nauwkeurig voorspellen zonder vast te zitten aan een vaste sequentielengte.
• Data-efficiëntie: De embeddings van PoET-2 tonen uitzonderlijke prestaties in few-shot learning, waarbij het met zeer weinig trainingsdata (bijv. <100 punten) prestaties bereikt die vergelijkbaar zijn met state-of-the-art modellen die duizenden data-punten nodig hebben.
• Openbare Beschikbaarheid: De auteurs hebben de code en modelgewichten beschikbaar gesteld via GitHub.

Resultaten

PoET-2 is geëvalueerd op de ProteinGym benchmark, die bestaat uit Deep Mutational Scanning (DMS) datasets en klinische datasets.

1. Zero-shot Variant Effect Prediction:

   • Indels: PoET-2 behaalt een state-of-the-art prestatie op indel-benchmarks, met een verbetering van ongeveer 20% ten opzichte van de vorige beste modellen (zoals VenusREM en S3F-MSA).
   • Hogere-orde mutaties: Het model presteert aanzienlijk beter bij mutaties met 3 of meer substituties, waar andere modellen vaak falen door hun aanname van additieve effecten.
   • Klinische varianten: PoET-2 verbetert de voorspelling van pathogeniciteit van menselijke mutaties (AUROC van 0.952 voor indels), wat een nieuw record is.
   • Ensembling: Een combinatie van PoET-2 en VenusREM overtreft alle individuele modellen op de DMS-substitutiebenchmark.

2. Supervised Learning (Few-shot):

   • Wanneer PoET-2 embeddings worden gebruikt als features voor een Gaussian Process (GP) regressiemodel, overtreft het de huidige state-of-the-art (Kermut) in alle cross-validatie scenario's (Random, Modulo, Contiguous).
   • Data-efficiëntie: In het "contiguous" split-scenario (waarbij testdata volledig nieuwe sequentiegebieden bevat) presteert PoET-2 met slechts 100 trainingspunten even goed als ESM C (een veel groter model) met ~2600 trainingspunten.

3. Invloed van Structuur:

   • Structuurconditionering verbetert de zero-shot voorspelling aanzienlijk, vooral voor stabiliteitsgerelateerde taken.
   • Voor supervised learning en klinische variantvoorspelling biedt expliciete structuurconditionering echter weinig tot geen extra voordeel, wat suggereert dat de embeddings van PoET-2 structurele informatie al impliciet bevatten.

Betekenis en Impact

PoET-2 markeert een verschuiving in de ontwikkeling van proteïne foundation modellen. Het bewijst dat het niet nodig is om modellen te schalen naar miljarden parameters om betere prestaties te behalen; in plaats daarvan is een slimme integratie van retrieval-augmentatie en multimodale conditionering effectiever.

• Efficiëntie: Met slechts 182M parameters is het model veel lichter en sneller te deployen dan concurrenten zoals ESM-3 of ProGen2 XL.
• Toepasbaarheid: De capaciteit om indels en complexe mutaties te modelleren opent nieuwe deuren voor proteïne-engineering, waar inserties en deleties vaak nodig zijn voor het ontwerpen van nieuwe functies.
• Generalisatie: De superieure prestaties in few-shot settings maken geavanceerde proteïne-engineering toegankelijker voor onderzoekers met beperkte experimentele datasets.

Kortom, PoET-2 biedt een robuust, efficiënt en veelzijdig framework dat de staat der kunst verbetert voor zowel het begrijpen van proteïne-evolutie als het ontwerpen van nieuwe proteïnen.