freivalds-pol

Cheap, sound verification of decentralized training steps — a research prototype aimed at Nous Research's Psyche / DisTrO network.

📄 Research writeup with figures: REPORT.md · design & derivations: docs/DESIGN.md · prototype→system log: CHANGES.md

Replace Psyche's ~2× redundant recompute with sub-1% probabilistic spot-checks, made sound on heterogeneous hardware by a floating-point model, hardened against adaptive adversaries, and stress-tested over full runs — including the finding that per-step verification matters most for the large, capacity-rich models where loss-monitoring alone is weakest.

Why

Psyche today verifies a node's training work by redundant recompute-and-compare across participants (plus Bloom filters to confirm DisTrO results were gossiped, and health checks for liveness). That costs roughly 2× the compute to verify 1× of work, scales poorly, and gives nodes no data privacy.

This project explores a cheaper verifier: catch a node that submits incorrect gradient work at far less than full-recompute cost, using

• Merkle commitments to the step transcript (model state, data shard, RNG seed, matmul outputs, submitted update),
• random challenges over the committed matmuls, and
• Freivalds' algorithm to check each challenged product in O(n²) instead of O(n³),

with a phase-2 path to a zero-knowledge spot-check (prove a challenged matmul is correct without shipping full activations or revealing a private shard).

See docs/DESIGN.md for the full threat model, protocol, and roadmap.

Layout

src/freivalds_pol/
  freivalds.py     Freivalds' probabilistic matmul check (right + left, explicit probes)
  numerics.py      floating-point error model: O(n^2) honest-noise bound, bf16/fp16 sim
  adaptive.py      adaptive adversaries (nullspace / fixed cheats)
  training.py      a real MLP training step -> MatMulRecords (+ finite-difference grad-check)
  transformer.py   a real Llama-style block step (RMSNorm/attention/GELU) -> 8 GEMMs, grad-checked
  compressor.py    simplified (1D-tiled) update compression + per-tile verifier
  demo.py          faithful 2D-chunk DeMo (decay/2D-DCT/top-k/error feedback) + per-block verifier
  trainer.py       multi-round DeMo training loop + budget-constrained adversaries
  curvature.py     Hessian-vector products + power iteration (flat/steep loss directions)
  model.py         multi-layer/multi-head transformer + AdamW (for scaling §8/§9)
  collusion.py     free-riding / update-copying detection + identity binding
  zk.py            non-interactive sumcheck argument for one matmul + PCS interface
  commitments.py   SHA-256 Merkle tree over transcript leaves
  transcript.py    StepTranscript / MatMulRecord (carries claimed dtype) + commitment
  challenge.py     random challenge sampling + Fiat-Shamir (commitment-derived) probes
  verifier.py      verify_step: commitment + shard + two-sided calibrated Freivalds + min-precision
  adversary.py     cheat transforms (lazy / wrong_compute / poison_shard)
experiments/
  run_detection.py detection rate vs #challenges; verifier cost vs recompute
  fp_crux.py       precision vs noise floor and the size of an undetectable cheat
  adaptive.py      predictable vs commit-then-sample probe; grinding; two-sided fix
  real_step.py     full verifier on a real, gradient-checked transformer-block step
  compressed.py    verify the DeMo-compressed update (DisTrO wire format) on a real gradient
  multiround.py    do sub-threshold cheats accumulate over a training run? (the hard one)
  curvature_attack.py  worst-case: aim the sub-threshold bias at the flattest loss direction
  backdoor.py      targeted backdoor: can a cheat evade per-step AND loss detection at once?
  backdoor_capacity.py  does over-parameterization open a stealthy backdoor? (yes)
  scale.py         re-run §8/§9 on a deep model + AdamW (do the findings survive?)
  grinding.py      expected work to grind an evading Fiat-Shamir probe vs cheat size
  zk_matmul.py     ZK sumcheck proof for one GEMM + honest cost vs recompute/Freivalds
  figures.py       regenerate every figure in figures/
tests/             15 pytest files (88 tests) across every module
docs/DESIGN.md     12-section design document; REPORT.md     research writeup with figures
figures/           generated plots; Makefile     test / lint / figures / experiments targets

Quickstart

git clone https://github.com/oscartiz/freivalds-pol && cd freivalds-pol
python -m venv .venv && source .venv/bin/activate
pip install -e ".[dev,viz]"

(For a bit-exact copy of the environment the release was verified with, use pip install -e . -r requirements-dev.txt instead of the extras.)

Verify a real training step in a few lines:

import numpy as np

from freivalds_pol import MLP, make_task, sample_challenges, step_transcript, verify_step
from freivalds_pol.adversary import wrong_compute

rng = np.random.default_rng(0)

# A node does one real training step (two-layer MLP -> five GEMMs) ...
X, Y = make_task(32, 64, 16, batch=8, rng=rng)
mlp = MLP.init(32, 64, 16, rng)
transcript, loss, grads = step_transcript(mlp, X, Y)

# ... and publishes a Merkle commitment before the challenge window opens.
commitment = transcript.commitment()

# The verifier audits 3 random GEMMs at O(n^2) each -- no recompute.
challenges = sample_challenges(len(transcript.matmuls), k=3, rng=rng)
honest = verify_step(transcript, challenges, commitment,
                     expected_shard_root=transcript.shard_root)
print(honest.accepted, honest.reason)   # True ok

# A node that corrupted its outputs is caught by the calibrated Freivalds check.
cheat = wrong_compute(transcript, rng=rng, frac=1.0)
caught = verify_step(cheat, challenges, cheat.commitment(),
                     expected_shard_root=transcript.shard_root)
print(caught.accepted, caught.reason)   # False freivalds(right) failed on 'fwd.Z1'

Development targets:

make test            # 88 tests
make lint            # ruff check + format check, clean
make figures         # regenerate figures/*.png
make experiments     # run all 12 experiment scripts

# or individually:
python -m experiments.run_detection      # detection rate / verifier cost
python -m experiments.fp_crux            # precision vs. noise floor & undetectable cheat
python -m experiments.adaptive           # predictable vs. commit-then-sample probe
python -m experiments.real_step          # full verifier on a real transformer-block step
python -m experiments.compressed         # verify the DeMo-compressed update (DisTrO format)
python -m experiments.multiround         # do sub-threshold cheats accumulate over a run?
python -m experiments.curvature_attack   # worst case: bias along the flattest loss direction
python -m experiments.backdoor           # targeted backdoor: stealth-vs-harm tradeoff
python -m experiments.backdoor_capacity  # over-parameterization widens the stealth window
python -m experiments.scale              # deep model + AdamW: do §8/§9 findings survive?
python -m experiments.grinding           # grinding cost vs cheat magnitude and rounds
python -m experiments.zk_matmul          # ZK sumcheck proof for one GEMM + honest costs

API overview

The protocol surface (everything importable from freivalds_pol):

• Prove side — MLP / TransformerBlock run a real, grad-checked training step; step_transcript(...) packages its GEMMs, data-shard binding, and submitted update as a StepTranscript; transcript.commitment() is the Merkle root the node publishes.
• Challenge side — sample_challenges(num_matmuls, k) picks which GEMMs to audit; fiat_shamir_probes(commitment, ...) derives the probe vectors from the commitment (plus a public beacon), so a node cannot aim a cheat at a probe it has already seen.
• Verify side — verify_step(transcript, challenges, expected_commitment, ...) runs the whole protocol and returns a VerifyResult (accepted, first failing reason, checks run). The pieces are exposed individually: freivalds_check_threshold / freivalds_residual[_with] (the O(n²) product check), calibrated_threshold / honest_bound_inf (the floating-point noise floor it is compared against), and merkle_proof / verify_merkle_proof (single-leaf openings).
• Compression layer — compressor.compress/verify_compressed (1D-tiled) and demo.encode/verify (faithful 2D-chunk DeMo) check the DisTrO wire format per tile/block.
• Attack toolkit — adversary (transcript cheats), adaptive.nullspace_cheat/fixed_cheat, trainer (multi-round budget-constrained adversaries, backdoors), curvature (flattest- direction attacks), collusion.detect_free_riders.
• ZK path — zk.prove / zk.verify_full: a non-interactive sumcheck argument for one matmul over a prime field (see limitations below).

Guarantees and limitations

Complexity. One Freivalds round on an (m×k)·(k×n) product costs O(mk + kn + mn) — three matrix–vector products — versus O(mkn) to recompute it. The calibrated acceptance threshold is also O(n²) (numerics.honest_bound_inf), so the whole per-challenge check stays quadratic.

Detection probability. Over exact arithmetic, a wrong product survives one random Rademacher probe with probability ≤ 1/2, so rounds probes drive escape below 2^-rounds. A node that corrupts a fraction f of its recorded matmuls survives a k-challenge audit with probability ≈ (1−f)^k. In floating point the guarantee is threshold-relative: cheats whose per-probe residual exceeds the calibrated honest-noise bound are caught; cheats below the noise floor are undetectable per-step by construction — the multi-round experiments (§8–§9b of docs/DESIGN.md) study exactly what such sub-threshold cheats can and cannot achieve over a full run.

Soundness conditions. The FP threshold is rigorous only when the challenged layer is computed at ≥ fp32 (the verifier enforces min_dtype); probes must be derived after the commitment from unpredictable randomness (Fiat-Shamir over commitment + fresh public beacon), and the two-sided (right + left) check is required to close the rank-1 adaptive edge.

Limitations. This is a research prototype, not a production verifier:

• Pure numpy at toy scale (single blocks up to 4-layer/8-head, d=64); no GPU kernels, no networking or consensus layer, and cost ratios are measured in FLOP counts, not wall-clock on heterogeneous fleets.
• The ZK path proves the sumcheck reduction end to end, but its polynomial commitment is an intentionally insecure reveal-based reference (zk.RevealCommitment) — sound, but neither hiding nor succinct. A real KZG/FRI PCS is future work.
• The security analysis is empirical beyond the stated bounds: adaptive/backdoor findings are measured on the models above and may not transfer to other architectures or objectives.
• No third-party security audit.

Status

Working prototype, end-to-end on a real (gradient-checked) transformer-block step — a Llama-style pre-norm block (RMSNorm, single-head causal attention, GELU MLP), not just synthetic matmuls. verify_step runs the full protocol over all 8 GEMMs of the step (including the data-dependent QKᵀ and PV): commitment + shard-root checks, Fiat-Shamir (commitment-derived) probes, a two-sided Freivalds check against the calibrated FP threshold, and a minimum-precision gate. Honest steps pass; lazy / wrong-compute / poison-shard / post-commitment-edit / below-precision cheats are rejected. And the DisTrO wire format — the DeMo-compressed update (momentum + per-tile DCT + top-k + error feedback) — is verified per tile on that real gradient. A multi-round DeMo trainer then settles the hardest question (do never-detected sub-threshold cheats accumulate? — no, not linearly; a worst-case curvature-targeted adversary gains no edge; and a targeted backdoor has no stealthy-and-effective regime at toy scale — but with AdamW + depth the backdoor becomes loss-stealthy, so loss monitoring alone fails and per-step verification is necessary; §8 findings hold at scale). 88 tests pass.

Findings folded in:

• FP crux — a rigorous honest-noise bound computable in O(n²), usable only at ≥ fp32 on the challenged layer; a statistical threshold whose smallest detectable cheat scales with the unit roundoff. (experiments/fp_crux.py, docs/DESIGN.md §5.)
• Adaptive adversary — a predictable probe is a total break (rank-1 nullspace cheat, unbounded impact); fixed by commit-then-sample; under fresh probes adaptivity buys nothing beyond the FP band and grinding is infeasible. The residual rank-1 edge is closed by two-sided probing (one-sided 0.000 → two-sided 1.000). (experiments/adaptive.py, §6.)
• Compressed update — the DeMo wire format decomposes into an elementwise momentum step, a linear DCT (Freivalds-checkable), and a top-k recompute, all checkable per tile; each compression cheat is caught with detection 1-(1-f)^c. (experiments/compressed.py, §7.)
• Multi-round (the hard one) — over a full run, a never-detected sub-threshold cheat does not accumulate linearly: drift grows sublinearly (p≈0.27 vs naive p=1) and the loss is barely moved, because the optimizer's restoring force bounds it. A directed bias has a real (but still sublinear) edge over random noise. (experiments/multiround.py, §8.)
• Worst-case curvature attack — no edge — aiming the bias at the Hessian's flattest direction (via Hessian-vector products + power iteration), even re-tracked as it moves, gives no drift or test-loss advantage over random. Flat directions are flat because the loss ignores them. (experiments/curvature_attack.py, §8.) Both this and the sublinear-drift finding hold at scale (4-layer/8-head + AdamW, experiments/scale.py).
• Targeted backdoor — loss-stealthy at scale (the M1 finding) — with the toy SGD step a backdoor wrecks the loss before it implants (no stealthy regime). But with AdamW + depth it becomes loss-stealthy: ~98% implanted at <1.1× test loss, so loss monitoring misses it — overturning the earlier single-block claim. Every effective budget is still ≫ the per-step Freivalds floor, so per-step verification catches it: it is necessary, not optional. (experiments/backdoor.py, experiments/scale.py, §9 / §9b.)

Next: nanoGPT scale + a language objective; fuse the gradient + compression checks over a committed accumulator chain across rounds; reproduce the AdamW backdoor on a richer objective.

Prior art

• zkFL — gradient aggregation for federated learning · https://arxiv.org/pdf/2502.18535
• A Survey of ZK-Proof-Based Verifiable ML · https://arxiv.org/html/2502.18535v2
• VeriLLM — publicly verifiable decentralized inference · https://arxiv.org/pdf/2509.24257
• Psyche network architecture · https://nousresearch.com/nous-psyche