Resumo Técnico: Escalando a Verificação de Redes Neurais com Paralelismo de Tensores e Paralelismo de Dados Totalmente Fragmentado

1. Definição do Problema

A verificação formal de redes neurais visa provar matematicamente que uma rede satisfaz propriedades de segurança para todas as entradas dentro de um domínio especificado. Embora algoritmos de propagação de limites como IBP, CROWN e $\alpha$-CROWN sejam o estado da arte para essa tarefa, sua aplicação prática é severamente limitada pela memória da GPU. Implementações padrão exigem que as matrizes de pesos e as matrizes $A$ de relaxação linear residam inteiramente em um único acelerador. Para redes grandes (por exemplo, uma única camada oculta com largura $h=262.144$), as matrizes de pesos sozinhas podem exceder 4 GB em float32, deixando memória insuficiente para as matrizes $A$, o processamento em lote (batching), ou tensores intermediários necessários para a verificação. Este artigo aborda a necessidade de remover esse gargalo de memória de GPU única, adaptando estratégias de fragmentação de parâmetros de treinamento de larga escala para o framework de verificação.

2. Metodologia

Os autores adaptam duas técnicas de paralelismo — Paralelismo de Tensores (TP) e Paralelismo de Dados Totalmente Fragmentado (FSDP) — ao framework de verificação auto LiRPA / $\alpha, \beta$-CROWN.

2.1 Paralelismo de Tensores (TP)

O TP divide as multiplicações de matrizes dentro de uma única camada entre $P$ GPUs.

• Implementação: Os autores implementaram operadores customizados BoundLinearTP Col e BoundLinearTP Row. Os pesos são fragmentados ao longo das dimensões de saída (coluna) ou entrada (linha).
• Mecanismo: Em um par Coluna–Linha, a computação é distribuída, exigindo uma operação AllReduce para recuperar o resultado.
• Limitação: Quando ativações ReLU ocorrem dentro de uma zona fragmentada (entre camadas de Coluna e Linha), a passagem reversa (backward pass) do CROWN requer limites intermediários para determinar as inclinações de relaxação ($\lambda_i$). Calcular esses limites exatamente exigiria comunicação entre GPUs em cada etapa. Para preservar a localidade, os autores substituem os limites por uma versão mais simples de Propagação de Limites de Intervalo (IBP) para esses limites intermediários. Isso introduz um "efeito de encapsulamento" (wrapping effect), fazendo com que a precisão dos limites se degrade conforme a profundidade da zona fragmentada aumenta.

2.2 Paralelismo de Dados Totalmente Fragmentado (FSDP)

O FSDP fragmenta os tensores de parâmetros igualmente entre $P$ GPUs, mas os reconstrói para a computação.

• Implementação: Implementado com aproximadamente 100 linhas de Python, o FSDP substitui os tensores BoundParams por fragmentos (shards). Antes da computação de uma camada, um AllGather monta a matriz de pesos completa; imediatamente após, o tensor completo é liberado, mantendo apenas o fragmento.
• Mecanismo: Esta abordagem garante que, no máximo, uma matriz de pesos completa resida na memória da GPU em qualquer instante. Crucialmente, o FSDP não altera a ordem ou o tipo das operações aritméticas; ele apenas altera o local de armazenamento dos pesos.
• Corretude: Como a aritmética permanece idêntica à linha de base de GPU única, o FSDP produz limites bit a bit idênticos ($\Delta = 0$).
• Integração: O FSDP integra-se perfeitamente com a verificação completa ($\beta$-CROWN + Branch-and-Bound) e camadas convolucionais (BoundConv).

2.3 Configuração Experimental

Os experimentos foram conduzidos em GPUs NVIDIA A40 usando PyTorch 2.x e a biblioteca auto LiRPA. Os autores introduziram uma flag bab.force_synchronous para garantir comparações justas entre os modos de GPU única e distribuído, desativando o redimensionamento automático de lote e a poda (pruning) no caminho de GPU única, garantindo cargas de trabalho idênticas.

3. Principais Contribuições

1. Implementação de TP para Verificação: Os autores implementaram TP para o auto LiRPA usando operadores customizados, alcançando aproximadamente $2\times$ de redução de memória de pico para $P=2$. Eles analisaram formalmente o compromisso, confirmando que, embora a integridade (soundness) seja preservada, a precisão dos limites degrada com o número de zonas fragmentadas devido à substituição por IBP.
2. Implementação de FSDP: Implementaram o FSDP para propagação de limites, alcançando 80–90% de redução de memória da linha de base e 34–39% de redução de memória de pico em MLPs largas. Crucialmente, este método produz limites bit a bit idênticos à linha de base de GPU única.
3. Integração com Verificação Completa: O FSDP foi integrado com sucesso ao $\beta$-CROWN e Branch-and-Bound (BaB), permitindo a primeira verificação completa (resultado UNSAT) de um modelo ResNet-large do CIFAR-100.
4. Identificação de Gargalos: Através de experimentação extensiva, os autores identificaram que, no modo $\alpha$-CROWN + BaB, o principal gargalo de memória não são as matrizes de pesos, mas os tensores alpha por neurônio (inclinações de ReLU).

4. Resultados

4.1 Redução de Memória

• TP: Alcançou um fator de redução de memória de pico de $\approx 1.985$ (próximo ao teórico $P=2$) para um MLP largo ($h=262.144$).
• FSDP:
  • Memória da Linha de Base: Reduzida em 80–90% em vários tamanhos de modelos, pois apenas os fragmentos são alocados antes da computação de limites.
  • Memória de Pico: Reduzida em 34–39% em MLPs largas ($h \ge 4096$). As economias são menores do que as economias da linha de base porque as matrizes $A$ não fragmentadas (que escalam com o quadrado da largura) permanecem na memória.

4.2 Precisão e Integridade dos Limites

• TP:
  • Modelos com uma zona fragmentada (sem ReLU intermediária) atingiram concordância de precisão de máquina ($\sim 10^{-8}$) com a referência de GPU única.
  • Modelos com múltiplas zonas fragmentadas mostraram desvio significativo (ex: $\sim 750$ unidades para um modelo de 3 zonas) devido à substituição por IBP, embora a integridade (superaproximação unilateral) tenha sido preservada.
• FSDP:
  • Alcançou exatamente zero de desvio ($\Delta lb = \Delta ub = 0.00$) em todos os modelos testados (MLPs e benchmarks ONNX) para ambos IBP e CROWN.
  • Verificou com sucesso um modelo Transformer (ViT) e ResNets convolucionais com resultados bit a bit idênticos à linha de base.

4.3 Verificação Completa (BaB)

No modo $\beta$-CROWN + BaB, o consumo de memória é dominado pelos tensores alpha (ex: 1.5 GB para tensores alpha vs. 1.5 MB para pesos em uma rede específica do MNIST).

• A fragmentação de pesos via FSDP proporcionou economias de memória de pico negligenciáveis neste modo porque os tensores alpha não foram fragmentados.
• No entanto, o FSDP permitiu com sucesso a verificação de um modelo ResNet-large (VNN-COMP 2024) que seria de outra forma inviável ou difícil de gerenciar, confirmando uma propriedade UNSAT sem o branching do BaB.

5. Significância e Limitações

O artigo demonstra que estratégias de fragmentação de parâmetros de treinamento podem ser transferidas para a verificação formal com um esforço de engenharia modesto.

• Significância: O FSDP oferece um caminho para verificar redes maiores com limites exatos, enquanto o TP oferece a maior redução de memória ao custo da precisão dos limites para redes profundas. O trabalho estabelece que o FSDP é compatível com os requisitos complexos de verificação completa (BaB) e diversas arquiteturas (CNNs, Transformers).
• Limitações:
  • Escalabilidade: Os experimentos foram limitados a $P=2$. Os autores observam que, à medida que $P$ aumenta, as matrizes $A$ não fragmentadas provavelmente se tornarão o custo dominante, levando a retornos decrescentes.
  • Precisão do TP: A degradação da precisão dos limites no TP para redes profundas limita sua utilidade para a verificação precisa de arquiteturas profundas.
  • Mudança de Gargalo: Os experimentos revelam que a pesquisa futura de verificação multi-GPU deve focar na fragmentação dos tensores alpha por neurônio, não apenas nos pesos, para alcançar um escalonamento significativo nos modos de verificação completa.

Os autores concluem que, embora o FSDP resolva o gargalo de armazenamento de pesos, o gargalo do tensor alpha continua sendo o desafio crítico para escalar a verificação completa para modelos maiores.