Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper addresses the limitations of GPU memory in formal neural network verification by adapting Tensor Parallelism (TP) and Fully Sharded Data Parallelism (FSDP) techniques from large-scale model training to the auto_LiRPA/α,β-CROWN verification framework. The authors demonstrate that TP achieves approximately 2x peak-memory reduction while maintaining soundness on MNIST-FC benchmarks, and FSDP reduces baseline memory by 80-90% while producing bitwise identical bounds to single-GPU implementations. Notably, the findings reveal that the memory bottleneck in α-CROWN+BaB mode is primarily due to per-neuron alpha tensors, indicating a crucial area for future research.
Memory constraints in neural network verification can be drastically reduced by leveraging parallelism techniques, enabling more extensive and efficient safety property checks.
Formal neural network verification -- proving that a network satisfies safety properties for \emph{all} inputs in a specified domain -- is bounded in practice by GPU memory: standard implementations of bound-propagation algorithms (IBP, CROWN, $α$-CROWN) require weight and relaxation-coefficient matrices to reside entirely on one accelerator. We adapt two parallelism techniques originally developed for large-scale model training to the \texttt{auto\_LiRPA}\,/\,$α,β$-CROWN verification framework. \textbf{Tensor Parallelism (TP)} shards both weight and $A$-matrices across GPUs, achieving ${\approx}2\times$ peak-memory reduction at $P{=}2$; soundness is confirmed on VNN-COMP 2022 MNIST-FC benchmarks, though bound tightness degrades with the number of sharded zones due to forced IBP substitution for intermediate bounds inside sharded zones. \textbf{Fully Sharded Data Parallelism (FSDP)} shards only weight matrices with a per-layer \texttt{AllGather}, producing bounds that are \emph{bitwise identical} to the single-GPU baseline: baseline memory drops by 80--90\%, peak memory by 34--39\% on wide MLPs. FSDP integrates cleanly with complete verification ($β$-CROWN + Branch-and-Bound) and with convolutional layers (\texttt{BoundConv}); a complete \emph{unsat} result is obtained for CIFAR-100 ResNet-large (VNN-COMP 2024) under FSDP. Across all experiments the memory bottleneck in $α$-CROWN+BaB mode proves to be per-neuron alpha tensors, not weight matrices, pointing to the key direction for future work.