論文「Enabling AI ASICs for Zero Knowledge Proof (MORPH)」の技術的サマリー

本論文は、ゼロ知識証明（ZKP）の生成コストを削減するため、AI 向け ASIC（特に Google の TPU）を ZKP カーネル（多倍長スカラー乗算 MSM と数論的変換 NTT）に最適化して適用するフレームワーク「MORPH」を提案するものです。既存の GPU や FPGA 向けアルゴリズムをそのまま TPU に移植すると、データ配置（レイアウト）変換のコストや演算精度のミスマッチにより性能が大幅に低下する課題を解決し、TPU を ZKP 加速のための最先端プラットフォームへと変貌させました。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細をまとめます。

---

1. 問題定義 (Problem)

ZKP プロトコル（zk-SNARKs など）において、証明生成（Prover）の計算コストは依然として高く、特に以下の 2 つのカーネルがボトルネックとなっています。

• 多倍長スカラー乗算 (MSM): 全体のレイテンシの約 70% を占める。
• 数論的変換 (NTT): 全体のレイテンシの約 20-30% を占める。

これらの計算は、大規模な素数体（256 ビット〜753 ビット）上の演算を必要としますが、AI 向け ASIC（TPU など）は主に低精度（8 ビット、16 ビット、32 ビット）の行列演算（GEMM）に特化しています。既存の ZKP 加速手法を TPU に移植する際、以下の 2 つの根本的な課題が存在しました。

1. 算術レベルの課題（精度とキャリーチェーン）:
   • 素数体は自然な剰余数システム（RNS）分解ができないため、高精度な乗算には「桁分解」と「キャリー伝搬（Carry Propagation）」の連鎖が必要。
   • TPU の SIMD 構造では、この逐次的なキャリー伝搬がボトルネックとなり、計算利用率が 1% 以下にまで低下する。
2. データフローレベルの課題（レイアウト変換コスト）:
   • TPU は 4KB のベクトルレジスタ（VReg）単位で動作し、データが連続していることを前提としている。
   • 既存のアルゴリズム（Presorted Pippenger MSM や Butterfly NTT など）は、微細なデータシャッフルや転置を頻繁に行う。
   • これらのレイアウト変換（Shuffle/Transpose）は、TPU 上で非常に高コストであり、従来のビッグオー（Big-O）計算量モデルでは見落とされがちだが、実際には実行時間の支配的な要因となる。

---

2. 手法とアプローチ (Methodology)

MORPH は、TPU のアーキテクチャ特性に合わせて ZKP カーネルを再定式化するフレームワークです。その基盤として、新しい複雑性モデル「Big-T」を導入し、算術およびデータフローの両面で最適化を行いました。

2.1 Big-T 複雑性モデル

従来の Big-O 計算量モデルは、均質な逐次マシンを前提としており、TPU のような異種混合パイプライン処理やレイアウト変換のコストを捉えきれません。MORPH はBig-T 複雑性を提案しました。

• 定義: 異種パイプラインユニット（VPU, MXU, XLU など）のボトルネックと、レイアウト変換の遅延を考慮した漸近的な実行時間の下限。
• 目的: 算術演算の計算量だけでなく、データ配置変換やメモリ帯域幅の制約を定量的に評価し、TPU での最適設計指針を与える。

2.2 算術最適化：MXU 中心の RNS 遅延削減 (MXU-centric RNS Lazy Reduction)

高精度な素数体演算を TPU の低精度行列演算エンジン（MXU）に適合させる手法です。

• 拡張有限体の利用: 素数体 $F_M$ 上の演算を、RNS 分解が可能なより大きな合成数体 $F_Q$（$Q > M^2$）に埋め込む。
• キャリーチェーンの排除: 乗算を独立した 32 ビットベクトル演算に分解し、キャリー伝搬を不要にする。
• MXU による削減: 従来の逐次的なモンゴメリー削減（Radix Mont）を、RNS 再構成と素数体への還元を、8 ビット行列乗算（MXU 実行）とベクトル演算（VPU 実行）に再定式化。
  • これにより、$O(D^2)$ の逐次演算を $O(D)$ の並列行列演算に変換し、キャリー伝搬をクリティカルパスから排除しました。

2.3 データフロー最適化

TPU 間の通信と内部のレイアウト再編成を最小化する手法です。

• MSM 向け：レイアウト不変のピッペンジャー (LS-PPG)
  • 従来の Presorted Pippenger アルゴリズムでは、スカラーのソートとバケット集約（Bucket Accumulation）の間でデータ再配置が必要でした。
  • MORPH は、ソートと集約の両段階で共通のデータシャードとレイアウトを維持し、TPU 間通信や TPU 内部のレイアウト変換を排除します。
  • バケット削減（Bucket Reduction）を木構造で並列化し、VPU の利用率を向上させます。
• NTT 向け：5 ステップ NTT
  • 既存の 3 ステップ NTT は、行列サイズが大きすぎると MXU のスパン（計算時間）が増大し、メモリ帯域幅のボトルネックになります。
  • MORPH は、行方向の NTT をさらに 3 ステップ NTT に再帰的に分割する「5 ステップ NTT」を提案。
  • これにより、大きな行列演算をより小さな行列演算の集合に分解し、MXU の利用率を維持しつつ、パラメータ（Twiddle Factor）のメモリ使用量とレイアウト変換コストを削減します。

---

3. 主要な貢献 (Key Contributions)

1. Big-T 複雑性モデルの提案:
   • 異種ハードウェアにおけるボトルネックとレイアウト変換コストを定量化する新しい指標。ZKP アルゴリズムの設計指針として機能し、MSM で 5.8 倍、NTT で 1.6 倍の高速化をもたらしました。
2. MXU 中心の RNS 遅延削減アルゴリズム:
   • 大整数乗算における逐次的なキャリー伝搬を排除し、MXU 加速された行列演算へ変換。Radix Montgomery 削減と比較して最大 90 倍の高速化を達成。
3. LS-PPG と 5 ステップ NTT の実装:
   • TPU 間通信を不要にし、内部レイアウト変換を最小化するデータフロー。これにより、大規模な ZKP 計算におけるメモリ不足（OOM）を回避し、スケーラビリティを向上させました。

---

4. 実験結果 (Results)

Google TPUv6e-8 上で実装され、NVIDIA V100 搭載の SotA 実装（GZKP）と比較評価されました。

• NTT の性能:
  • 753 ビットの NTT において、GZKP（V100）と比較して最大 10 倍のスループット向上を達成。
  • 精度が上がる（256 ビット→753 ビット）につれて、GPU のレイテンシは 6〜7 倍増加するのに対し、TPU は 1.3〜3 倍の増加に留まり、RNS 演算の MXU への効率的なマッピングが有効であることを示しました。
• MSM の性能:
  • 753 ビット MSM において、GZKP と同等かそれ以上のスループット（最大 1.2 倍）を達成。
  • 256 ビットでは、微細なポイントの散乱（Scatter）によるメモリ帯域幅の制約により性能が劣りましたが、753 ビットのような高精度領域では TPU の優位性が明確になりました。
• アブレーション研究:
  • 算術最適化のみで MSM/NTT のレイテンシを 50〜90 倍削減。
  • データフロー最適化を組み合わせることで、さらに 3.1 倍（MSM）の高速化を実現。
  • バッチサイズを増やすことで VReg の利用率が向上し、さらに性能が改善することが確認されました。

---

5. 意義と結論 (Significance)

MORPH は、AI 向け ASIC（特に TPU）が ZKP 加速において GPU を凌駕しうることを実証した最初のフレームワークです。

• 技術的意義:
  • 「Big-O 計算量」だけでなく「レイアウト変換コスト」と「ハードウェア固有の制約」を考慮した設計思想の重要性を示しました。
  • 高精度な暗号計算を、低精度の行列演算エンジン（AI ASIC）で効率的に実行するための新しいパラダイム（RNS 埋め込みとレイアウト不変データフロー）を確立しました。
• 実用的意義:
  • ZKP の証明生成コストを大幅に削減し、スケーラブルなブロックチェーンやプライバシー保護された検証可能計算の実用化を加速します。
  • AI ASIC を ZKP 加速の標準的なプラットフォームとして位置づけ、今後の ZKP ハードウェア設計の方向性を示唆しています。

結論として、MORPH は AI 向け ASIC を ZKP カーネルに適用する際の障壁を克服し、TPU を高精度 NTT 処理における最先端のスループットマシンとして確立しました。