Sampling at intermediate temperatures is optimal for training large language models in protein structure prediction

Titel: Sampling bij intermediaire temperaturen is optimaal voor het trainen van grote taalmodellen in eiwitstructuurvoorspelling

Auteurs: L. Ghiringhelli, A. Zambon, G. Tiana
Instituut: Departement Natuurkunde, Universiteit van Milaan (Italië)

---

1. Het Probleem

Grote taalmodellen (LLM's), gebaseerd op de Transformer-architectuur, hebben de velden van machine learning en biofysica getransformeerd, met name door de mogelijkheid om de native conformatie van eiwitten uitsluitend op basis van hun aminozuursequentie te voorspellen (bijv. AlphaFold). Ondanks hun succes blijft het fundamentele mechanisme achter hun superieure prestaties, vergeleken met traditionele feedforward-netwerken, onvoldoende begrepen.

Terwijl er veel bekend is over het werkingsprincipe en de generalisatie-eigenschappen van feedforward-netwerken (zoals de vorm van het oplossingsruimte-manifold), ontbreekt er een diepgaand inzicht in hoe Transformers leren en waarom ze zo effectief zijn in het vastleggen van co-evolutie en ruimtelijke contacten in eiwitten.

2. Methodologie

De auteurs benaderen het probleem vanuit het perspectief van statistische mechanica. Ze behandelen het trainen van neurale netwerken als het verkennen van een energielandschap (de loss-functie) in een canoniek ensemble.

• Dataset: Een synthetische dataset van aminozuursequenties die zijn gegenereerd om te vouwen tot de native structuur van het acyl-coenzyme A binding protein (PDB-code: 2abd). De sequenties zijn gegenereerd via een energie-minimalisatieproces met een willekeurige interactiematrix.
• Taken: Masked Language Modeling (MLM), waarbij de taak is om gemaskeerde aminozuren in een sequentie te voorspellen.
• Architecturen:
  • Transformer (TF): Varianten met 1 blok (TF1) en 4 blokken (TF4).
  • Feedforward (FF): Een netwerk zonder attention-mechanisme, maar met een vergelijkbaar aantal parameters als TF4, gebruikt als controlegroep.
• Sampling-algoritme: In plaats van alleen te vertrouwen op gradient descent (Adam), gebruiken de auteurs een pseudo-Langevin dynamics-methode. Hierbij wordt de loss-functie behandeld als energie en worden de parameters gesampled bij verschillende temperaturen ($T$). Dit maakt het mogelijk om het "low-loss manifold" te karakteriseren zonder bias.
• Analyse: De auteurs analyseren de generalisatiefout, de variabiliteit van de parameters tussen verschillende gesamplede modellen, en de voorspellende waarde van de attention-matrix voor eiwitcontacten.

3. Belangrijkste Bijdragen en Resultaten

A. Optimaliteit van Intermediaire Temperaturen

• Feedforward-netwerken: Vertonen een scherke overgang (vergelijkbaar met een fase-overgang van de eerste orde) in de loss-functie bij een specifieke temperatuur. Hierdoor is er geen breed temperatuurbereik met goede generalisatie-eigenschappen; ze springen van een goed getrainde maar slecht generaliserende staat naar een slecht getrainde staat.
• Transformers: Vertonen een gladde overgang in de loss-functie. Dit resulteert in een breed bereik van intermediaire temperaturen waar de generalisatiefout minimaal is.
• Conclusie: De beste generalisatie wordt bereikt bij temperaturen die iets hoger liggen dan de temperatuur waar de trainingstfout minimaal is (ongeveer $1.1 - 1.5 \times T_{1/2}$, waarbij $T_{1/2}$ de temperatuur is bij de piek van de warmtecapaciteit). Dit verklaart waarom grote taalmodellen, die vaak "under-trained" blijven (omdat het minimaliseren van de loss tot het absolute minimum onpraktisch is), toch uitstekend presteren: ze opereren in dit gunstige intermediaire regime.

B. Behoud van Parameters en Embedding-dimensie

• Bij het samplingen bij de optimale temperatuur zijn de parameters van de meeste lagen in de Transformer hoogst geconserveerd (d.w.z. zeer vergelijkbaar tussen verschillende gesamplede modellen), mits de embedding-dimensie ($d$) optimaal is gekozen.
• De auteurs ontdekten dat bepaalde gewichten in normalisatielagen een bimodale verdeling vertonen: een piek rond nul (onbruikbare dimensies) en een piek rond een positieve waarde.
• Operatieve methode: Door de embedding-dimensie te verkleinen tot het aantal niet-nul waarden (bijv. van $d=64$ naar $d=26$ voor TF1 en $d=14$ voor TF4), behouden de modellen hun generalisatievermogen voor de sequentie-taak, maar worden de parameters in alle lagen nog sterker geconserveerd. Dit biedt een manier om de minimale, optimale dimensie van een model te bepalen.

C. Attention-matrix en Contactvoorspelling

• Een cruciale eigenschap van Transformers in de eiwitvoorspelling is dat de attention-matrix lijkt op de contactkaart (contact map) van het eiwit.
• Verrassende bevinding: De attention-matrix is voorspellender voor de echte contacten van het eiwit bij temperaturen die hoger liggen dan de temperatuur die optimaal is voor het leren van de sequentie.
• Dimensie-effect: Grotere embedding-dimensies leiden tot betere contactvoorspelling, zelfs als de generalisatie voor de sequentietaken vergelijkbaar blijft bij kleinere dimensies.
• Meting: Met een nieuwe metric ($P@L^\Delta_{NT}$) die rekening houdt met contacten die verder dan 7 residu's uit elkaar liggen, bleek dat de gesamplede modellen (bij intermediaire/hoge temperaturen) betere contactkaarten produceren dan modellen die puur via Adam zijn geoptimaliseerd (Early Stopped of Over Trained).

4. Significatie en Conclusie

Dit onderzoek biedt een fundamenteel inzicht in waarom Transformers zo effectief zijn voor biologische toepassingen:

1. Statistische Mechanica als Tool: Het gebruik van Langevin-dynamica om het parameter-ruimte te verkennen, onthult een breed "sweet spot" van temperaturen voor generalisatie dat niet zichtbaar is bij standaard training.
2. Onafhankelijkheid van Contactvoorspelling: De capaciteit om de native contacten van een eiwit te voorspellen (een emergente eigenschap) is niet perfect gecorreleerd met de capaciteit om de sequentie te voorspellen. Het optimaliseren voor de laagste loss (sequentietaak) is dus niet noodzakelijk de beste strategie om de structuurvoorspelling te maximaliseren.
3. Praktische Implicatie: De auteurs suggereren dat het gebruik van temperatuur-gestuurde sampling (Langevin) een efficiëntere optimalisatiestrategie kan zijn voor grote taalmodellen dan traditionele gradient descent, omdat het modellen genereert met superieure voorspellende eigenschappen voor zowel sequenties als structurele contacten.

Kortom, het paper stelt dat het "onder-traineren" van grote modellen in feite een voordeel is, omdat het hen in een intermediaire temperatuurregime houdt waar ze het beste generaliseren en de meest informatieve attention-matrices genereren voor structurele voorspelling.