Resumo Técnico: Design-CP

Declaração do Problema

Modelos generativos de proteínas de todos os átomos, como o RFdiffusion 3 (RFD3), possuem o potencial de projetar grandes complexos multiméricos ao modelar conjuntamente todas as cadeias. No entanto, as arquiteturas subjacentes dependem de representações de pares (token-token e átomo-átomo) cujos custos de memória escalam quadraticamente ($O(I^2)$ e $O(L^2)$) com o número de tokens ($I$) e átomos ($L$). À medida que o tamanho do conjunto alvo cresce, esses tensores de características quadráticos e ativações intermediárias excedem rapidamente a capacidade de memória de uma única GPU. Este "teto de dispositivo único" força os pipelines atuais a dependerem do docking de corpos rígidos de oligômeros projetados independentemente, em vez de uma geração conjunta de ponta a ponta, particularmente para grandes montagens simétricas como nanopartículas icosaédricas ou octaédricas. Embora clusters multi-GPU sejam padrão em campos adjacentes, como a modelagem de grandes modelos de linguagem, a aplicação de estratégias de paralelismo semelhantes a modelos estruturais foi limitada pelos padrões específicos de memória e comunicação dos mecanismos de atenção de pares.

Metodologia

Os autores introduzem o Design-CP, um framework que implementa duas estratégias de inferência de paralelismo de contexto (CP) para o RFD3, distribuindo as ativações quadráticas através de uma malha multi-GPU enquanto preserva os pesos pré-treinados (ou seja, sem exigir retreinamento ou ajuste fino).

1. Estratégias de Inferência

O artigo implementa e compara dois esquemas de particionamento (sharding):

• Particionamento de Linha 1D (1D Row-Sharding): Uma estratégia leve que distribui a representação de par ($Z$ e $P$) através de $P$ GPUs ao longo de um único eixo (queries). As chaves (keys) e valores (values) são replicados em cada GPU. Este esquema reformula a autoatenção como atenção cruzada, exigindo um único ALLGATHER das ativações de trilha única após cada bloco de atenção para garantir que todos os dispositivos tenham chaves e valores idênticos para o próximo passo.
• Particionamento de Grade 2D (2D Grid Sharding): Uma estratégia adaptada do Fold-CP, que organiza o tensor de pares sobre uma grade de dispositivos $\sqrt{P} \times \sqrt{P}$. Esta abordagem utiliza Ring Attention (Liu et al., 2023) para computar escores de atenção em tempo real. As queries são particionadas ao longo das linhas, enquanto as chaves e valores são particionados ao longo das colunas e circulados via deslocamentos de anel (ring shifts). Isso permite que os dispositivos computem a atenção entre as linhas de query locais e os fragmentos de chave/valor recebidos sem jamais materializar a matriz de atenção completa $I \times I$. O esquema 2D restringe a comunicação coletiva a subgrupos de tamanho $\sqrt{P}$, reduzindo o volume de comunicação e o overhead de memória transitória em comparação ao particionamento 1D.

2. Restrições de Simetria

Um componente crítico da metodologia é a aplicação de fortes restrições de simetria de grupo de pontos durante a amostragem. Para alvos altamente simétricos (ex: icosaédricos, octaédricos), o modelo de difusão roda no complexo completo, mas resimetriza sua predição em cada passo de denoising. Apenas as coordenadas de uma única Subunidade Assimétrica (ASU) são retidas; as demais subunidades são sobrescritas por cópias determinísticas geradas via operações de grupo. Isso efetivamente colapsa o espaço de design, reduzindo os graus de liberdade a uma única ASU, enquanto ainda modela o conjunto completo de interações entre subunidades. Este mecanismo permite que o modelo permaneça autoconsistente com o contexto da montagem completa sem exceder os limites nativos de corte de treinamento.

3. Detalhes de Implementação

• Otimização de Memória: Ambos os esquemas habilitam automaticamente um "incorporador de pares em blocos" (chunked pairwise embedder - modo de baixa memória) que constrói o tensor de par de átomos $P$ sob demanda em índices esparsos de $k$-vizinhos mais próximos (kNN), evitando a materialização densa.
• Determinismo: Os autores implementam guardas de comunicação específicos (ex: broadcasting de características de tokens agregadas do rank 0) para garantir equivalência numérica com a inferência de GPU única, abordando o não-determinismo em reduções de ponto flutuante atômicas.
• Hardware: Os métodos são testados tanto em GPUs de data center de alta memória (NVIDIA HG200, 95 GB) quanto em GPUs de nível de estação de trabalho (NVIDIA RTX A4000, 16 GB).

Principais Contribuições

1. Framework Design-CP: A introdução de dois esquemas de inferência de paralelismo de contexto (particionamento de linha 1D e particionamento de grade 2D) para o RFD3 que permitem o denoising conjunto de grandes proteínas multiméricas através de múltiplas GPUs sem modificar os pesos do modelo.
2. Caracterização de Escalabilidade: Uma análise mostrando que o tamanho máximo viável da Subunidade Assimétrica (ASU) escala com a raiz quadrada do número de GPUs ($\sqrt{P}$), consistente com a distribuição dos custos de memória quadrática. O esquema 2D demonstra uma escalabilidade de tempo de execução (wall-clock) superior devido à sobreposição de comunicação com computação.
3. Design Habilitado por Simetria: A demonstração de que as fortes restrições de simetria de grupo de pontos tornam o paralelismo de contexto "utilizável imediatamente" para o design de todos os átomos de ponta a ponta de nanopartículas icosaédricas. Esta abordagem mitiga problemas de desvio de distribuição que tipicamente degradam a qualidade da amostra ao amostrar fora do regime de treinamento nativo.
4. Democratização do Design: O design bem-sucedido de nanocages octaédricas em um pequeno cluster de GPUs de nível de estação de trabalho (16 GB de memória), provando que o design de grandes montagens proteicas é viável sem acesso a hardware de alto desempenho de data center.

Resultados

• Design Icosaédrico: Usando o particionamento 2D em um cluster de 4 GPUs (grade 2x2), os autores geraram nanopartículas icosaédricas com 60 cadeias (12.600 resíduos no total). A qualidade do backbone por cadeia (quebras de cadeia, colisões, estrutura secundária) correspondeu amplamente à do RFD3 monômero de GPU única, embora as montagens simétricas tenham mostrado uma tensão geométrica ligeiramente maior (ex: desvios mais amplos em ligações C$\alpha$–C$\alpha$). As métricas de interface com consciência de simetria (colisões, contagem de contatos) acompanharam bem uma base funcional natural (Lumazine Synthase, PDB: 1NQX), indicando interfaces intersubunidades fisicamente razoáveis.
• Desempenho de Escalonamento: O esquema de particionamento 2D alcançou tempos de inferência melhores que o esquema 1D, frequentemente completando a inferência em aproximadamente metade do tempo com o mesmo número de GPUs, devido à redução da latência de comunicação. O comprimento máximo da ASU antes de erros de Out-Of-Memory (OOM) aumentou de ~58 resíduos em uma única GPU para ~237 resíduos em 16 GPUs.
• Design Octaédrico em Hardware Comum: Os autores projetaram com sucesso gaiolas octaédricas (24 cadeias, ~4.272 resíduos) em 16 GPUs NVIDIA RTX A4000 (16 GB cada). Os designs resultantes exibiram métricas de backbone saudáveis e interfaces simétricas bem formadas, confirmando a viabilidade da abordagem em hardware modesto.
• Ablação de Simetria: Análises visuais e quantitativas confirmaram que remover as restrições de simetria para sistemas grandes (ex: 10.800 resíduos) resultou em estruturas degeneradas e não proteicas, ressaltando que a simetria é essencial para manter a qualidade da amostra nessas escalas.

Significância e Alegações

O artigo alega que o Design-CP remove a barreira de memória de GPU única que anteriormente tornava necessário o docking de corpos rígidos para grandes montagens proteicas. Ao permitir o denoising conjunto de todos os átomos em uma única trajetória, ele muda o paradigma de design de "dock-and-design" para "geração de todos os átomos de ponta a ponta".

Os autores posicionam este trabalho como um passo em direção à democratização do design de grandes montagens. Ao demonstrar que o design de nanopartículas simétricas de alta qualidade é possível em GPUs de nível de estação de trabalho, o método reduz a barreira de hardware, potencialmente ampliando a participação no design de proteínas além de grupos com muitos recursos. Os autores explicitamente observam que suas contribuições estão na camada de inferência e paralelização; eles não introduzem novos pesos de modelo ou dados de treinamento. Eles também reconhecem limitações, especificamente que forçar além dos cortes de treinamento nativos sem restrições de simetria ainda pode degradar os resultados para alvos assimétricos, e que a validação experimental das montagens projetadas permanece um passo futuro necessário.