Technische Samenvatting: Few-step Cofolding met All-Atom Flow Maps (DECAF)

Probleemstelling

De nauwkeurige computationele modellering van 3D-biomoleculaire complexen, met name voor eiwit-ligand cofolding, is een dominant paradigma geworden binnen de structuurgebaseerde medicijnontwerp (SBDD). Hoewel moderne all-atom generatieve modellen (bijv. AlphaFold 3, Boltz-1) structuren kunnen voorspellen met atomaire resolutie, vertrouwen zij op dure iteratieve diffusie-rollouts. Het genereren van hoogwaardige structuren vereist doorgaans honderden neural function evaluations (NFEs), wat een significante computationele bottleneck vormt voor zowel implementatie als inference-time zoektechnieken.

Bovendien is het genereren van fysiek valide structuren op atomair niveau cruciaal; onnauwkeurige lokale modellering leidt tot catastrofale fouten zoals sterische botsingen, onjuiste bindingslengtes/hoeken en gespannen zijketens. Strikte fysieke beperkingen verschuiven vaak de computationele last naar inference-time methoden, waarbij reward-gestuurde zoekopdrachten (bijv. Sequential Monte Carlo, MCTS) worden gebruikt om samples te verfijnen. Deze zoekmethoden verhogen echter de inference-kosten, omdat ze meerdere reward-queries vereisen over meerdere deeltjes tijdens het traject. De centrale onderzoeksvraag is of het mogelijk is om een all-atom cofolding-model te trainen dat efficiënt reward-geoptimaliseerde samples genereert met slechts een klein aantal NFEs tijdens inference.

Methodologie: Het DECAF-framework

De auteurs introduceren DENOISER COFOLDING ALL-ATOM FLOWMAP (DECAF), een framework ontworpen om state-of-the-art all-atom cofolding diffusiemodellen te distilleren naar few-step flow maps. De methodologie adresseert specifieke technische uitdagingen bij het aanpassen van standaard diffusiemodellen naar flow maps voor all-atom eiwitmodellering.

1. $\sigma$-Ruimte Flow Map Reparametrisatie

Standaard diffusiemodellen gebruiken vaak de EDM (Elucidating the Design Space) parametrisatie, die conditioneert op een enkele ruisniveau $\sigma$. Deze formulering laat zich niet eenvoudig uitbreiden naar dual-time flow maps, en de niet-lineaire tijd-reparametrisatie $\sigma(t)$ die in EDM wordt gebruikt, leidt tot numerieke instabiliteit bij het berekenen van de tijd-partiële afgeleiden ($\partial \sigma / \partial t$) die vereist zijn voor flow map objectieven.

Om dit te omzeilen, definieert DECAF de flow map direct in de $\sigma$-ruisruimte. In plaats van de tijd $t$ te mappen, mapt het tussen ruisniveaus $\rho \to \sigma$ (waarbij $\rho > \sigma$).

• $\sigma$-Velocity: Het framework definieert een instantane $\sigma$-velocity $v_\sigma(x_\sigma) = dx_\sigma/d\sigma$.
• Gemiddelde Velocity: De flow map wordt geparametriseerd als een verplaatsing met behulp van de gemiddelde $\sigma$-velocity tussen twee ruisniveaus, waardoor het model directe "sprongen" kan maken langs de probability flow ODE-trajectorie in de ruisruimte.
• Stabiliteit: Deze herparametrisatie elimineert de problematische $\partial \sigma / \partial t$ factor, wat stabiele training en directe distillatie van vooraf getrainde EDM-stijl cofolding-modellen mogelijk maakt.

2. Denoiser Parametrisatie met SE(3) Uitlijning

Een cruciale innovatie is het gebruik van een denoiser-gebaseerde formulering in plaats van een directe velocity-voorspeller.

• Endpoint Voorspelling: Het model voorspelt de doel-endpoint $\hat{x}_{tgt}$ (de schone structuur) in plaats van de instantane velocity.
• SE(3) Rigide Uitlijning: De verliesfunctie incorporeert een minimalisatie over de SE(3)-groep (met behulp van het Kabsch-algoritme) om de voorspelde structuur $\hat{x}_{tgt}$ uit te lijnen met de grondwaarheid (of de voorspelling van de teacher). Dit dwingt zachtjes SE(3)-symmetrie af, wat cruciaal is voor het trainen van stabiele modellen en het verminderen van variantie.
• Verliesfunctie: Het trainingsdoel komt de voorspelde endpoint overeen met de voorspelling van de teacher na rigide uitlijning, wat zowel diagonale (score matching) als off-diagonale (jump) training ondersteunt.

3. DECAF-SEARCH: Inference-Time Framework

Om de flow map te benutten voor reward-optimalisatie, introduceren de auteurs DECAF-SEARCH.

• Lookahead Mechanisme: In tegenstelling tot standaard diffusiemodellen, waarbij reward-estimatie een volledige denoising vereist, biedt de DECAF flow map een "lookahead" voorspelling van de schone endpoint $\hat{x}_0$ vanuit elke tussenliggende ruisige staat $x_\sigma$. Dit maakt hogere-getrouwe reward-estimatie ($R(\hat{x}_0)$) mogelijk zonder het volledige traject te hoeven voltooien.
• Zoekalgoritme: Het framework past inference-time sturingsmethoden (Feynman-Kac, MCTS) aan voor de flow map setting. Het gebruikt een stochastische $\gamma$-sampler om deeltjes te genereren, evalueert hun lookahead rewards, en voert clean-space gradiënt-ascent of resampling uit om de generatie te sturen naar fysiek valide, hoog-reward configuraties.

Belangrijkste Bijdragen

1. Eerste All-Atom Eiwit Flow Map: Het ontwerp van DECAF-Boltz, de eerste all-atom eiwit flow map die er succesvol in slaagt een vooraf getrainde Boltz-1 teacher te distilleren.
2. Methodologische Innovaties:
   • Een nieuwe $\sigma$-ruimte herparametrisatie die stabiele training van flow maps voor all-atom modellering mogelijk maakt.
   • Een denoiser-gebaseerde parametrisatie die SE(3) gewogen rigide uitlijning mogelijk maakt, cruciaal voor stabiliteit en nauwkeurigheid.
3. Inference-Time Zoeken: De introductie van DECAF-SEARCH, een framework dat gebruikmaakt van de flow map lookahead om reward-estimatie en uitlijning te verbeteren, wat efficiënt zoeken met minder NFEs mogelijk maakt.
4. Empirische Validatie: Demonstratie dat DECAF beter presteert dan diffusie-gebaseerde baselines in lage-NFE regimes en de performance van een full-budget evenaart met aanzienlijk minder berekeningen.

Experimentele Resultaten

De auteurs evalueerden DECAF op de Runs N' Poses (RnP) en PoseBusters benchmarks, waarbij zij vergeleken met Boltz-1x (een variant van Boltz-1 met een fysiek potentiaal) en andere frontier-modellen zoals AlphaFold 3 en Pearl.

Prestaties bij Lage NFE Budgets

• RnP Benchmark: Bij strikte NFE-budgetten (bijv. 20–50 NFEs), presteert DECAF-SEARCH significant beter dan Boltz-1x (getuned) over alle metrieken (PB-Valid, Success Rate, lDDT-PLI). Bijvoorbeeld, bij 20 NFEs behaalt DECAF een PB-Valid score van 83,5% vergeleken met 75,9% voor Boltz-1x.
• Generalisatie: De winst in prestaties houdt stand over alle kwartielen van pocket-gelijkenis, inclusclusief out-of-distribution targets, wat de robuustheid van het model bevestigt.
• Efficiëntie: DECAF evenaart de prestaties van Boltz-1x bij 600 NFEs met slechts 30 NFEs (een 20x reductie in compute) op PoseBusters.

Distillatie van State-of-the-Art Modellen

• DECAF-Pearl: Het framework werd toegepast om het state-of-the-art Pearl cofolding-model te distilleren. DECAF-Pearl evenaart de success rate van de teacher terwijl het 5x minder NFEs gebruikt (40 NFEs vs. volledige simulatie).
• Vergelijking: DECAF-Pearl presteert beter dan andere krachtige baselines (inclusief volledige simulatie van Pearl) in termen van fysieke validiteit en success rate bij lage NFE-budgetten.

Pareto Frontiere

Analyse van de compute-optimale frontiere laat zien dat DECAF-SEARCH de Boltz-1x domineert over alle inference compute-budgetten. De optimale zoekstrategie varieert per budget:

• Lage NFE ($\le$ 30): Particle-based SMC (FK-steering) is optimaal.
• Gemiddelde NFE (100–250): Monte Carlo gradiënt-estimatie (MC-GRAD) is het meest effectief.
• Hoge NFE ($\ge$ 142): MCTS-gebaseerd zoeken levert de beste resultaten op.

Betekenis en Claims

Het artikel claimt dat DECAF een belangrijke stap voorwaarts is om all-atom biomoleculaire generatie computationeel haalbaar te maken voor high-throughput toepassingen.

• Versnelling: Door diffusietrajecten te comprimeren naar few-step flow maps, reduceert DECAF de inference latency met ordes van grootte (bijv. 5x tot 20x reducties in NFEs) zonder in te boeten op structurele accuratesse of fysieke validiteit.
• Reward Optimalisatie: De flow map lookahead maakt efficiëntere en nauwkeurigere reward-gestuurde zoekopdrachten mogelijk, wat de "inference tax" aanpakt die geassocieerd wordt met meerdere reward-queries in standaard diffusiemodellen.
• Praktische Impact: De versnelling maakt nieuwe workflows mogelijk, zoals het cofolden van volledige ligand-bibliotheken tegen een target voor virtuele screening en het genereren van grootschalige synthetische datasets voor het trainen van downstream modellen, wat voorheen onpraktisch was vanwege de kosten van volledige diffusie-simulaties.
• Bescheidenheid: De auteurs erkennen dat fouten (bijv. niet-planare sp2-bindingen) worden overgenomen van de reward-signalen van de teacher-modellen (Boltz-1x) in plaats van van de zoekprocedure zelf. Ze merken ook op dat de relatie tussen het NFE-budget en de success rate niet-monotoon is, wat wijst op toekomstig werk in adaptief zoeken en gezamenlijke training met taakspecifieke rewards.

De code is vrijgegeven op https://github.com/genesistherapeutics/decaf.