技術要約：テンソル並列性と完全シャード・データ並列によるニューラルネットワーク検証のスケーリング

1. 問題提起

形式的なニューラルネットワーク検証は、特定のドメイン内のすべての入力に対して、ネットワークが安全性を満たすことを数学的に証明することを目的としている。IBP、CROWN、および $\alpha$-CROWN などの境界伝搬アルゴリズムは、このタスクにおける最先端の手法であるが、実用においては GPU メモリによって深刻に制約されている。標準的な実装では、重み行列と線形緩和 $A$ 行列が単一のアレキセイタ上に完全に存在する必要がある。大規模なネットワーク（例：幅 $h=262,144$ の単一隠れ層）の場合、重み行列だけで float32 で 4 GB を超えることがあり、検証に必要な $A$ 行列、バッチング、または中間テンソルのためのメモリが不足する。本論文は、大規模なモデル学習からパラメータ・シャーディング戦略を適応させることで、この単一 GPU メモリのボトルネックを取り除く必要性に取り組んでいる。

2. 手法

著者らは、2 つの並列化技術であるテンソル並列（TP）と完全シャード・データ並列（FSDP）を auto LiRPA / $\alpha, \beta$-CROWN 検証フレームワークに適応させている。

2.1 テンソル並列 (TP)

TP は、単一レイヤー内での行列乗算を $P$ 個の GPU に分割する。

• 実装: 著者らはカスタム BoundLinearTP Col および BoundLinearTP Row オペレータを実装した。重みは出力（列）または入力（行）の次元に沿って分割される。
• メカニズム: 列–行ペアにおいて、計算は分散され、結果を復元するために 1 回の AllReduce 操作を必要とする。
• 制限: シャード化されたゾーン内（Column と Row レイヤーの間）で ReLU 活性化が発生する場合、CROWN のバックワードパスは、緩和スロープ（$\lambda_i$）を決定するために中間境界を必要とする。これらを正確に計算するには、各ステップで GPU 間の通信が必要となる。局所性を維持するため、著者らはこれらの中間境界に対して、より単純な区間境界伝搬（IBP）を代用している。これにより「ラッピング効果」が生じ、シャード化されたゾーンの深さが増すにつれて境界のタイトさが低下する。

2.2 完全シャード・データ並列 (FSDP)

FSDP は、パラメータテンソルを $P$ 個の GPU に等しく分割するが、計算のためにそれらを再構成する。

• 実装: 約 100 行の Python コードで実装されており、FSDP は BoundParams テンソルをシャードに置き換える。レイヤーの計算の前に、AllGather によってフルウェイト行列を組み立て、直後にフルテンソルを解放して、シャードのみを保持する。
• メカニズム: このアプローチは、いかなる瞬間においても最大 1 つのフルウェイト行列のみが GPU メモリ内に存在する状態を保証する。決定的なことは、FSDP は演算の順序や種類を変更せず、単に重みの格納場所を変更するだけであることである。
• 正当性: 算術演算がシングル GPU のベースラインと同一であるため、FSDP はビット単位で同一の境界（$\Delta = 0$）を生成する。
• 統合: FSDP は、完全な検証（$\beta$-CROWN + Branch-and-Bound）および畳み込みレイヤー（BoundConv）とシームレスに統合される。

2.3 実験設定

実験は、PyTorch 2.x および auto LiRPA ライブラリを使用し、NVIDIA A40 GPU 上で行われた。著者らは、自動バッチ・リサイズおよびプルーニングを無効にすることで、シングル GPU パスと分散モードの間で公平な比較を行うために、bab.force_synchronous フラグを導入し、同一のワークロードを保証した。

3. 主な貢献

1. 検証のための TP 実装: 著者らはカスタムオペレータを用いて auto LiRPA 用の TP を実装し、$P=2$ において約 $2\times$ のピークメモリ削減を達成した。彼らはトレードオフを正式に分析し、健全性は維持されるものの、シャード化されたゾーンの数が増えると IBP 代用により境界のタイトさが低下することを確認した。
2. FSDP 実装: 著者らは境界伝搬用の FSDP を実装し、幅広い MLP において 80–90% のベースラインメモリ削減と 34–39% のピークメモリ削減を達成した。決定的なことに、この手法はシングル GPU ベースラインとビット単位で同一の境界を生成する。
3. 完全な検証との統合: FSDP は $\beta$-CROWN および Branch-and-Bound (BaB) と正常に統合され、FSDP 環境下での CIFAR-100 ResNet-large モデルの初の完全な検証（UNSAT 結果）を可能にした。
4. ボトルネックの特定: 広範な実験を通じて、著者らは $\alpha$-CROWN + BaB モードにおいて、主要なメモリボトルネックは重み行列ではなく、ニューロンごとのアルファテンソル（ReLU スロープ）であることを特定した。

4. 結果

4.1 メモリ削減

• TP: 幅広い MLP ($h=262,144$) に対して、ピークメモリ削減係数 $\approx 1.985$（理論値 $P=2$ に近い）を達成した。
• FSDP:
  • ベースラインメモリ: 様々なモデルサイズにおいて 80–90% 削減された。これは、境界計算の前にシャードのみが割り当てられるためである。
  • ピークメモリ: 幅広い MLP ($h \ge 4096$) において 34–39% 削減された。節約幅がベースラインの削減幅よりも低いのは、アンシャードの $A$ 行列（幅の二乗に比例する）が依然としてメモリに保持されるためである。

4.2 境界の精度と健全性

• TP:
  • 1 つのシャード化ゾーンを持つモデル（中間 ReLU なし）は、シングル GPU リファレンスとマシン精度での一致（$\sim 10^{-8}$）を達成した。
  • 複数のシャード化ゾーンを持つモデルは、IBP 代用により顕著な偏差を示したが（例：3 ゾーンモデルで $\sim 750$ ユニット）、健全性（片側過大近似）は維持された。
• FSDP:
  • テストされたすべてのモデル（MLP および ONNX ベンチマーク）において、IBP と CROWN の両方で正確にゼロの偏差（$\Delta lb = \Delta ub = 0.00$）を達成した。
  • Transformer モデル (ViT) および畳み込み ResNet を、ベースラインとビット単位で同一の結果で正常に検証した。

4.3 完全な検証 (BaB)

$\beta$-CROWN + BaB モードでは、メモリ消費量はアルファテンソルによって支配される（例：特定の MNIST ネットワークにおいて、重みの 1.5 MB に対しアルファテンソルは 1.5 GB）。

• FSDP による重みのシャーディングは、アルファテンソルがシャード化されていないため、このモードにおけるピークメモリ削減にはほとんど寄与しなかった。
• しかし、FSDL は ResNet-large モデル（VNN-COMP 2024）の検証を成功させ、BaB 分岐なしで UNSAT 性を確定させた。これは、本来であれば不可能または管理困難であったものである。

5. 意義と限界

本論文は、学習からのパラメータ・シャーディング戦略が、適度なエンジニアリング努力によって形式的検証に転用できることを示している。

• 意義: FSDP は、正確な境界を持つより大きなネットワークを検証する道を提供し、TP は精度の代償として最高のメモリ削減を提供する。本研究は、FSDP が複雑な要件を持つ完全な検証（BaB）や多様なアーキテクチャ（CNN、Transformer）と互換性があることを確立した。
• 制限:
  • スケーラビリティ: 実験は $P=2$ に限定されている。著者らは、$P$ が増加するにつれて、アンシャードの $A$ 行列が支配的なコストとなり、収穫逓減が起こる可能性が高いと指摘している。
  • TP のタイトさ: 深いネットワークにおける TP の境界のタイトさの低下は、精密な検証におけるその有用性を制限する。
  • ボトルネックのシフト: 実験は、完全な検証をスケールアップさせるためには、単に重みだけでなく、ニューロンごとのアルファテンソルをシャード化することに焦点を当てる必要があることを明らかにしている。

著者らは、FSDP は重みのストレージ・ボトルネックを解決するが、アルファテンソルのボトルネックが、より大きなモデルへの完全な検証をスケールさせるための決定的な課題として残っていると結論づけている。