Riepilogo Tecnico: Design-CP

Definizione del Problema

I modelli generativi di proteine all-atom, come RFdiffusion 3 (RFD3), hanno il potenziale per progettare grandi complessi multimerici modellando congiuntamente tutte le catene. Tuttavia, le architetture sottostanti si basano su rappresentazioni di coppia (token-token e atomo-atomo) le cui spese di memoria scalano quadraticamente ($O(I^2)$ e $O(L^2)$) rispetto al numero di token ($I$) e di atomi ($L$). Man mano che la dimensione dell'assemblaggio target cresce, questi tensori di feature quadratici e le attivazioni intermedie superano rapidamente la capacità di memoria di una singola GPU. Questo "limite del singolo dispositivo" costringe le attuali pipeline a fare affidamento sul docking di corpi rigidi di oligomeri progettati indipendentemente, piuttosto che su una generazione congiunta end-to-end, in particolare per grandi assemblaggi simmetrici come le nanoparticelle icosaedriche o ottaedriche. Sebbene i cluster multi-GPU siano standard in campi adiacenti come la modellazione di grandi modelli linguistici, l'applicazione di strategie di parallelismo simili ai modelli strutturali è stata limitata dalle specifiche modalità di memoria e comunicazione dei meccanismi di attenzione di coppia.

Metodologia

Gli autori introducono Design-CP, un framework che implementa due strategie di inferenza di parallelismo di contesto (CP) per RFD3 che distribuiscono le attivazioni quadratiche su un mesh multi-GPU preservando i pesi pre-addestrati (ovvero, senza richiedere ri-addestramento o fine-tuning).

1. Strategie di Inferenza

Il documento implementa e confronta due schemi di partizionamento (sharding):

• 1D Row-Sharding: Una strategia leggera che suddivide la rappresentazione di coppia ($Z$ e $P$) su $P$ GPU lungo un singolo asse (query). Le chiavi (keys) e i valori (values) sono replicati su ogni GPU. Questo schema riformula l'auto-attenzione come cross-attenzione, richiedendo un singolo ALLGATHER delle attivazioni a traccia singola dopo ogni blocco di attenzione per garantire che tutti i dispositivi abbiano chiavi e valori identici per il passaggio successivo.
• 2. 2D Grid Sharding: Una strategia adattata da Fold-CP, che suddivide il tensore di coppia su una griglia di dispositivi $\sqrt{P} \times \sqrt{P}$. Questo approccio utilizza la Ring Attention (Liu et al., 2023) per calcolare i punteggi di attenzione on-the-fly. Le query sono partizionate lungo le righe, mentre le chiavi e i valori sono partizionati lungo le colonne e fatti circolare tramite spostamenti ad anello (ring shifts). Ciò consente ai dispositivi di calcolare l'attenzione tra le righe di query locali e gli shard di chiavi/valori in entrata senza mai materializzare la matrice di attenzione completa $I \times I$. Lo schema 2D restringe la comunicazione collettiva a sottogruppi di dimensione $\sqrt{P}$, riducendo il volume di comunicazione e l'overhead di memoria transitorio rispetto allo sharding 1D.

2. Vincoli di Simmetria

Un componente critico della metodologia è l'applicazione di forti vincoli di simmetria del gruppo puntuale durante il campionamento. Per target altamente simmetrici (ad es., icosaedrici, ottaedrici), il modello di diffusione opera sull'intero complesso ma ri-sistematizza la sua predizione ad ogni passo di denoising. Solo le coordinate di una singola Sottounità Asimmetrica (ASU) vengono mantenute; le restanti subunità sono sovrascritte da copie deterministiche generate tramite operazioni di gruppo. Questo collassa efficacementmente lo spazio di progettazione, riducendo i gradi di libertà a una singola ASU pur modellando l'intero set di interazioni inter-subunità. Questo meccanismo permette al modello di rimanere coerente con l'intero contesto dell'assemblaggio senza superare i limiti nativi di crop dell'addestramento.

3. Dettagli di Implementazione

• Ottimizzazione della Memoria: Entrambi gli schemi attivano automaticamente un "chunked pairwise embedder" (modalità a bassa memoria) che costruisce il tensore di coppia atomica $P$ on-the-fly su indici k-nearest neighbor (kNN) sparsi, evitando la materializzazione densa.
• Determinismo: Gli autori implementano specifici guardiani di comunicazione (ad es., il broadcasting di feature di token aggregate dal rank 0) per garantire l'equivalenza numerica con l'inferenza su singola GPU, affrontando il non-determinismo nelle riduzioni di floating-point atomiche.
• Hardware: I metodi sono testati sia su GPU per data center ad alta memoria (NVIDIA HG200, 95 GB) che su GPU di classe workstation (NVIDIA RTX A4000, 16 GB).

Contributi Chiave

1. Framework Design-CP: L'introduzione di due schemi di inferenza di parallelismo di contesto (1D row-sharding e 2D grid-sharding) per RFD3 che consentono il denoising congiunto di grandi proteine multimeriche su più GPU senza modificare i pesi del modello.
2. Caratterizzazione della Scalabilità: Un'analisi che mostra come la dimensione massima fattibile della Sottounità Asimmetrica (ASU) scali con la radice quadrata del numero di GPU ($\sqrt{P}$), in modo coerente con la distribuzione dei costi di memoria quadratici. Lo schema 2D dimostra una superiore scalabilità del tempo di esecuzione (wall-clock) grazie alla sovrapposizione della comunicazione con il calcolo.
3. Progettazione Basata sulla Simmetria: La dimostrazione che i forti vincoli di simmetria del gruppo puntuale rendono il parallelismo di contesto "utilizzabile immediatamente" per la progettazione all-atom end-to-end di nanoparticelle icosaedriche. Questo approccio mitiga i problemi di shift della distribuzione che tipicamente degradano la qualità del campione quando si campiona al di fuori del regime di addestramento nativo.
4. Democratizzazione della Progettazione: Il successo nella progettazione di nanoparticelle ottaedriche su un piccolo cluster di GPU di classe workstation (16 GB di memoria), provando che la progettazione di grandi assemblaggi proteici è fattibile senza l'accesso a hardware di alto livello.

Risultati

• Progettazione Icosaedrica: Utilizzando lo sharding 2D su un cluster a 4 GPU (griglia 2x2), gli autori hanno generato nanoparticelle icosaedriche con 60 catene (12.600 residui totali). La qualità del backbone per catena (rotture della catena, clash, struttura secondaria) ha ampiamente eguagliato quella dei monomeri RFD3 su singola GPU, sebbene gli assemblaggi simmetrici abbiano mostrato una leggera tensione geometrica (ad es., deviazioni più ampie della lunghezza del legame C$\alpha$–C$\alpha$). Le metriche di interfaccia consapevole della simmetria (clash, conteggio contatti) hanno seguito bene un baseline funzionale naturale (Lumazine Synthase, PDB: 1NQX), indicando interfacce inter-subunità fisicamente ragionevoli.
• Prestazioni di Scalabilità: Lo schema di sharding 2D ha ottenuto tempi di inferenza migliori rispetto allo schema 1D, completando spesso l'inferenza in circa la metà del tempo a parità di numero di GPU, grazie alla riduzione della latenza di comunicazione. La lunghezza massima dell'ASU prima degli errori Out-Of-Memory (OOM) è aumentata da circa 58 residui su una singola GPU a circa 237 residui su 16 GPU.
• Progettazione Ottaedrica su Hardware Comune: Gli autori hanno progettato con successo gabbie ottaedriche (24 catene, ~4.272 residui) su 16 GPU NVIDIA RTX A4000 (16 GB ciascuna). I design risultanti hanno esibito metriche di backbone sane e interfacce simmetriche ben formate, confermando la viabilità dell'approccio su hardware modesto.
• Ablazione sulla Simmetria: L'analisi visiva e quantitativa ha confermato che la rimozione dei vincoli di simmetria per sistemi grandi (ad es., 10.800 residui) ha prodotto strutture degenerate e non simili a proteine, sottolineando che la simmetria è essenziale per mantenere la qualità del campione a queste scale.

Significato e Rivendicazioni

L'articolo afferma che Design-CP rimuove la barriera della memoria della singola GPU che precedentemente rendeva necessario il docking per grandi assemblaggi proteici. Abilitando il denoising congiunto di tutti gli atomi in un singolo tragitto, esso sposta il paradigma della progettazione dal "dock-and-design" alla "generazione all-atom end-to-end".

Gli autori pongono questo lavoro come un passo verso la democratizzazione della progettazione di grandi assemblaggi proteici. Dimostrando che la progettazione di nanoparticelle simmetriche di alta qualità è possibile su GPU di classe workstation, il metodo abbassa la barriera hardware all'ingresso, ampliando potenzialmente la partecipazione alla progettazione proteica oltre i gruppi dotati di grandi risorse. Gli autori notano esplicitamente che i loro contributi risiedono nello strato di inferenza e parallelizzazione; non introducono nuovi pesi del modello o dati di addestramento. Riconoscono anche i limiti, specificamente che spingersi oltre i crop di addestramento nativi senza vincoli di simmetria può comunque degradare gli output per target asimmetrici, e che la validazione sperimentale degli assemblaggi progettati rimane un passo futuro necessario.