Resumen Técnico: Design-CP

Planteamiento del Problema

Los modelos generativos de proteínas de todos los átomos, como RFdiffusion 3 (RFD3), tienen el potencial de diseñar complejos multiméricos grandes mediante el modelado conjunto de todas las cadenas. Sin embargo, las arquitecturas subyacentes dependen de representaciones por pares (token-token y átomo-átomo) cuyas costes de memoria escalan cuadráticamente ($O(I^2)$ y $O(L^2)$) con el número de tokens ($I$) y de átomos ($L$). A medida que crece el tamaño del ensamblaje objetivo, estos tensores de características cuadráticos y las activaciones intermedias exceden rápidamente la capacidad de memoria de una sola GPU. Este "techo de un solo dispositivo" obliga a los pipelines actuales a depender del acoplamiento de cuerpos rígidos de oligómeros diseñados independientemente, en lugar de una generación conjunta de extremo a extremo, particularmente para ensamblajes simétricos grandes como las nanopartículas icosaédricas u octaédricas. Aunque los clústeres multi-GPU son estándar en campos adyacentes como el modelado de lenguaje de gran escala, la aplicación de estrategias de paralelismo similares a los modelos estructurales se ha visto limitada por los patrones específicos de memoria y comunicación de los mecanismos de atención por pares.

Metodología

Los autores presentan Design-CP, un marco que implementa dos estrategias de inferencia de paralelismo de contexto (CP) para RFD3 que distribuyen las activaciones cuadráticas a través de una malla multi-GPU preservando los pesos preentrenados (es decir, sin requerir reentrenamiento o ajuste fino).

1. Estrategias de Inferencia

El artículo implementa y compara dos esquemas de fragmentación (sharding):

• Fragmentación de filas 1D (1D Row-Sharding): Una estrategia ligera que fracciona la representación de pares ($Z$ y $P$) a través de $P$ GPUs a lo largo de un único eje (consultas/queries). Las claves (keys) y valores (values) se replican en cada GPU. Este esquema reformula la autoatención como atención cruzada, requiriendo un único ALLGATHER de las activaciones de vía única después de cada bloque de atención para asegurar que todos los dispositivos tengan claves y valores idénticos para el siguiente paso.
• Fragmentación de rejilla 2D (2D Grid Sharding): Una estrategia adaptada de Fold-CP, que organiza el tensor de pares en una rejilla de dispositivos de $\sqrt{P} \times \sqrt{P}$. Este enfoque utiliza Ring Attention (Liu et al., 2023) para computar las puntuaciones de atención sobre la marcha. Las consultas se fragmentan a lo largo de las filas, mientras que las claves y los valores se fragmentan a lo largo de las columnas y circulan mediante desplazamientos de anillo (ring shifts). Esto permite que los dispositivos computen la atención entre las filas de consulta locales y los fragmentos entrantes de clave/valor sin llegar a materializar la matriz de atención completa de $I \times I$. El esquema 2D restringe la comunicación colectiva a subgrupos de tamaño $\sqrt{P}$, reduciendo el volumen de comunicación y la sobrecarga de memoria transitoria en comparación con la fragmentación 1D.

2. Restricciones de Simetría

Un componente crítico de la metodología es la aplicación de fuertes restricciones de simetría de grupo puntual durante el muestreo. Para objetivos altamente simétricos (por ejemplo, icosaédricos, octaédricos), el modelo de difusión se ejecuta sobre el complejo completo pero resimetriza su predicción en cada paso de eliminación de ruido (denoising). Solo se retienen las coordenadas de una única Subunidad Asimétrica (ASU); las demás subunidades se sobrescriben mediante copias deterministas generadas vía operaciones de grupo. Esto colapsa efectivamente el espacio de diseño, reduciendo los grados de libertad a una sola ASU, mientras sigue modelando el conjunto completo de interacciones entre subunidades. Este mecanismo permite que el modelo se mantenga autoconsistente con el contexto del ensamblaje completo sin exceder los límites de recorte de entrenamiento nativo.

3. Detalles de Implementación

• Optimización de Memoria: Ambos esquemas habilitan automáticamente un "codificador de pares por fragmentos" (modo de baja memoria) que construye el tensor de pares de átomos $P$ sobre la marcha en índices dispersos de $k$-vecinos más cercanos (kNN), evitando la materialización densa.
• Determinismo: Los autores implementan guardas de comunicación específicas (por ejemplo, la transmisión de características de tokens agregadas desde el rango 0) para asegurar la equivalencia numérica con la inferencia de una sola GPU, abordando el no-determinismo en las reducciones de punto flotante atómico.
• Hardware: Los métodos se prueban tanto en GPUs de centro de datos de alta memoria (NVIDIA HG200, 95 GB) como en GPUs de nivel de estación de trabajo (NVIDIA RTX A4000, 16 GB).

Contribuciones Clave

1. Marco Design-CP: La introducción de dos esquemas de inferencia de paralelismo de contexto (fragmentación de filas 1D y fragmentación de rejilla 2D) para RFD3 que permiten el denoising conjunto de proteínas multiméricas grandes a través de múltiples GPUs sin modificar los pesos del modelo.
2. Caracterización de la Escalabilidad: Un análisis que muestra que el tamaño máximo factible de la Subunidad Asimétrica (ASU) escala con la raíz cuadrada del número de GPUs ($\sqrt{P}$), consistente con la distribución de los costes de memoria cuadráticos. El esquema 2D demuestra una escalabilidad de tiempo de ejecución superior debido al solapamiento de la comunicación con la computación.
3. Diseño Habilitado por Simetría: La demostración de que las fuertes restricciones de simetría de grupo puntual hacen que el paralelismo de contexto sea "utilizable de inmediato" para el diseño de todos los átomos de extremo a extremo de nanopartículas icosaédricas. Este enfoque mitiga los problemas de desplazamiento de distribución que típicamente degradan la calidad de la muestra cuando se muestrea fuera del régimen de entrenamiento nativo.
4. Democratización del Diseño: El diseño exitoso de jaulas octaédricas en un pequeño clúster de GPUs de nivel de estación de trabajo (16 GB de memoria), probando que el diseño de ensamblajes grandes es factible sin acceso a hardware de alto nivel.

Resultados

• Diseño Icosaédrico: Utilizando la fragmentación 2D en un clúster de 4 GPUs (rejilla 2x2), los autores generaron nanopartículas icosaédricas con 60 cadenas (12,600 residuos en total). La calidad del esqueleto por cadena (rupturas de cadena, choques, estructura secundaria) coincidió ampliamente con la de los monómeros de RFD3 de una sola GPU, aunque los ensamblajes simétricos mostraron una tensión geométrica ligeramente mayor (por ejemplo, desviaciones más amplias en la longitud de enlace C$\alpha$–C$\alpha$). Las métricas de interfaz con simetría (choques, recuentos de contactos) siguieron bien una base funcional natural (Lumazine Synthase, PDB: 1NQX), indicando interfaces entre subunidades físicamente razonables.
• Rendimiento de Escalamiento: El esquema de fragmentación 2D logró mejores tiempos de inferencia que el esquema 1D, completando a menudo la inferencia en aproximadamente la mitad del tiempo con el mismo número de GPUs, debido a la reducción de la latencia de comunicación. La longitud máxima de la ASU antes de los errores de falta de memoria (OOM) aumentó de ~58 residuos en una sola GPU a ~237 residuos en 16 GPUs.
• Diseño Octaédrico en Hardware Comercial: Los autores diseñaron con éxito jaulas octaédricas (24 cadenas, ~4,272 residuos) en 16 GPUs NVIDIA RTX A4000 (16 GB cada una). Los diseños resultantes exhibieron métricas de esqueleto saludables e interfaces simétricas bien formadas, confirmando la viabilidad del enfoque en hardware modesto.
• Ablación sobre la Simetría: El análisis visual y cuantitativo confirmó que eliminar las restricciones de simetría para sistemas grandes (por ejemplo, 10,800 residuos) resultaba en estructuras degeneradas y no similares a proteínas, subrayando que la simetría es esencial para mantener la calidad de la muestra a estas escalas.

Significado y Reivindicaciones

El artículo afirma que Design-CP elimina la barrera de memoria de una sola GPU que anteriormente hacía necesario el acoplamiento de cuerpos rígidos para grandes ensamblajes de proteínas. Al permitir el denoising conjunto de todos los átomos en una sola trayectoria, desplaza el paradigma de diseño de "acoplar y diseñar" a la "generación de todos los átomos de extremo a extremo".

Los autores posicionan este trabajo como un paso hacia la democratización del diseño de grandes ensamblajes. Al demostrar que el diseño de nanopartículas simétricas de alta calidad es posible en GPUs de nivel de estación de trabajo, el método reduce la barrera de hardware, permitiendo potencialmente una mayor participación en el diseño de proteínas más allá de los grupos con muchos recursos. Los autores señalan explícitamente que sus contribuciones están en la capa de inferencia y paralelización; no introducen nuevos pesos de modelo ni datos de entrenamiento. También reconocen las limitaciones, específicamente que forzar más allá de los recortes de entrenamiento nativos sin restricciones de simetría aún puede degradar los resultados para objetivos asimétricos, y que la validación experimental de los ensamblajes diseñados sigue siendo un paso futuro necesario.