🤖 machine learning

Scaling Neural Network Verification with Tensor Parallelism and Fully Sharded Data Parallelism

本文将张量并行（Tensor Parallelism）和全分片数据并行（Fully Sharded Data Parallelism）应用于 $\alpha$ -CROWN 验证框架，以显著降低 GPU 显存占用，从而实现了对以往因显存限制而无法进行的、如 CIFAR-100 上 ResNet-large 等大规模神经网络的形式化验证。

原作者： Sergei Vorobyov, Eugene Ilyushin

发布于 2026-06-09

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Sergei Vorobyov, Eugene Ilyushin

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图证明一辆自动驾驶汽车永远不会发生碰撞，无论天气如何变化，也无论行人如何突然跳出。你不能仅仅把汽车测试一百万次；你需要一个数学上的“证明”，以确保它在每一种可能的情况下都是安全的。这被称为形式化神经网络验证（Formal Neural Network Verification）。

问题在于，进行这种证明对计算机内存的要求极高。这就像是在解一个巨大的拼图，但所有的拼图块（数据）和规则（数学逻辑）都必须放在一张单一的小桌子（单个图形处理器/GPU）上。如果拼图太大，桌子就会溢出，导致证明失败。

这篇论文介绍了两种解决这个拼图问题的方案，它们通过使用多张“桌子”（多个 GPU）协同工作来解决问题，借鉴了当今训练巨型 AI 模型的方法。

以下是他们两种主要解决方案的拆解，使用了简单的类比进行解释：

1. “拆分拼图”法（张量并行 - Tensor Parallelism）

核心思想： 想象你有一个巨大的拼图。与其让一个人拿着整个拼图，不如把它切成两半。A 人拿着左半部分，B 人拿着右半部分。他们各自处理自己的部分，并向对方喊出计算结果。

运作方式： 研究人员将“权重”（拼图块）和“规则”（数学逻辑）分散到两个 GPU 上。
好消息： 这将每个计算机所需的内存减少了近一半（约 2 倍缩减）。对于小型或浅层拼图，这种方法非常高效。
代价： 当拼图变得很深（有很多层）时，这两个人必须在看不见全貌的情况下，去猜测两部分之间的连接关系。为了节省时间，他们使用了一种“快速且粗略”的估计方法（称为 IBP）。
结果： 最终的证明仍然是安全的（即：如果车实际上是危险的，它不会误判为安全），但随着拼图变深，答案会变得有些“模糊”或不够精确。这就像是通过观察地平线来估算山的高度，而不是进行精确测量。

2. “共享图书馆”法（全分片数据并行 - FSDP）

核心思想： 想象一个图书馆，书实在太大了，一个书架放不下。与其为每个读者都复印整本书，不如把书拆分成一页一页。

运作方式： 研究人员将“权重”（书的页面）分散到不同的 GPU 上。
神奇之处： 当一台计算机需要进行计算时，它会迅速从其他计算机那里收集它需要的所有页面，完成数学运算，然后立即将这些页面放回原处。在任何特定时刻，没有任何一台计算机持有“整本书”。
好消息：
- 完美的准确度： 因为计算过程与单台计算机持有整本书时的做法完全一致，所以结果在“位级”（bit-for-bit）上是完全相同的。没有“模糊感”。
- 内存节省： 它节省了大量的内存（基础设置下节省 80–90%，峰值使用时节省 34–39%）。
代价： 它需要计算机之间进行一些“沟通”来收集页面，这会消耗一点时间，但内存的节省是非常值得的。

大惊喜：到底是什么堵塞了内存？

研究人员原本以为“权重”（拼图块或书页）是主要问题。但他们错了。

一旦他们使用这些新方法释放了权重所占用的空间，他们发现了真正的瓶颈：一种被称为**“Alpha 张量”（alpha tensors）**的特定数据。

类比： 想象你在解拼图。“权重”是拼图块，而“Alpha 张量”则是你为了记录进度，必须在每一个拼图块上贴上的“便利贴”。
发现： 在最先进的验证模式下（即使用一种称为“分支定界法/Branch-and-Bound”的方法来检查是否发生碰撞），这些“便利贴”占据了 99% 的内存，而不是拼图块本身。
结论： 尽管他们成功地将拼图块分散到了不同的计算机上，但“便利贴”仍然太大，无法装进内存。要解决最大的问题（例如验证用于自动驾驶的复杂 AI），未来的工作需要研究如何也将这些“便利贴”分散到多台计算机上。

结果总结

张量并行（Tensor Parallelism）： 非常擅长节省内存，但会让深度网络的答案变得不太精确。
FSDP： 保持答案完全精确，并节省大量内存。它成功验证了一个复杂的图像识别模型（ResNet），而该模型之前因为体积太大而无法被验证。
未来方向： 验证更大规模 AI 系统的关键不再仅仅是拆分权重，而是要研究如何拆分那些追踪验证过程的“便利贴”（Alpha 张量）。

简而言之，这篇论文展示了如何利用多台计算机来验证 AI 的安全性，但也揭示了我们在验证最庞大、最复杂的 AI 系统之前，仍有一个主要的内存障碍需要跨越。

技术摘要：利用张量并行与全分片数据并行实现神经网络验证的扩展

1. 问题陈述

形式化神经网络验证旨在从数学上证明网络在指定定义域内的所有输入都满足安全性属性。虽然像 IBP、CROWN 和 $\alpha$ -CROWN 这样的边界传播算法是该领域的先进技术，但它们的实际应用受到 GPU 显存的严重限制。标准实现要求权重矩阵和线性松弛 $A$ -矩阵必须完全驻留在单个加速器上。对于大型网络（例如，具有宽度 $h=262,144$ 的单隐藏层），仅权重矩阵就可能超过 4 GB（float32 格式），导致没有足够的显存用于 $A$ -矩阵、批处理或验证所需的中间张量。本文通过将大规模模型训练中的参数分片策略应用于验证框架，旨在解决这一单 GPU 显存瓶颈问题。

2. 方法论

作者将两种并行技术——张量并行（TP）和全分片数据并行（FSDP）——适配到了 auto LiRPA / $\alpha, \beta$ -CROWN 验证框架中。

2.1 张量并行 (TP)

TP 将单个层内的矩阵乘法拆分到 $P$ 个 GPU 上。

实现： 作者实现了自定义的 BoundLinearTP Col 和 BoundLinearTP Row 算子。权重沿输出（列）或输入（行）维度进行分片。
机制： 在一个“列-行”对中，计算是分布式的，需要一次 AllReduce 操作来恢复结果。
局限性： 当 ReLU 激活函数出现在分片区域内（位于 Column 层和 Row 层之间）时，CROWN 的反向传播过程需要中间边界来确定线性松弛斜率 ( $\lambda_i$ )。为了精确计算这些值，需要在每一步都进行 GPU 间通信。为了保持局部性，作者在这些中间边界处使用了更简单的区间边界传播（IBP）。这引入了“包裹效应”（wrapping effect），随着分片区域深度的增加，边界紧致度会下降。

2.2 全分片数据并行 (FSDP)

FSDP 将参数张量均匀地分片到 $P$ 个 GPU 上，但在计算时重新构建它们。

实现： 通过约 100 行 Python 代码实现，FSDP 将 BoundParams 张量替换为分片。在层计算之前，通过 AllGather 组装完整的权重矩阵；紧接着，释放完整的张量，仅保留分片。
机制： 这种方法确保在任何时刻最多只有一个完整的权重矩阵驻留在 GPU 显存中。至关重要的是，FSDP 不会改变算术操作的顺序或类型；它仅仅改变了权重的存储位置。
正确性： 由于算术运算与单 GPU 基准完全一致，FSDP 产生的边界是位一致的（ $\Delta = 0$ ）。
集成： FSDP 可以无缝集成到完整的验证（ $\beta$ -CROWN + 分支定界法/Branch-and-Bound）以及卷积层（BoundConv）中。

2.3 实验设置

实验在配备 NVIDIA A40 GPU 的 PyTorch 2.x 和 auto LiRPA 库上进行。作者引入了一个 bab.force_synchronous 标志，通过禁用单 GPU 路径中的自动批次调整和剪枝，确保单 GPU 与分布式模式之间的公平比较，从而保证工作负载完全一致。

3. 核心贡献

面向验证的 TP 实现： 作者使用自定义算子为 auto LiRPA 实现了 TP，在 $P=2$ 时实现了约 $2\times$ 的峰值显存减少。他们正式分析了这种权衡，确认虽然保证了完备性（soundness），但由于引入了 IBP 替代方案，边界紧致度会随分片区域数量的增加而下降。
FSDP 实现： 他们为边界传播实现了 FSDP，在宽 MLP 上实现了 80–90% 的基准显存减少和 34–39% 的峰值显存减少。至关重要的是，该方法产生的边界与单 GPU 基准是位一致的。
与完整验证的集成： FSDP 已成功集成到 $\beta$ -CROWN 和分支定界法（BaB）中，实现了在 FSDP 模式下对 CIFAR-100 ResNet-large 模型进行首次完整验证（得到 UNSAT 结果）。
瓶颈识别： 通过广泛的实验，作者发现，在 $\alpha$ -CROWN + BaB 模式下，主要的显存瓶颈不是权重矩阵，而是每个神经元的 alpha 张量（ReLU 斜率）。

4. 结果

4.1 显存减少

TP： 对于宽 MLP ( $h=262,144$ )，实现了约 $1.985 $的峰值显存减少因子（接近理论值$ P=2$）。
FSDP：
- 基准显存： 在各种模型规模下减少了 80–90%，因为仅在边界计算前分配分片。
- 峰值显存： 在宽 MLP ( $h \ge 4096$ ) 上减少了 34–39%。由于未分片的 $A$ -矩阵（其规模随宽度平方增长）仍驻留在内存中，因此节省的比例低于基准节省比例。

4.2 边界精度与完备性

TP：
- 具有一个分片区域（无中间 ReLU）的模型达到了与单 GPU 参考值机器精度的吻合（ $\sim 10^{-8}$ ）。
- 具有多个分片区域的模型显示出显著偏差（例如，对于 3 区模型，偏差约为 $\sim 750$ 单位），这是由于 IBP 替代导致的，尽管完备性（单侧过近似）得到了保持。
FSDP：
- 在所有测试的模型（MLP 和 ONNX 基准）中，对于 IBP 和 CROWN 均实现了完全零偏差 ( $\Delta lb = \Delta ub = 0.00$ )。
- 成功验证了 Transformer 模型（ViT）和卷积 ResNet，且结果与基准位一致。

4.3 完整验证 (BaB)

在 $\beta$ -CROWN + BaB 模式下，显存消耗主要由 alpha 张量决定（例如，在特定的 MNIST 网络中，alpha 张量为 1.5 GB，而权重仅为 1.5 MB）。

通过 FSDP 对权重进行分片在此时对峰值显存节省微乎其微，因为 alpha 张量并未被分片。
然而，FSDP 成功实现了对 ResNet-large 模型（VNN-COMP 2024）的验证，该模型在其他情况下难以管理或无法完成，并在没有 BaB 分支的情况下确认了 UNSAT 属性。

5. 意义与局限性

本文证明了来自训练阶段的参数分片策略可以以较低的工程成本迁移到形式化验证中。

意义： FSDP 提供了一条利用精确边界验证更大网络的路径，而 TP 则提供了最高的显存减少能力，但代价是深度网络中边界紧致度的下降。这项工作确立了 FSDP 与复杂验证需求（BaB）以及多样化架构（CNN、Transformer）的兼容性。
局限性：
- 可扩展性： 实验仅限于 $P=2$ 。作者指出，随着 $P$ 的增加，未分片的 $A$ -矩阵可能会成为主要的成本，从而导致收益递减。
- TP 紧致度： TP 在深度网络中导致的边界紧致度下降限制了其在精密验证深度架构方面的效用。
- 瓶颈转移： 实验表明，未来的多 GPU 验证研究必须专注于分片每个神经元的 alpha 张量，而不只是权重，才能在完整验证模式下实现显著的扩展。

作者总结道，虽然 FSDP 解决了权重存储的瓶颈，但 alpha 张量的瓶颈仍然是扩展完整验证以处理更大模型的关键挑战。