用并行计算突破神经网络验证的内存瓶颈,让大模型安全验证成为可能。
Scaling Neural Network Verification with Tensor Parallelism and Fully Sharded Data Parallelism
- 引入张量并行与全分片数据并行,拆分权重和矩阵以降低内存占用。
- 在宽层多层感知机上实现80%-90%内存下降,峰值内存减少34%-39%。
- 支持完整验证流程和卷积层,可验证CIFAR-100 ResNet-large模型。
形式化神经网络验证——证明网络在指定输入范围内满足安全属性——受GPU内存限制:标准的界传播算法(IBP、CROWN、α-CROWN)要求权重和松弛系数矩阵完全驻留在单个加速器上。本文将大规模模型训练中发展的两种并行技术引入auto_LiRPA / α,β-CROWN验证框架。张量并行(TP)将权重和A矩阵跨GPU切分,在P=2时实现约2倍峰值内存降低;在VNN-COMP 2022 MNIST-FC基准上验证了其正确性,但因在切分区间的中间界被迫采用IBP,导致界紧度下降。全分片数据并行(FSDP)仅切分权重矩阵,并通过每层AllGather操作,生成与单卡基线比特相同的验证结果:在宽层MLP上,基线内存降低80–90%,峰值内存降低34–39%。FSDP可无缝集成完整验证(β-CROWN + Branch-and-Bound)和卷积层(BoundConv),成功完成对CIFAR-100 ResNet-large模型(VNN-COMP 2024)的非满足性验证。所有实验表明,在α-CROWN+BaB模式下,内存瓶颈主要来自每神经元的alpha张量,而非权重矩阵,指明未来工作重点。
原文摘要 · Abstract (English)
Formal neural network verification -- proving that a network satisfies safety properties for *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 auto_LiRPA / $α,β$-CROWN verification framework. 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. Fully Sharded Data Parallelism (FSDP) shards only weight matrices with a per-layer AllGather, producing bounds that are 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 (BoundConv); a complete 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。