Rubyで簡易STARK/zkVMを実装してCRubyとSpinelで速度比較してみた

STARKとzkVMの仕組みについて解像度上げるために、外部のgem依存ゼロのpure RubyでSTARK 証明系と 4 命令のzkVM「BabyStark」(BabyBear体上の最小STARKの意)を実装してみた。

もう 1 つの実験として、同じソースをCRubyとSpinel(MatzのRuby AOTコンパイラ)の両方で動かし、生成される証明がバイト単位で一致することをCIで検証。CRubyとSpinelで生成したバイナリのベンチマークも行ってみた。

なお本実装は教材用で、STARK 101*1系の古典的な構成(商多項式の合成 + FRI)を使っており、現在のSTARK(ethSTARK、Plonky3、RISC Zero 等)が使う DEEP-ALI*2は含まない。主にコードと理論の対応の分かりやすさを優先。

STARK が証明するもの

STARK(Scalable Transparent ARgument of Knowledge)は「ある計算を正しく実行した」ことを、計算をやり直すより高速に検証できる証明を作る仕組み。zkVMはこれをVMに適用したもので、「このプログラムをこの入力で実行したら、この出力になった」ことを証明する。

証明の流れを一言でまとめると:

計算の実行過程を表(実行トレース)にする
「トレースが正しい」ことを多項式の言葉に翻訳する(AIR: Algebraic Intermediate Representation)。正しいトレースであること ⟺ ある多項式たちが特定の点で 0 になること
「特定の点で 0 になる」⟺「ゼロ化多項式で割り切れる」なので、その商多項式が低次の多項式として存在することを示せばよい
低次性の証明には FRI(Fast Reed-Solomon IOP of Proximity)を使う
対話をなくすために Fiat-Shamir 変換(チャレンジをハッシュで導出)を適用する

BabyStarkではこの各ステップを 1 ファイルずつの Ruby コードにしている。

体の選択: BabyBear

すべての計算は有限体 ${F_p}$ 上で行う。BabyStarkではPlonky3やRISC Zeroと同じBabyBear素数 p = 2^31 − 2^27 + 1 = 2013265921 を使った。この理由は、

NTT フレンドリー性。p − 1 = 2^27 × 15 なので、位数 2^k(k ≤ 27)の乗法的部分群が存在し、2 冪サイズの NTT(数論変換)がそのまま使える。多項式の評価・補間が O(n log n) になる。NTTについては以前の記事参照。
Ruby実装(主にSpinel化)の要で、全要素が31ビットに収まるため、乗算の中間値がint64に収まる。つまり多倍長整数に頼らず、機械語の整数演算だけで体演算が完結する。

module Field
  def self.mul(a, b)
    (a * b) % Params::P   # 中間値 < 2^62
  end
end

体の要素は「0 以上 p 未満に正規化した Integer」で表し、FieldElementのようなクラスでラップしない。これは後述するSpinelサブセットへの適合でもある。

チャレンジは 4 次拡大体から取る

Fiat-Shamir で導出するチャレンジ(制約合成の乱数 α、FRI の折り畳み係数 β)を基礎体から取ると、チャレンジ空間が 2³¹ しかなく、「偽の多項式同士が偶然打ち消し合う」確率が ${2^{-31}}$ で頭打ちになる。そこで4次拡大体 ${F_p \lbrack X \rbrack / (X^{4} − 11) (≈ 2^{124})}$ を実装してチャレンジをそこから取る。これも Plonky3 と同じ構成。

証明系のコア:

BabyStark の証明系は2段構えになっている。

まず決まった1つの計算だけを証明する専用実装(stark.rb)を作り、
その後それを任意のプログラム用に一般化してzkVMにした(zkvm.rb)

証明系の仕組みは前者のほうが見通しがよいので、この章では冒頭の5ステップが実際のコードでどうなるかをフィボナッチ専用版で追う。

題材はSTARK 101と同じフィボナッチ平方数列 ${a_{n+2} = a_{n+1}^{2} + a_n^{²}}$ で、 ${a_0 = 1}$ が公開入力、 ${a_1}$ が証明者だけが知るwitness、 ${a_{1022}}$ が公開出力。「 ${a_{1022}}$ がこの値になる ${a_1}$ を知っている」ことを、 ${a_1}$ を明かさず検証者に1022ステップの再計算をさせずに証明する。

ステップ 1：実行トレースを作る

トレースの生成はただの計算実行で、今回の場合、漸化式を1023要素の配列に展開するだけ：

def build_trace(witness)
  trace = Array.new(1023, 0)
  trace[0] = 1              # 公開入力 a_0
  trace[1] = witness        # 秘密の a_1
  i = 2
  while i < 1023
    prev  = trace[i - 1]
    prev2 = trace[i - 2]
    trace[i] = Field.add(Field.mul(prev, prev), Field.mul(prev2, prev2))
    i += 1
  end
  trace
end

ステップ 2：トレースを多項式にしてコミットする

次に、この数列を「多項式の値の並び」とみなす。位数1024の乗法的部分群の生成元をgとして、 ${f(g^{i}) = a_i}$ となる多項式fを補間で作る。この補間・評価を担うのが NTT。係数 → 評価値が NTT（= 多項式の一括評価）、評価値 → 係数がその逆変換 INTT（= 補間）。BabyBear を選んだ理由のひとつはこの補間を利用するため（2 冪サイズの部分群の存在）。フィボナッチの例だと配列の要素である評価値から多項式を計算する。

# INTTで多項式の係数を導出(次数 ≤ 1022)
coeffs = interpolate_trace(trace)

わざわざ多項式にするのは、1つの点の嘘が全体に波及するから。数列のまま 1 要素を改竄しても矛盾は隣としか起きないけど、1023個の点で値を指定された次数1022以下の多項式は一意なので、1 点でも違う値を通したければ丸ごと別の多項式になるしかない。そして低次多項式同士は広いドメインで見るとほとんどの点で食い違う(次数dの多項式は根をd個しか持てない)。

続いて、この多項式fを実行トレースのドメイン（評価に使う値の集合*3）の8倍のサイズの8192個の評価ドメインで評価し直す。このような小さいドメイン上の値の列を、その多項式を介して、大きいドメイン上の評価値の列に引き伸ばすことをLDE（low-degree extension）と呼ぶ。情報量という意味ではLDEでは何も増えていないけど、低次多項式を過剰な数の点で評価した冗長表現はリードソロモン符号そのものであるため、LDEによりトレースをリードソロモン符号化したことになる。

lde = Poly.coset_lde(coeffs, 31, 8192)  # コセット 31·<h> 上の8192個の点で再評価

続いて、評価値でマークルツリーにコミットし、そのルートハッシュを検証者に送る。

trace_tree = MerkleTree.new(lde.map { |v| encode_field(v) })
transcript.absorb("trace_root", trace_tree.root) # ルートだけ検証者に送る

以降、証明者はどの位置の値を聞かれても、このツリーと整合する値しか答えられない。

ステップ３：制約を商多項式に変える

上記の例で実行トレースが正しいことは、

公開入力が ${f(g^{0}) = 1}$ つまり ${a_0 = 1}$ であり、
公開出力が ${f(g^{1022})}$ であり
公開入力から公開出力までの計算の各ステップが[tex: {f(g^{2}x) - f(gx)^{2} - f(x)² = 0}]が成立する
- ${x = g^{n}}$ の場合、 ${gx = g^{n+1}, g^{2}x = g^{n+2}}$ であるため、 ${a_{n+2} = a_{n+1}^{2} + a_n^{²}}$ つまり ${a_{n+2} - a_{n+1}^{2} - a_n^{²} = 0}$ から上記の式が成立する

という制約を満たすことになる。この制約を多項式が割り切れるという形に変換する。compute_cp_evalsより、

# f(x) — x は LDE 上の i 番目の点
v0 = lde[i]
# f(gx) — g はコセット上では「8 つ隣」にあたる
v1 = lde[(i + 8) % 8192]
# f(g²x)
v2 = lde[(i + 16) % 8192]

# (f(x) − 1) / (x − g⁰)
q0 = Field.mul(Field.sub(v0, 1), inv_d0[i])
# (f(x) − 公開出力) / (x − g¹⁰²²)
q1 = Field.mul(Field.sub(v0, output), inv_d1[i]) 
numer = Field.sub(v2, Field.add(Field.mul(v1, v1), Field.mul(v0, v0)))
# (f(g²x) − f(gx)² − f(x)²) / Z(x)
q2 = Field.mul(numer, z_inv)

# チャレンジ α で3つの商多項式の低次性の主張を1つの線形結合に合成
cp[i] = ExtField.add(
  ExtField.add(ExtField.scale(alphas[0], q0), ExtField.scale(alphas[1], q1)),
  ExtField.scale(alphas[2], q2)
)

制約がすべて成立していれば q0, q1, q2 は低次多項式の評価列になり、チャレンジαで合成した合成多項式（CP）も低次になる。逆に1点でも制約が破れていると商は多項式にならず、CPの評価列は「どんな低次多項式とも一致しない」ものになる。ここで「計算が正しい」という主張が「CPが低次」という主張に翻訳される。

ステップ４：低次性の証明

残りは、コミットしたCPの評価列が低次多項式に近いことの証明で、これを行うのがFRI。FRIは次数dの問題を次数d/2の問題に再帰的に還元していく。

多項式f(x)について、次数が偶数の項と奇数の項に分解する：
- たとえば、 ${f(x) = c_0 + c_1x + c_2x^{2} + c_3x^{3}}$ は、偶数次は ${c_0 + c_2x^{2}}$ で、奇数次は ${c_1x + c_3x^{3} = x\cdot(c_1 + c_3x^{2})}$ で、 ${x^{2}}$ を変数に取ることで上記のように次数が半分の２つの多項式に割ることができる。
- この分解は係数を知らなくても評価値だけでできる。分解した式にxと-xをそれぞれ代入すると、 ${f(x) = f_e(x^{2}) + x \cdot f_o(x^{2})}$ 、 ${f(-x) = f_e(x^{2}) - x \cdot f_o(x^{2})}$ という連立になり、和と差を取れば ${f_e(x^{2}) = \frac{f(x) +f(-x)}{2}, \quad f_o(x^{2}) = \frac{f(x) - f(-x)}{2x}}$ と、fの±のペアの値から ${f_e, f_o}$ の値が求まる。
ランダムなチャレンジβを使って、１の２つの多項式を１つに束ねる ${f'(y) = f_e(y) + \beta f_o(y)}$ 。 ${f_e, f_o}$ が両方とも低次であればf'も低次の多項式になる。
f(x)が次数d未満であれば、2のf'は次数d/2未満となり、これをlog d回繰り返すと定数（もしくは十分小さい次数）になり、検証者が直接チェックできるようになる。

2の式にy = x²を代入して1つの復元式にまとめると、次の層の各点の値は、元の層の±ペアの値のみから直接計算できる：

$f_{j+1}(x^{2}) = \frac{f_j(x) + f_j(-x)}{2} + \beta_j \cdot \frac{f_j(x) - f_j(-x)}{2x}$

（ ${f_0 = CP}$ 、 ${f_{j+1}}$ は ${f_j}$ を ${\beta_j}$ で畳んだ層）。証明者の畳み込みも、後述する検証者のクエリ検算も、実態はこの式の評価になる。この畳込みの処理がfri.rbの

# 8192 点 → 64 点まで 7 回畳む
while cur.length > 64
  tree = MerkleTree.new(cur.map { |v| FRI.encode_ext(v) }) # この層における評価列の各点の値をリーフとしてマークルツリーでコミット
  transcript.absorb("fri_root", tree.root)  # この層をコミットしてから
  beta = transcript.challenge_ext("fri_beta")  # β を入手
  cur = FRI.fold(cur, cur_shift, beta) # 畳み込み
  cur_shift = Field.mul(cur_shift, cur_shift)
end

64点になるまで畳み込むと、次数８未満の多項式の評価列が残るので、係数8個を検証者に送る。

ここで証明者が検証者に送るデータを整理しておくと、

ステップ2のトレースのマークルルート
上記FRIのマークルルートのリスト（畳み込みの回数分7個）
最終層の係数
ランダムな30個のクエリ位置の開示データ：
- マークルルートだけでは正しく畳み込まれたかは分からないので、検証者は評価ドメイン 8192 点の中から30個*4ランダムに選んで（クエリ位置）、その位置のデータを要求し、証明者はその値とマークルパスを提供する。位置をxとした場合、1箇所につき以下のデータを提供する。
  - トレースの3点 ${f(x), f(gx), f(g^{2}x)}$ の値＋マークルパス
  - FRI各層jでのペア ${f_j(x), f_j(-x)}$ の値＋マークルパス

各層jのペアが与えられると検証者は自分で畳み込みを計算できるため、クエリの検証は、

トレースの開示3点のマークルパスを検証し、
その3点からステップ3の式でCP(x)を再計算する。これがFRI層0の「あるべき値」
層0の開示ペアのマークルパスを検証し、クエリ位置側の値が2の再計算値と一致するか照合し、
ペアから畳み込みを再計算し、その結果が層1の開示値（マークルパス検証済み）と一致するか照合する。これを最終層まで 7 回繰り返す。
最終の折り畳み結果が「係数8個から直接評価した値」と一致するか照合する。

検証者は「トレース開示→CP→層0→層1→ … → 係数」という 1 本の計算チェーンを、コミット済みの値と自分で再計算した値だけでなぞり直す。証明者がどこかで嘘をコミットしていれば、βが各層のコミット後に決まる以上、ランダムな検査点でこの連鎖のどこかが高確率で破れる。

ステップ 5：Fiat-Shamirで非対話型に

ここまで「検証者のチャレンジ」として α、β、クエリ位置が登場したけど、実際には対話しない。チャレンジは「それまでに送信したすべてのデータのハッシュ」から決める。トランスクリプトの実装はハッシュチェーンそのもの(transcript.rb):

def absorb(label, data)   # 公開値・コミットメントを状態に混ぜる
  @state = Hashing.hash_hex(TAG_ABSORB + @state + label + data)
end

def squeeze(label)        # 状態からチャレンジを導出(状態も進む)
  @state = Hashing.hash_hex(TAG_CHALLENGE + @state + label)
end

zkVM化

上記はフィボナッチ専用版だったけど、次は任意のプログラムの実行を証明対象にする。ISAは以下の4命令に絞る。

MOVI：定数をレジスタにロードする（例：MOVI rd, imm = rd ← imm）。
ADD：体の加算（例：ADD rd, rs1, rs2 = rd ← rs1 + rs2）
MUL：体の乗算（例：MUL rd, rs1, rs2= rd ← rs1 × rs2）
JNZ：Jump if Not Zero(非ゼロならジャンプ)（例：JNZ rs, addr = rs ≠ 0 なら pc ← addr、rs = 0 なら次の命令へ）

rdが結果の書き込み先レジスタ（destination）、rs / rs1 / rs2 が読み出し元レジスタ（source）、immは命令に埋め込まれた定数(immediate)、addrはジャンプ先の ROM アドレスを表す。JNZはこのISA唯一の制御フロー命令で、ループも条件分岐も停止もすべてこれで表現する。

先程のフィボナッチ平方数列をこの4命令で書くとfib_square.asmより、

# フィボナッチ平方数列: a_{n+2} = a_{n+1}^2 + a_n^2
#
# 公開入力(初期レジスタ): r0 = a_0 = 1、r1 = a_1(witness 相当)
# ループを 1021 回まわすと r0 = a_1021, r1 = a_1022。
# 停止時の r1 が公開出力。
#
# レジスタ割り当て:
#   r0 = a_n, r1 = a_{n+1}, r2 = ループカウンタ, r3/r4 = 平方の一時値,
#   r5 = halt 用の非ゼロ定数, r6 = -1(デクリメント用), r7 = 0(コピー用)

        MOVI r2, 1021      # 残り反復回数
        MOVI r6, -1        # r6 <- p-1(カウンタのデクリメント用)
        MOVI r5, 1         # halt 用の非ゼロ定数
loop:   MUL  r3, r0, r0    # r3 <- a_n^2
        MUL  r4, r1, r1    # r4 <- a_{n+1}^2
        ADD  r0, r1, r7    # r0 <- a_{n+1}(r7 = 0 との加算でコピー)
        ADD  r1, r3, r4    # r1 <- a_n^2 + a_{n+1}^2 = a_{n+2}
        ADD  r2, r2, r6    # カウンタをデクリメント
        JNZ  r2, loop      # 残りがあれば継続
halt:   JNZ  r5, halt      # 自己ループ = 停止

zkVMのAIR化

フィボナッチ専用版では、多項式が特定の点で0になるという形の制約にすることで計算規則が１つの式に固定されていたけど、VMでは規則がROM上のプログラム次第で変わる。その可変規則を代数の言葉に翻訳する工程がAIR化。

実用的なzkVMでは、Lookup Argumentsと呼ばれる、トレースからプログラムROMへのルックアップ証明を使う。ただ、BabyStarkでは対象プログラムが数命令〜数十命令と短いことを前提にpc（プログラムカウンタ）をワンホット列に展開する方式*5を取った。

BabyStarkのzkVMで実行する場合は↓

require "babystark"

program = BabyStark::Assembler.assemble(File.read("examples/fib_square.asm"))
input = [1, 3_141_592, 0, 0, 0, 0, 0, 0]  # r0 = a_0, r1 = a_1(初期レジスタ)

proof = BabyStark::ZKVM.prove(program: program, public_input: input, trace_length: 8192)
proof.public_output  # 停止時の r0..r7(r1 = a_1022 = 1525593042)

BabyStark::ZKVM.verify(program: program, public_input: input,
                       public_output: proof.public_output, proof: proof)  # => true

Spinel版

ここからがもう 1 つの実験。Spinel はRubyのソースを型推論してCにコンパイルし、ランタイム依存なしの単一バイナリを吐くAOTコンパイラで、対応する Ruby サブセットで書けばネイティブ性能が得られる。

BabyStarkは最初から「lib/以下は Spinelサブセット準拠」という方針の下でコードを書いた。ただSpinel自体はまだ開発中のプロダクトということもあり、いくつか挙動に問題があったけどPR作ってマージされたのでコミットa9610564で正常に動作する。この辺りRubyの結果と比較しつつ開発できるのは便利。

ベンチマーク結果

実際にSpinelでビルドすると、高速化の効果が得られるのか計測してみた。計測対象はzkVM版のフィボナッチ平方数列で、6命令ループの反復回数をトレース長の枠に収まる最大値 (T−4)/6 に変えながら試してみた。

トレース長 ${2^{10} / 2^{12} / 2^{14}}$ の証明生成・検証を CRuby / CRuby+YJIT / Spinel で計測した（AMD Ryzen 9 5950X、CRuby 4.0.5）

証明生成

T	CRuby	+YJIT	Spinel
1,024	972 ms	446 ms	590 ms
4,096	3,588 ms	1,376 ms	3,163 ms
16,384	14,811 ms	5,503 ms	15,544 ms

検証

T	CRuby	+YJIT	Spinel
1,024	12 ms	7 ms	5 ms
4,096	16 ms	11 ms	7 ms
16,384	23 ms	16 ms	10 ms

検証はSpinel 全サイズ最速(CRubyの2.3倍)。ところが、証明の生成はYJITが全サイズ最速で、SpinelはT=16,384だとCRubyより遅かった。事前の予想(タイトな体演算ループのマイクロベンチではSpinelがCRubyの10〜24 倍)とはまるで違う結果になった。

その後、フェーズ別プロファイルやGCの実験で調査したところ：

YJIT が速いのはJITそのもので、GCではない。 CRuby で GC.disable して計測しても時間は変わらない(割り当て 245 万オブジェクト、minor GC 28 回が走っているのに)。証明のホットパスは Field.add/Field.mul のような小メソッド呼び出しを 1 点あたり数百回積む構造で、YJITのインライン化がまさにここに効く。CRubyの世代別GCにとって短命オブジェクトの嵐はほぼタダ。
Spinelは割り当ての多いフェーズで超線形に劣化する。 計算タイトなNTTではYJITと互角なのに、制約評価(毎点小配列を大量生成)は仕事量16倍に対し35.9 倍(YJIT は 17.4 倍でほぼ線形)にスケールする。SpinelのGCは世代を持たないマーク・アンド・スイープGCで、コレクションのたびに生存ヒープ全体(LDE 24 列 × 131k 要素)を走査するため、割り当て O(T) × マーク O(T) で O(T²) 成分が生じる。
追試として、拡大体値を「4要素 Array」から「4フィールドの不変クラス」に変える実験をした。Spinel は 8 スカラーフィールド以下の不変クラスを C 構造体としてスタック割り当てするため、生成CからGCへの割り当てが消えて2.5 倍高速化し YJIT と同着になる。ところが同じ変更が CRubyでは 1.4 倍の減速になる(オブジェクト生成 + インスタンス変数アクセスは配列より重い)。Spinelの最適とCRubyの最適が真逆を向いたわけで、BabyStark は「CRuby が正」の方針から配列表現を維持した。

まとめると: YJIT は「Ruby らしい書き方」をそのまま速くする。Spinel は「Cらしく書ける計算カーネル」を最速にするが、割り当てが激しいワークロードではGCの性格が支配する。検証だけ行うノード(検証はSpinel 最速・バイナリ 354KB・ランタイム依存なし)と、証明を作る側(YJIT)で使い分けるのが現状の最適解、というのは当初まったく予想していなかった結果となった。

*1:StarkWareが公開しているSTARK証明系をゼロから実装するハンズオンチュートリアル。

*2:コミット後に評価ドメイン外のランダム点zでの多項式の値を証明者に申告させる方式(ethSTARKで導入)。

*3:位数1024の部分群（ ${g^{0}, ..., g^{1023}}$ ）

*4:30というのは目標のセキュリティ水準から導出した値。

*5:N個の候補の内どれか１つという情報を、1つの数値ではなくN本の0/1の並びで表す方式

Develop with pleasure!

福岡でCloudとかBlockchainとか。