Technische Samenvatting: Design-CP

Probleemstelling

All-atom generatieve eiwitmodellen, zoals RFdiffusion 3 (RFD3), hebben het potentieel om grote multimerische complexen te ontwerpen door alle ketens gezamenlijk te modelleren. De onderliggende architecturen vertrouwen echter op paarrepresentaties (token-token en atoom-atoom) waarvan de geheugenkosten kwadratisch schalen ($O(I^2)$ en $O(L^2)$) met het aantal tokens ($I$) en atomen ($L$). Naarmate de grootte van de doelassemblage groeit, overschrijden deze kwadratische feature-tensoren en tussenliggende activaties snel de geheugencapaciteit van een enkele GPU. Deze "single-device ceiling" dwingt huidige pipelines om te vertrouwen op rigid-body docking van onafhankelijk ontworpen oligomeren in plaats van end-to-end gezamenlijke generatie, met name voor grote symmetrische assemblages zoals icosahedrale of octahedrale nanodeeltjes. Hoewel multi-GPU-clusters standaard zijn in aangrenzende velden zoals large language modeling, is de toepassing van vergelijkbare parallellismestrategieën op structuurmodellen beperkt gebleven door de specifieke geheugen- en communicatiepatronen van pairwise attention-mechanismen.

Methodologie

De auteurs introduceren Design-CP, een framework dat twee context-parallel (CP) inferentiestrategieën voor RFD3 implementeert die de kwadratische activaties over een multi-GPU-mesh distribueren terwijl de voorgetrainde gewichten behouden blijven (oftewel, geen hertraining of fine-tuning vereist).

1. Inferentiestrategieën

Het paper implementeert en vergelijkt twee sharding-schema's:

• 1D Row-Sharding: Een lichtgewicht strategie die de paarrepresentatie ($Z$ en $P$) langs een enkele as (queries) over $P$ GPU's verdeelt. Keys en values worden op elke GPU gerepliceerd. Dit schema herformuleert self-attention als cross-attention, waarbij een enkele ALLGATHER van de single-track activaties na elke attention-block nodig is om ervoor te zorgen dat alle apparaten identieke keys en values hebben voor de volgende stap.
• 2D Grid Sharding: Een strategie aangepast van Fold-CP, die de paar-tensor tegelt over een $\sqrt{P} \times \sqrt{P}$ device-grid. Deze aanpak maakt gebruik van Ring Attention (Liu et al., 2023) om attention-scores on-the-fly te berekenen. Queries worden langs rijen geshard, terwijl keys en values langs kolommen worden geshard en via ring-shifts worden gecirculeerd. Dit stelt apparaten in staat om attention te berekenen tussen lokale query-rijen en binnenkomende key/value-shards zonder ooit de volledige $I \times I$ attention-matrix te materialiseren. Het 2D-schema beperkt collectieve communicatie tot subgroepen van grootte $\sqrt{P}$, wat de communicatievolumes en de transiënte geheugenoverhead vermindert in vergelijking met 1D-sharding.

2. Symmetriebeperkingen

Een cruciaal onderdeel van de methodologie is de toepassing van sterke puntgroep-symmetriebeperkingen tijdens de sampling. Voor hoog-symmetrische targets (bijv. icosahedra, octahedra) draait het diffusiemodel op het volledige complex, maar resymmetriseert het de voorspelling bij elke denoising-stap. Alleen de coördinaten van een enkele Asymmetrische Subunit (ASU) worden behouden; de resterende subunits worden overschreven door deterministische kopieën gegenereerd via groepstransformaties. Dit reduceert effectief de ontwerpruimte door de vrijheidsgraden te verkleinen tot een enkele ASU, terwijl nog steeds de volledige set interacties tussen subunits wordt gemodelleerd. Dit mechanisme zorgt ervoor dat het model zelfconsistent blijft met de volledige assemblage-context zonder de native training crop-limieten te overschrijden.

3. Implementatiedetails

• Geheugenoptimalisatie: Beide schema's activeren automatisch een "chunked pairwise embedder" (low-memory mode) die de atoom-paar-tensor $P$ on-the-fly construeert op sparse $k$-nearest neighbor (kNN) indices, waardoor de dense materialisatie wordt vermeden.
• Determinisme: De auteurs implementeren specifieke communicatiebewakers (bijv. het broadcasten van geaggregeerde token-features van rank 0) om numerieke equivalentie met single-GPU inferentie te garanderen, wat de non-determinisme in atonische floating-point reducties aanpakt.
• Hardware: De methoden worden getest op zowel high-memory data-center GPU's (NVIDIA HG200, 95 GB) als workstation-grade GPU's (NVIDIA RTX A4000, 16 GB).

Belangrijkste Bijdragen

1. Design-CP Framework: De introductie van twee context-parallel inferentie-schema's (1D row-sharding en 2D grid-sharding) voor RFD3 die de gezamenlijke denoising van grote multimerische eiwitten over meerdere GPU's mogelijk maken zonder de modelgewichten te wijzigen.
2. Schaalbaarheidskarakterisering: Een analyse die aantoont dat de maximale haalbare Asymmetrische Subunit (ASU) grootte schaalt met de vierkantswortel van het aantal GPU's ($\sqrt{P}$), consistent met de distributie van kwadratische geheugenkosten. Het 2D-schema vertoont superieure wall-clock schaling vanwege het overlappen van communicatie met berekening.
3. Symmetrie-gestuurd Ontwerp: De demonstratie dat sterke puntgroep-symmetriebeperkingen context-parallellisme "out of the box bruikbaar" maken voor end-to-end, all-atom ontwerp van icosahedrale nanodeeltjes. Deze aanpak verzacht de distributieverschuiving die normaal gesproken de kwaliteit van samples verslechtert wanneer er buiten het native training-regime wordt gesampled.
4. Democratisering van Ontwerp: Het succesvolle ontwerp van octahedrale nanodeeltjes op een kleine cluster van workstation-grade GPU's (16 GB geheugen), wat bewijst dat ontwerp van grote assemblages mogelijk is zonder toegang tot high-end data-center hardware.

Resultaten

• Icosahedra Ontwerp: Met behulp van 2D sharding op een 4-GPU-cluster (2x2 grid) genereerden de auteurs icosahedrale nanodeeltjes met 60 ketens (12.600 residuen in totaal). De per-keten backbone-kwaliteit (ketenbreuken, botsingen, secundaire structuur) kwam grotendeels overeen met die van single-GPU RFD3 monomeren, hoewel de symmetrische assemblages een iets hogere geometrische spanning vertoonden (bijv. bredere C$\alpha$–C$\alpha$ bindingslengte-afwijkingen). Symmetrie-bewuste interface-metrieken (clashes, contactaantallen) volgden goed de natuurlijke functionele baseline (Lumazine Synthase, PDB: 1NQX), wat duidt op fysiek redelijke interacties tussen subunits.
• Schaalprestaties: Het 2D sharding-schema behaalde betere inferentietijden dan het 1D-schema, en voltooide de inferentie vaak in ongeveer de helft van de tijd bij hetzelfde aantal GPU's, dankzij de lagere communicatielatentie. De maximale ASU-lengte voordat Out-Of-Memory (OOM) fouten optraden, nam toe van ~58 residuen op een enkele GPU naar ~237 residuen op 16 GPU's.
• Octahedra Ontwerp op Commodity Hardware: De auteurs ontwierpen succesvol octahedrale kooien (24 ketens, ~4.272 residuen) op 16 NVIDIA RTX A4000 GPU's (16 GB elk). De resulterende ontwerpen vertoonden gezonde backbone-metrieken en goed gevormde symmetrische interfaces, wat de levensvatbaarheid van de aanpak op bescheiden hardware bevestigt.
• Ablatie op Symmetrie: Visuele en kwantitatieve analyse bevestigde dat het verwijderen van symmetriebeperkingen voor grote systemen (bijv. 10.800 residuen) resulteerde in degeneratieve, niet-eiwitachtige structuren, wat onderstreept dat symmetrie essentieel is voor het behoud van de kwaliteit van de samples op deze schalen.

Betekenis en Claims

Het paper claimt dat Design-CP de single-GPU geheugenbarrière wegneemt die voorheen rigid-body docking voor grote eiwitassemblages noodzakelijk maakte. Door de gezamenlijke denoising van alle atomen in een enkele traject mogelijk te maken, verschuift het de ontwerpparadigma van "dock-and-design" naar "end-to-end all-atom generatie".

De auteurs positioneren dit werk als een stap naar de democratisering van het ontwerp van grote assemblages. Door aan te tonen dat hoogwaardig ontwerp van symmetrische nanodeeltjes mogelijk is op workstation-grade GPU's, verlaagt de methode de hardwarematige drempel voor deelname, wat potentieel de deelname aan eiwitontwerp buiten zeer kapitaalkrachtige groepen vergroot. De auteurs merken expliciet op dat hun bijdragen zich bevinden op de inferentie- en parallelliseringslaag; zij introduceren geen nieuwe modelgewichten of trainingsdata. Ze erkennen ook beperkingen, specifelijk dat het pushen voorbij de native training crops zonder symmetriebeperkingen nog steeds de output voor asymmetrische targets kan verslechteren, en dat experimentele validatie van de ontworpen assemblages een noodzakelijke toekomstige stap blijft.