Technische Samenvatting: Yeti – Een Compacte Tokenizer voor Proteïne-structuur voor Reconstructie en Multi-Modale Generatie

1. Probleemstelling

De integratie van proteïne-sequenties, 3D-structuren en functionele annotaties in een uniforme multi-modale representatie is cruciaal voor de vooruitgang van de computationele biologie. Echter, huidige benaderingen staan voor drie hoofduitdagingen:

1. Data-heterogeniteit: Biologische data omvat diverse modaliteiten (biofysische beperkingen, tekst, beeldvorming, volumetrische data) met onderscheidende statistische eigenschappen.
2. Onvolledigheid en Ruis: Real-world datasets lijden vaak aan ontbrekende modaliteiten, ruis in metingen en ongelijke dekking over de structuur-functieruimte.
3. Schaalbaarheid en Koppeling: De meeste modellen opereren op één schaal of modaliteit, en missen mechanismen om informatie over te dragen tussen niveaus (van atomaar tot cellulair) of om gekoppelde processen te verwerken.

Specifiek, terwijl recente generatieve modellen proteïnes behandelen als 3D-objecten met behulp van continue coördinatenruimtes (bijv. SE(3) diffusie), is er een overtuigend alternatief: het coderen van 3D-atomaire structuren als discrete tokens. Dit maakt schaalbaar, modaliteits-onafhankelijk trainen mogelijk, vergelijkbaar met taalmodellen. Echter, bestaande tokenizers voor proteïne-structuren (bijv. ESM3, DPLM-2, Kanzi) geven vaak prioriteit aan reconstructienauwkeurigheid boven generatieve capaciteiten. Er is een misalignering waarbij hoge reconstructiefideliteit geen garantie biedt voor hoge kwaliteit van token-diversiteit of ondersteuning voor onvoorwaardelijke co-generatie van sequentie en structuur. Bovendien vertrouwen veel state-of-the-art modellen op enorme parameteraantallen of twee-staps trainingsparadigma's waarbij encoders worden bevroren, wat end-to-end optimalisatie potentieel beperkt.

2. Methodologie

2.1 Architectuur: Yeti

Yeti (Yielding Encoded Tokens for Intermodality) is een compacte, single-stage tokenizer voor proteïne-structuren die is ontworpen om continue 3D-atomaire coördinaten af te beelden op discrete tokens.

• Input: Proteïne-structuren weergegeven als gemiddeld gecentreerde $C_\alpha$-coördinaten ($x \in \mathbb{R}^{L \times 3}$).
• Encoder: Een Transformer-gebaseerde encoder die gebruikmaakt van multi-head attention met rotatoire positionele codering (RoPE) om inputs af te beelden op continue latente embeddings $Z \in \mathbb{R}^{L \times D}$.
• Quantisatie: De embeddings worden geprojecteerd in een quantisatieruimte en verwerkt door een Lookup-Free Quantizer (LFQ). De LFQ-ruimte is een Cartesisch product van binaire sub-codeboeken ($C_i = \{-1, +1\}$). Quantisatie wordt uitgevoerd via een tekenoperatie, die vectoren afbeeldt op discrete indices zonder de noodzaak van een lookup-tabel, wat resulteert in een vocabulairegrootte $K = 2^D$.
• Decoder: Een StripedHyena-architectuur (een hybride signaalverwerkingsmodel) fungeert als decoder. In tegenstelling tot traditionele decoders die vertrouwen op SE(3)-invariante componenten zoals Invariant Point Attention (IPA), leert de StripedHyena-decoder een conditioneel snelheidsveld om structuren te reconstrueren uit ruisige coördinaten.

2.2 Trainingsdoel: Flow Matching

Yeti wordt end-to-end getraind met behulp van een Flow Matching-doelstelling, die interpoleert tussen een Gaussische ruisverdeling ($p_0$) en de doelwitverdeling van proteïne-structuren ($p_{data}$).

• Functie voor Verlies: Het model minimaliseert het verschil tussen het voorspelde conditionele snelheidsveld $v_\theta(x_t, t, c)$ en het doelwit-vectorveld ($x_1 - x_0$), waarbij $x_1$ de schone structuur is en $x_0$ Gaussische ruis.
• Entropie-regularisatie: Om codeboek-collaps te voorkomen en hoog gebruik te garanderen, wordt een entropiestraf toegevoegd aan de verliesfunctie: $L = L_{FM} + \lambda L_{entropy}$.
• Inferentie: Classifier-Free Guidance (CFG) wordt ingezet tijdens inferentie. Door tijdens het trainen willekeurig token-embeddings op nul te zetten, leert het model zowel conditionele als onvoorwaardelijke snelheidsvelden, wat gestuurde decoding mogelijk maakt om de structuurnauwkeurigheid te verfijnen.

2.3 Multi-Modale Co-Generatie

Om het generatieve nut van Yeti-tokens te valideren, trainden de auteurs een Masked Diffusion Model (MDM) vanaf nul. Dit model genereert gezamenlijk aminozuur-sequenties en Yeti-structuurtokens zonder enige vooraf getrainde initialisatie (in tegenstelling tot DPLM-2, dat warm-start vanuit een sequentiemodel, of ESM3, dat gekaskadeerde pipelines gebruikt).

3. Belangrijkste Bijdragen

1. Compacte, End-to-End Tokenizer: Introductie van Yeti, een single-stage tokenizer die end-to-end wordt getraind met flow matching, waarmee het twee-staps "freeze encoder"-paradigma dat gebruikelijk is in eerdere werken wordt vermeden.
2. Hoog Codeboekgebruik: Yeti bereikt superieur codeboekgebruik en token-diversiteit in vergelijking met grotere modellen (ESM3, DPLM-2) en andere quantisatiemethoden (FSQ, VQ-VAE), ondanks aanzienlijk minder parameters.
3. Parameter-efficiëntie: Yeti bereikt concurrerende reconstructienauwkeurigheid met slechts 62,5 miljoen parameters, ongeveer 10x minder dan ESM3 (648M) en 2x minder dan DPLM-2 (118M).
4. Proof-of-Concept Co-Generatie: Demonstratie dat een compact multi-modaal model, vanaf nul getraind op Yeti-tokens, onvoorwaardelijke co-generatie van proteïne-sequenties en -structuren kan uitvoeren, met resultaten vergelijkbaar met modellen die 10x groter zijn.
5. Analyse van Vouwingsdynamiek: Identificatie van een hiërarchische emergentie van secundaire structuur tijdens de flow-matching decodeertrajectorie, gekenmerkt door vroege globale compactie gevolgd door een "late commitment"-fenomeen waarbij de globale topologie in de laatste stappen consolideert.

4. Resultaten

4.1 Codeboekanalyse

• Gebruik: Op een codeboek van 8.192 entries bereikt Yeti hoge perplexiteit en entropie, aanzienlijk beter dan ESM3 en Kanzi. Op de CATH-dataset gebruikt Yeti ongeveer 66% van zijn capaciteit (Perplexiteit = 5414), met een intra-structuur diversiteit ($\bar{U}$) van 98,8%, wat de metric bijna verzadigt en wijst op zeer onderscheidende token-toewijzingen voor residuen.
• Vergelijking: Yeti presteert beter dan DPLM-2 in codeboekgebruik, ondanks dat beide LFQ gebruiken, toegeschreven aan het end-to-end trainingsparadigma versus de bevroren encoder-aanpak van DPLM-2.

4.2 Reconstructienauwkeurigheid

• Metingen: Yeti bereikt een TM-score van 0,96 op CAMEO, 0,97 op CASP14 en 0,95 op CASP15 en CATH.
• Efficiëntie: Hoewel ESM3 een iets lagere RMSD toont (vanwege zijn enorme schaal), zijn de TM-scores van Yeti vergelijkbaar of superieur, wat aangeeft dat het effectief de correcte globale topologie herstelt. De flow-gebaseerde decoder van Yeti produceert samples met variatie in lokale atomaire posities (wat RMSD bestraft) maar consistente globale topologie, een eigenschap die gunstig is voor multi-conformatie-generatie.

4.3 Onvoorwaardelijke Co-Generatie

• Prestatie: Een model van 224 miljoen parameters, vanaf nul getraind op 2,6 miljoen proteïnes (met Yeti-tokens), bereikte een self-consistency TM-score (scTM) van 0,70 en een pLDDT van 76,12.
• Vergelijking: Deze prestatie is vergelijkbaar met DPLM-2 (scTM 0,87) en aanzienlijk beter dan ESM3-Open (scTM 0,46), ondanks dat DPLM-2 vertrouwt op een vooraf getraind sequentiemodel van 650 miljoen parameters en ESM3 een gekaskadeerde pipeline gebruikt.
• Ontwerpbaarheid: ProteinMPNN herstelde 30,6% van de oorspronkelijk gegenereerde sequenties uit de gegenereerde structuren, wat de geometrische consistentie en ontwerpbaarheid van de co-gegenereerde paren bevestigt.

4.4 Dynamiek van Flow-trajectoires

Analyse van het decodeertrajectoie onthulde een niet-monotoon ontstaan van kenmerken:

• Vroege Fase ($t < 0,6$): Snelle globale compactie (afnemende Traagheidsstraal).
• Late Fase ($t > 0,9$): Snelle coalescentie van secundaire structuurelementen ($\alpha$-helices en $\beta$-sheets) en consolidatie van de globale topologie.
  Dit suggereert dat de tokenizer effectief grofkorrelige dichtheid ontkoppelt van hoogfrequente structurele details.

5. Betekenis en Claims

Het artikel beweert dat Yeti dient als een compacte en expressieve basis voor het trainen van multi-modale proteïnemodellen. De primaire betekenis ligt in de demonstratie dat:

1. Discrete Tokenisatie Leefbaar is voor Generatie: Hoogwaardige, onvoorwaardelijke co-generatie van sequentie en structuur is mogelijk zonder enorme vooraftraining of gekaskadeerde pipelines, mits de tokenizer expressief is en de trainingsdoelstelling (flow matching) is afgestemd op generatieve doelen.
2. Efficiëntie boven Schaal: Een zorgvuldig ontworpen, kleiner model (62,5 miljoen parameters) kan grotere modellen (600M+) in token-diversiteit en reconstructie overtreffen of evenaren, wat de notie uitdaagt dat schaal de enige weg is naar prestaties in proteïnemodeling.
3. Unificatie van Representatie: Het vermogen om een model vanaf nul te trainen op gezamenlijke sequentie-structuurdata valideert het potentieel voor een universele moleculaire representatie die de intrinsieke koppeling tussen aminozuur-sequenties en 3D-structuren vastlegt.

De auteurs positioneren Yeti niet als een definitieve oplossing voor alle proteïne-ontwerptaken, maar als een proof-of-concept dat de levensvatbaarheid van hun tokenisatiestrategie vaststelt. Zij erkennen beperkingen, zoals de afhankelijkheid van structuren met één conformatie en de noodzaak van verdere schaalvergroting om de prestaties van gespecialiseerde, zwaar getune modellen zoals La-Proteina te evenaren. Echter, zij stellen dat Yeti een schaalbare, efficiënte en effectieve ruggengraat biedt voor toekomstige multi-modale proteïnemodellen die in staat zijn tot co-generatie van hoogst plausible sequenties en structuren.