Technische Samenvatting: ENSEMBITS

Probleemstelling

Proteïne-taalmodellen (PLMs) en structurele voorspellingstools hebben grotendeels vertrouwd op statische structurele representaties of eenvoudige aminozuursequenties. Hoewel bestaande proteïne-structuretokenizers (PST's) effectief de lokale geometrie van statische momentopnamen vastleggen, slagen ze er niet in om de gecorreleerde bewegingen en alternatieve conformationele toestanden in te coderen die inherent zijn aan proteïne-ensembles. Biologische functie wordt vaak niet gecodeerd in een enkele structuur, maar in de verdeling van conformaties die een proteïne bemonsterd (bijvoorbeeld katalytische lussen die de toegang van substraten regelen). Het vakgebied verschuift naar ensemble-generatie, maar er ontbreekt een discrete vocabulaire die deze dynamische ensembles kan tokeniseren voor gebruik in transformer-gebaseerde architecturen. Bovendien zijn dynamische data vaak schaars, en bestaande methoden slagen er ofwel niet in om gecorreleerde beweging over residuen gezamenlijk te modelleren, of ze vereisen volledige trajecten die bij inferentie mogelijk niet beschikbaar zijn.

Methodologie: ENSEMBITS

ENSEMBITS wordt geïntroduceerd als de eerste tokenizer die specifiek is ontworpen voor proteïne-conformationele ensembles. Het beeldt een ongeordende multiset van structurele frames (conformeren) af op een discrete token via een Residual Vector-Quantized Variational Autoencoder (RVQ-VAE) architectuur.

1. Probleemformulering

Het model behandelt ensemble-tokenisatie als een multiset-naar-token compressieprobleem.

• Invoer: Een ongeordende multiset $P = \{x_1, \dots, x_P\}$ van $P$ frames voor een enkel proteïne, waarbij elk frame per-atom coördinaten bevat.
• Uitvoer: Een discrete token uit een codeboek $\mathcal{C}$ die het lokale dynamische motief van een residu over het ensemble vertegenwoordigt.
• Beperking: De tokenizer moet $P$-agnostisch zijn, ensembles van elke positieve kardinaliteit accepteren en permutatie-invariant zijn met betrekking tot de volgorde van de frames.

2. Ontwerp van Descriptoren

Om dynamiek vast te leggen, berekent ENSEMBITS SE(3)-invariante descriptoren over frames. De auteurs maken gebruik van een ESM3-stijl relatief-frame-descriptor (de productsetting) die de lokale ruggegraatomgeving codeert als SE(3)-transformaties tussen het frame van een residu en zijn $K$ ruimtelijk dichtstbijzijnde buren.

• Dynamische Modus: Cruciaal zijn de identiteiten van de buren ($N^p_r$) onafhankelijk voor elke frame $p$ opnieuw berekend. Dit stelt de descriptor in staat om op natuurlijke wijze contactvorming en contactbreuk langs het traject te coderen, een signaal dat onzichtbaar is voor vaste-benaderingsmethoden.

3. Architectuur

• Set-encoder: Een PerceiverIO-stijl architectuur die de multiset van per-residu-descriptoren afbeeldt op een enkele latente vector $z$. Het gebruikt een gedeelde per-element MLP gevolgd door cross-attention met leerbare query-tokens, waardoor permutatie-invariantie per constructie wordt gewaarborgd.
• Residuale Quantizer: Een multi-stadia RVQ-VAE met $K=3$ niveaus discretiseert de latente $z$ in een tuple van codeboek-indexen $(c_1, \dots, c_K)$. De uiteindelijke token is de som van de embeddingen van deze niveaus.
• Decoder & Verlies: De decoder reconstrueert de multiset van descriptoren. Het reconstructieverlies maakt gebruik van Hongaarse matching om voorspelde en doel-descriptoren uit te lijnen, waardoor permutatie-invariantie wordt gewaarborgd zonder een vaste volgorde te vereisen.

4. Single-Frame-naar-Token Distillatie (SFTD)

Om de schaarste aan dynamische data aan te pakken en inferentie op enkele statische structuren mogelijk te maken, stellen de auteurs SFTD voor.

• Mechanisme: Tijdens het trainen verwerkt het model twee ensembles van hetzelfde proteïne: een volledig ensemble ($P_{max}$) en een willekeurig bemonsterd sub-ensemble (tot $P=1$).
• Doel: De latente representatie van het sub-ensemble wordt naar de latent van het volledige ensemble getrokken (met een stop-gradient op het volledige ensemble).
• Resultaat: Dit stelt de tokenizer in staat om bij het testen een "dynamische token" te voorspellen vanuit een enkele voorspelde structuur, waardoor op ensemble-niveau informatie effectief wordt gedistilleerd naar een single-frame representatie.

Belangrijkste Bijdragen

1. Formulering en Pipeline: De eerste end-to-end formulering van proteïne-dynamiek-tokenisatie als een multiset-naar-token compressieprobleem, met gebruik van SE(3)-invariante descriptoren en permutatie-invariante set-encoders.
2. SFTD-doel: Een nieuwe trainingsstrategie die embeddingen van sub-ensembles uitlijnt met die van volledige ensembles, waardoor het model kan fungeren als een dynamische tokenizer zelfs wanneer slechts één statische structuur beschikbaar is.
3. Empirische Validatie: Uitgebreide validatie die aantoont dat ENSEMBITS ensemble-niveau dynamiek vastlegt en effectief overdraagt naar downstream functionele taken, vaak prestaties van statische tokenizers overtreffend ondanks het gebruik van aanzienlijk minder pretraining-data.

Resultaten

1. Dynamische Representatie (RMSF & ANOVA)

• RMSF-voorspelling: ENSEMBITS domineert benchmarks voor de voorspelling van Root Mean Square Fluctuation (RMSF).
  • Multi-frame: ENSEMBITS ($P=full$) presteert beter dan de sterkste baseline (ProtProfileMD) met ongeveer 2,5 Spearman-punten op mdCATH-div en ongeveer 11–18 punten op MISATO-splits.
  • Single-frame: Zelfs met $P=1$ (gedistilleerd) presteert ENSEMBITS beter dan alle andere single-frame tokenizers (inclusief AminoAseed en ESM3struct) bij RMSF-voorspelling.
• ANOVA-test: Een eenweg-ANOVA op bewegingsamplitude ($s_1$) toont aan dat ENSEMBITS-tokens 37,1% van de variantie in lokale bewegingsamplitude verklaren ($\eta^2 = 0,371$). Dit is ongeveer 3x hoger dan de volgende beste multi-frame baseline (Vote_3Di, $\eta^2 \approx 0,128$) en 7x hoger dan single-frame statische tokenizers. Dit bevestigt dat tokens onderscheidende lokale dynamiek coderen en niet alleen statische vouwklassen of posities.

2. Downstream Functionele Taken

Ondanks het gebruik van veel minder pretraining-data (alleen MD-ensembles, geen PDB-schaal inverse folding pretraining), komt ENSEMBITS overeen met of overtreft statische tokenizers op verschillende taken:

• Voorspelling van Mutatie-effecten: Op de PROTEINGYM-benchmark presteert ENSEMBITS ($P=1$) in combinatie met ESM2 beter dan alle andere structurele tokenizers die zijn gecombineerd met ESM2, met een winst van 6,9% ten opzichte van ESM2 alleen.
• Bindingplaats & Affiniteit: ENSEMBITS leidt bij de voorspelling van bindingplaatsen (AUROC 0,750) en regressie van bindingsaffiniteit over alle splits. Opmerkelijk is dat het ESM3struct en AminoAseed overtreft bij het detecteren van bindingplaatsen, wat suggereert dat lokale conformationele flexibiliteit een kritiek signaal is voor interface-residuen dat statische tokenizers missen.
• EC & GO-voorspelling: ENSEMBITS is concurrerend met ESM3struct bij Enzyme Commission (EC) en Gene Ontology (GO) taken, hoewel ESM3struct een voorsprong behoudt bij taken die sterk gecorreleerd zijn met vouwklassen (waarschijnlijk vanwege zijn inverse-folding pretraining).

Betekenis en Claims

Het artikel positioneert ENSEMBITS als een noodzakelijke stap in de evolutie van proteïne-modellering van statische structurele voorspelling naar ensemble-generatie.

• Overbrugging van de Kloof: Het biedt de discrete vocabulaire die nodig is om proteïne-dynamiek te brengen in taalmodellering en ontwerp.
• Data-efficiëntie: Het toont aan dat hoogwaardige dynamische representaties kunnen worden geleerd uit relatief kleine MD-corpora (in vergelijking met de enorme sequentie/structuur-corpora die worden gebruikt door modellen zoals ESM3), mits de representatie expliciet is ontworpen voor dynamiek.
• Praktisch Nut: Het SFTD-mechanisme verlicht de bottleneck van de beschikbaarheid van dynamische data, waardoor het model kan worden bevraagd op enkele statische structuren (bijvoorbeeld AlphaFold-voorspellingen) terwijl het nog steeds gebruikmaakt van de geleerde dynamische priors.

De auteurs merken bescheiden op dat de kwaliteit van de tokens wordt begrensd door de dekking en fideliteit van de onderliggende trainings-trajecten en dat het doel is om een methodologie te introduceren voor het inbrengen van dynamische informatie in discrete tokens, in plaats van te claimen dat ze alle proteïne-dynamiek hebben opgelost. Toekomstig werk omvat schaling naar grotere ensemble-corpora en het direct integreren van deze tokens in proteïne-taalmodellen voor multi-frame structurele generatie.