Technische Samenvatting: ABOPD – Antibody CDR Design via On-Policy Distillation

1. Probleemstelling

Het computationeel ontwerpen van antilichaam Complementarity-Determining Regions (CDR's) is een kritieke uitdaging in de eiwitengineering, aangezien deze flexibele lussen bepalen hoe een antigeen wordt herkend. Hoewel recente diffusie-gebaseerde generatieve modellen het veld hebben geavanceerd, bestaat er een fundamentele mismatch tussen het trainingsdoel en het inferentieproces:

• Trainingsdistributie: Standaard denoising-training ziet toe op "reference-derived states", wat ruisige structuren zijn die zijn gegenereerd door natuurlijke (ground truth) CDR's te verstoren.
• Inferentiedistributie: Tijdens generatie produceert het model recursief toestanden op basis van zijn eigen eerdere voorspellingen.
• De Mismatch: Naarmate het generatietraject vordert, accumuleren voorspellingsfouten, waardoor het model toestanden bezoekt die significant afwijken van de reference-derived distributie die tijdens de training werd gebruikt. Dit is bijzonder schadelijk voor flexibele CDR's zoals CDR-H3, waarbij kleine afwijkingen in de backbone kunnen voortplanten, wat leidt tot sterische botsingen en onjuiste geometrieën gericht op het antigeen.

Bestaande post-training strategieën, zoals Supervised Fine-Tuning (SFT) of offline distillatie, blijven opereren op de reference-derived distributie, waardoor ze er niet in slagen fouten te corrigeren die specifiek ontstaan langs het eigen generatietraject van het model.

2. Methodologie: ABOPD Framework

De auteurs stellen ABOPD voor (Antibody design framework built on On-Policy Distillation), een post-training strategie die is ontworpen om de structurele supervisie af te stemmen op de toestanden die daadwerkelijk worden aangetroffen tijdens recursieve generatie. Het framework bestaat uit drie hoofdfasen:

A. Hybride Pretraining (H-DiffAb)

Om een robuuste initialisatie te vestigen, pretrainen de auteurs een enkel model, H-DiffAb, over diverse masking-scenario's: single-CDR, multi-CDR en alle zes de CDR-target masks. Dit model dient als de initialisatie voor zowel de student als de teacher, wat een eerlijke vergelijking garandeert.

B. Backbone-Aware Teacher Adaptation

Een "Teacher"-model wordt afgeleid van H-DiffAb door het integreren van privileged information die alleen beschikbaar is tijdens de training: de natuurlijke geometrie van de doelresiduen.

• Architectuur: De teacher gebruikt lichtgewicht residuele MLP-adapters om natuurlijke $C_\alpha$-coördinaten, atoomvaliditeitsmasks (N, $C_\alpha$, C, O, $C_\beta$) en target-to-complex afstandskenmerken te coderen.
• Functie: Deze geprivilegieerde beschrijvers conditioneren de denoising-voorspelling zonder de diffusie-toestandrepresentatie te wijzigen. De teacher wordt geoptimaliseerd op reference-derived states om coördinaat- en oriëntatiefouten te minimaliseren, wat resulteert in een model met superieure geometrische voorspellingsmogelijkheden vergeleken met het basis H-DiffAb.

C. On-Policy Distillation (De Kerninnovatie)

Het "Student"-model wordt post-getraind met een mixed-objective benadering:

1. Student Rollout: De student genereert reverse-denoising trajecten (rollouts) beginnend vanaf volledig ruisige toestanden.
2. State Matching: Op specifieke tijdstappen ($t \in \{80, 70, \dots, 5\}$) evalueert de bevroren Teacher dezelfde door de student gegenereerde toestand ($R^S_t$) die de student op dat moment verwerkt.
3. Coordinate-Transition Distillation: De student wordt getraind om de voorspelde coördinaatruis ($\hat{\epsilon}^{pos}$) van de Teacher op deze rollout-toestanden te matchen. Dit wordt geformuleerd als een reverse-KL divergentie minimalisatie, vereenvoudigd tot een $L_2$-loss op coördinaatruis-voorspellingen:
   $$\mathcal{L}_{OPD}^{pos} = \mathbb{E} \left[ \sum_{t \in \mathcal{T}} \sum_{i \in \mathcal{G}} \| \hat{\epsilon}^{pos}_{S, t, i} - \hat{\epsilon}^{pos}_{T, t, i} \|^2_2 \right]$$
   Cruciaal is dat deze supervisie alleen wordt toegepast op $C_\alpha$-coördinaten, aangezien deze direct de backbone-geometrie en loop-plaatsing controleren. Sequentie- en oriëntatievoorspellingen worden in deze stap niet gedistilleerd om overfitting op de mogelijk imperfecte sequentievoorspellingen van de teacher te voorkomen.
4. Offline Denoising Anchor: Om de gezamenlijke sequentie-structuur-oriëntatie doelstelling die tijdens de pretraining is geleerd te behouden, wordt de originele denoising loss ($\mathcal{L}_{anchor}$) behouden op onafhankelijk gesamplede reference-derived states.
5. Finale Doelstelling: $\mathcal{L}_{ABOPD} = \mathcal{L}_{anchor} + \beta \mathcal{L}_{OPD}^{pos}$ (met $\beta=0.6$).

3. Belangrijkste Bijdragen

1. Identificatie van State-Distribution Mismatch: Het artikel belicht dat de kloof tussen reference-derived states en student rollout states tijdens het denoising-proces groter wordt, vooral in de latere verfijningsstadia waar geometrische correcties het meest kritiek zijn.
2. ABOPD Framework: Het voorstel van een on-policy distillatiemethode die $C_\alpha$ coördinaat-transitie-distillatie op student rollouts combineert met een offline denoising anchor.
3. Backbone-Aware Teacher: De demonstratie dat een teacher-model, verrijkt met geprivilegieerde natuurlijke geometrie, significant nauwkeurigere geometrische targets biedt dan standaardmodellen, zelfs als de sequentienauwkeurigheid niet superieur is.

4. Experimentele Resultaten

De methode werd geëvalueerd op twee benchmarks: Simultaneous multi-CDR redesign (19 SAbDab complexen) en RAbD CDR-H3 generatie (60 target gevallen).

• RAbD CDR-H3 Prestaties: ABOPD verminderde de $C_\alpha$ RMSD van 2.37 Å naar 1.95 Å (een verbetering van 0.42 Å) vergeleken met de H-DiffAb baseline. Het presteerde ook beter dan zowel SFT als offline distillatie controles.
• Multi-CDR Redesign: ABOPD verbeterde de RMSD over alle zes CDR's, met de meest significante winst op CDR-H3 (2.94 Å $\to$ 2.48 Å). Het verminderde ook de intra-CDR en CDR–antigeen clash-ratio's aanzienlijk (van 0.93% naar 0.42% en 1.58% naar 0.67%, respectievelijk).
• Geometrische Getrouwheid: Verbeteringen werden waargenomen in lDDT, TM-score en DockQ, wat duidt op een betere lokale geometrie, globale heavy-chain structuur en interface-plaatsing.
• Mechanisme Analyse: Een analyse op tijdstapniveau toonde aan dat het "Teacher voordeel" (de foutreductie die de teacher biedt ten opzichte van het basismodel) pas in latere tijdstappen ($t \le 50$) positief wordt. Dit bevestigt dat on-policy supervisie het meest effectief is tijdens de latere stadia van geometrische verfijning, wanneer het traject van de student al is afgeweken van de referentiedistributie.
• Ablatie Studies: Het distilleren van enkel de coördinaat-tak leverde de beste resultaten op. Het toevoegen van sequentie- of oriëntatie-distillatie bood geen verdere voordelen en degradeerde in sommige gevallen de prestaties.

5. Betekenis en Claims

Het artikel beweert dat ABOPD een pad biedt naar hogere-getrouwheid eiwitontwerp door de specifieke beperkingen van standaard post-training strategieën aan te pakken.

• Complementaire Aanpak: De auteurs positioneren ABOPD als een complementaire richting voor grootschalige pretraining en nieuwe generatieve architecturen. Het demonstreert dat effectieve structurele post-training niet alleen afhangt van het supervisiesignaal, maar van het toepassen van dat signaal op de juiste toestanden (namelijk, de toestanden die tijdens de generatie worden bezocht).
• Bescheiden Omvang: De auteurs geven expliciet aan dat de methode de backbone-reconstructie en geometrische nauwkeurigheid verbetert, maar niet noodzakelijkerwijs claimt alle aspecten van antilichaamontwerp op te lossen (bijv. de winsten in sequentie-reconstructie waren bescheiden). De primaire bijdrage is de validatie van on-policy distillatie voor het corrigeren van geaccumuleerde geometrische fouten in diffusie-gebaseerde antilichaamgeneratie.