Resumen Técnico: Escalamiento de la Verificación de Redes Neuronales con Paralelismo de Tensores y Paralelismo de Datos Completamente Fragmentado

1. Declaración del Problema

La verificación formal de redes neuronales tiene como objetivo demostrar matemáticamente que una red satisface propiedades de seguridad para todas las entradas dentro de un dominio especificado. Aunque los algoritmos de propagación de límites como IBP, CROWN y $\alpha$-CROWN son el estado del arte para esta tarea, su aplicación práctica está severamente limitada por la memoria de la GPU. Las implementaciones estándar requieren que las matrices de pesos y las matrices $A$ de relajación lineal residan enteramente en un solo acelerador. Para redes grandes (por ejemplo, una sola capa oculta con anchura $h=262,144$), las matrices de pesos por sí solas pueden exceder los 4 GB en float32, dejando memoria insuficiente para las matrices $A$, el procesamiento por lotes (batching) o los tensores intermedios requeridos para la verificación. Este artículo aborda la necesidad de eliminar este cuello de botella de memoria de una sola GPU adaptando las estrategias de fragmentación de parámetros de la formación de modelos a gran escala al marco de verificación.

2. Metodología

Los autores adaptan dos técnicas de paralelismo —Paralelismo de Tensores (TP) y Paralelismo de Datos Completamente Fragmentado (FSDP)— al marco de verificación auto LiRPA / $\alpha, \beta$-CROWN.

2.1 Paralelismo de Tensores (TP)

El TP divide las multiplicaciones de matrices dentro de una sola capa a través de $P$ GPUs.

• Implementación: Los autores implementaron operadores personalizados BoundLinearTP Col y BoundLinearTP Row. Los pesos se fragmentan a lo largo de las dimensiones de salida (columna) o entrada (fila).
• Mecanismo: En un par Columna–Fila, la computación se distribuye, requiriendo una operación AllReduce para recuperar el resultado.
• Limitación: Cuando ocurren activaciones ReLU dentro de una zona fragmentada (entre capas de Columna y Fila), la pasada hacia atrás (backward pass) de CROWN requiere límites intermedios para determinar las pendientes de relajación ($\lambda_i$). Calcular estos exactamente requeriría comunicación entre GPUs en cada paso. Para preservar la localidad, los autores sustituyen los límites intermedios por la más simple Propagación de Límites de Intervalo (IBP). Esto introduce un "efecto de envoltorio" (wrapping effect), lo que provoca que la precisión de los límites se degrade a medida que aumenta la profundidad de la zona fragmentada.

2.2 Paralelismo de Datos Completamente Fragmentado (FSDP)

FSDP fragmenta los tensores de parámetros equitativamente entre $P$ GPUs pero los reconstruye para la computación.

• Implementación: Implementado con aproximadamente 100 líneas de Python, FSDP reemplaza los tensores BoundParams con fragmentos (shards). Antes de la computación de una capa, un AllGather ensambla la matriz de pesos completa; inmediatamente después, el tensor completo se libera, manteniendo solo el fragmento.
• Mecanismo: Este enfoque asegura que, en cualquier instante, exista como máximo una matriz de pesos completa en la memoria de la GPU. Crucialmente, FSDP no altera el orden ni el tipo de operaciones aritméticas; simplemente cambia la ubicación de almacenamiento de los pesos.
• Corrección: Debido a que la aritmética es idéntica a la línea base de una sola GPU, FSDD produce límites bit a bit idénticos ($\Delta = 0$).
• Integración: FSDP se integra perfectamente con la verificación completa ($\beta$-CROWN + Branch-and-Bound) y las capas convolucionales (BoundConv).

2.3 Configuración Experimental

Los experimentos se realizaron en GPUs NVIDIA A40 utilizando PyTorch 2.x y la biblioteca auto LiRPA. Los autores introdujeron un flag bab.force_synchronous para asegurar comparaciones justas entre los modos de una sola GPU y distribuidos, desactivando el redimensionamiento automático de lotes y la poda en la ruta de una sola GPU, asegurando así cargas de trabajo idénticas.

3. Contribuciones Clave

1. Implementación de TP para Verificación: Los autores implementaron TP para auto LiRPA utilizando operadores personalizados, logrando aproximadamente una reducción de memoria pico de $2\times$ con $P=2$. Analizaron formalmente el compromiso, confirmando que, si bien la solidez (soundness) se preserva, la precisión de los límites se degrada con el número de zonas fragmentadas debido a la sustitución por IBP.
2. Implementación de FSDP: Implementaron FSDP para la propagación de límites, logrando una reducción de memoria de la línea base del 80–90% y una reducción de memoria pico del 34–39% en MLPs anchos. Crucialmente, este método produce límites bit a bit idénticos a la línea base de una sola GPU.
3. Integración con Verificación Completa: FSDP se integró con éxito con $\beta$-CROWN y Branch-and-Bound (BaB), permitiendo la primera verificación completa (resultado UNSAT) de un modelo ResNet-large de CIFAR-100.
4. Identificación de Cuellos de Botella: A través de una experimentación extensiva, los autores identificaron que en el modo $\alpha$-CROWN + BaB, el principal cuello de botella de memoria no son las matrices de pesos, sino los tensores alfa por neurona (pendientes de ReLU).

4. Resultados

4.1 Reducción de Memoria

• TP: Logró un factor de reducción de memoria pico de $\approx 1.985$ (cercano al teórico $P=2$) para un MLP ancho ($h=262,144$).
• FSDP:
  • Memoria de la Línea Base: Se redujo en un 80–90% en varios tamaños de modelos, ya que solo se asignan los fragmentos antes de la computación de límites.
  • Memoria Pico: Se redujo un 34–39% en MLPs anchos ($h \ge 4096$). Los ahorros son menores que los de la línea base porque las matrices $A$ no fragmentadas (que escalan con el cuadrado de la anchura) permanecen en memoria.

4.2 Precisión y Solidez de los Límites

• TP:
  • Los modelos con una zona fragmentada (sin ReLU intermedia) lograron una concordancia de precisión de máquina ($\sim 10^{-8}$) con la referencia de una sola GPU.
  • Los modelos con múltiples zonas fragmentadas mostraron una desviación significativa (por ejemplo, $\sim 750$ unidades para un modelo de 3 zonas) debido a la sustitución por IBP, aunque se preservó la solidez (sobreaproximación de un solo lado).
• FSDP:
  • Logró una desviación exactamente cero ($\Delta lb = \Delta ub = 0.00$) en todos los modelos probados (MLPs y benchmarks de ONNX) tanto para IBP como para CROWN.
  • Verificó con éxito un modelo Transformer (ViT) y ResNets convolucionales con resultados bit a bit idénticos a la línea base.

4.3 Verificación Completa (BaB)

En el modo $\beta$-CROWN + BaB, el consumo de memoria está dominado por los tensores alfa (por ejemplo, 1.5 GB para tensores alfa frente a 1.5 MB para pesos en una red específica de MNIST).

• El fragmentar los pesos mediante FSDP proporcionó ahorros de memoria pico insignificantes en este modo porque los tensores alfa no fueron fragmentados.
• Sin embargo, FSDP permitió con éxito la verificación de un modelo ResNet-large (VNN-COMP 2024) que de otro modo sería inviable o difícil de gestionar, confirmando una propiedad UNSAT sin ramificación de BaB.

5. Significación y Limitaciones

El artículo demuestra que las estrategias de fragmentación de parámetros de la formación pueden transferirse a la verificación formal con un esfuerzo de ingeniería modesto.

• Significación: FSDP ofrece un camino para verificar redes más grandes con límites exactos, mientras que TP ofrece la mayor reducción de memoria a costa de la precisión de los límites para redes profundas. El trabajo establece que FSDP es compatible con los requisitos complejos de la verificación completa (BaB) y diversas arquitecturas (CNNs, Transformers).
• Limitaciones:
  • Escalabilidad: Los experimentos se limitaron a $P=2$. Los autores señalan que a medida que $P$ aumenta, las matrices $A$ no fragmentadas probablemente se convertirán en el costo dominante, lo que llevará a rendimientos decrecientes.
  • Precisión de TP: La degradación de la precisión de los límites en TP para redes profundas limita su utilidad para la verificación precisa de arquitecturas profundas.
  • Desplazamiento del Cuello de Botella: Los experimentos revelan que la investigación futura de multi-GPU en verificación debe centrarse en fragmentar los tensores alfa por neurona, no solo los pesos, para lograr un escalamiento significativo en los modos de verificación completa.

Los autores concluyen que, si bien FSDP resuelve el cuello de botella del almacenamiento de pesos, el cuello de botella de los tensores alfa sigue siendo el desafío crítico para escalar la verificación completa a modelos más grandes.