Skip to content

Commit 4c4b4dc

Browse files
perf: prove Add/Sub/Neg without a sumcheck (non-ZK path) (#252)
* perf: prove Add/Sub/Neg without a sumcheck on the non-ZK path (#174) Add/Sub/Neg are linear element-wise ops: their output MLE equals a fixed linear combination of the operand MLEs over the same hypercube. The identity therefore holds at the random output opening point r by Schwartz-Zippel, so the per-op log(n)-round sumcheck is redundant. Convert their OperatorProofTrait impls to the no-sumcheck opening-propagation form used by Identity: the prover opens the operands at the node's reduced output point r and returns no execution proof; the verifier checks the linear relation (left+right / left-right / -operand == output) directly. Operand openings are discharged by the existing batched HyperKZG opening. The ZK (BlindFold) path is unchanged and still proves these ops via its batched sumcheck, so the AddParams/AddProver/AddVerifier structs are retained for it; a follow-up can add a BlindFold algebraic-identity constraint to drop the sumcheck there too. The Sub soundness test is re-expressed to forge an operand opening directly (the new direct check rejects it). Verified: add/sub/neg unit tests, full jolt-atlas-core suite incl. nanoGPT e2e, soundness_tests, and the --features zk suite all pass; fmt + clippy (default and zk) clean. * Update jolt-atlas-core/src/onnx_proof/ops/eval_reduction.rs Co-authored-by: Antoine Douchet <antoinedcht@hotmail.com> --------- Co-authored-by: Antoine Douchet <antoinedcht@hotmail.com>
1 parent 21729b8 commit 4c4b4dc

8 files changed

Lines changed: 230 additions & 221 deletions

File tree

jolt-atlas-core/src/onnx_proof/malicious_prover.rs

Lines changed: 2 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -10,17 +10,12 @@ use atlas_onnx_tracer::{
1010
};
1111
use joltworks::{
1212
field::JoltField,
13-
poly::{
14-
commitment::commitment_scheme::CommitmentScheme,
15-
opening_proof::ProverOpeningAccumulator,
16-
unipoly::{CompressedUniPoly, UniPoly},
17-
},
13+
poly::{commitment::commitment_scheme::CommitmentScheme, unipoly::UniPoly},
1814
subprotocols::{
1915
evaluation_reduction::{EvalReductionProof, ReducedInstance},
2016
sumcheck::SumcheckInstanceProof,
21-
sumcheck_prover::SumcheckInstanceProver,
2217
},
23-
transcripts::{AppendToTranscript, Transcript},
18+
transcripts::Transcript,
2419
};
2520

2621
use crate::onnx_proof::{
@@ -183,55 +178,6 @@ impl MaliciousONNXProof {
183178
}
184179
}
185180

186-
/// Variant of `Sumcheck::prove` that also returns the final claim and leaves
187-
/// opening caching to the caller (for adversarial experiments).
188-
pub fn malicious_sumcheck_prove<F: JoltField, ProofTranscript: Transcript>(
189-
sumcheck_instance: &mut dyn SumcheckInstanceProver<F, ProofTranscript>,
190-
opening_accumulator: &mut ProverOpeningAccumulator<F>,
191-
transcript: &mut ProofTranscript,
192-
) -> (
193-
SumcheckInstanceProof<F, ProofTranscript>,
194-
Vec<F::Challenge>,
195-
F,
196-
) {
197-
let num_rounds = sumcheck_instance.num_rounds();
198-
199-
// Append input claims to transcript
200-
let input_claim = sumcheck_instance.input_claim(opening_accumulator);
201-
transcript.append_scalar(&input_claim);
202-
let mut previous_claim = input_claim;
203-
let mut r_sumcheck: Vec<F::Challenge> = Vec::with_capacity(num_rounds);
204-
let mut compressed_polys: Vec<CompressedUniPoly<F>> = Vec::with_capacity(num_rounds);
205-
for round in 0..num_rounds {
206-
let univariate_poly = sumcheck_instance.compute_message(round, previous_claim);
207-
// append the prover's message to the transcript
208-
let compressed_poly = univariate_poly.compress();
209-
compressed_poly.append_to_transcript(transcript);
210-
let r_j = transcript.challenge_scalar_optimized::<F>();
211-
r_sumcheck.push(r_j);
212-
213-
// Cache claim for this round
214-
previous_claim = univariate_poly.evaluate(&r_j);
215-
sumcheck_instance.ingest_challenge(r_j, round);
216-
compressed_polys.push(compressed_poly);
217-
}
218-
219-
let final_claim = previous_claim;
220-
221-
// Allow the sumcheck instance to perform any end-of-protocol work (e.g. flushing
222-
// delayed bindings) after the final challenge has been ingested and before we cache
223-
// openings.
224-
sumcheck_instance.finalize();
225-
226-
// Deliberately do not call `cache_openings` here. The caller controls how
227-
// openings and operand claims are cached for attack experiments.
228-
(
229-
SumcheckInstanceProof::new(compressed_polys),
230-
r_sumcheck,
231-
final_claim,
232-
)
233-
}
234-
235181
/// Simulate the evaluation reduction protocol, only we just consider the case where num_use = 1 for simplicity.
236182
pub fn malicious_eval_reduction_prove<F: JoltField, T: Transcript>(
237183
node: &ComputationNode,

jolt-atlas-core/src/onnx_proof/ops/add.rs

Lines changed: 67 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
use crate::{
2-
impl_standard_sumcheck_proof_api,
3-
onnx_proof::{ops::OperatorProofTrait, ProofId, ProofType, Prover, Verifier},
2+
onnx_proof::{ops::OperatorProofTrait, ProofId, Prover, Verifier},
43
utils::opening_access::{AccOpeningAccessor, Target},
54
};
65
use atlas_onnx_tracer::{
@@ -16,7 +15,9 @@ use joltworks::{
1615
field::{IntoOpening, JoltField},
1716
poly::{
1817
eq_poly::EqPolynomial,
19-
multilinear_polynomial::{BindingOrder, MultilinearPolynomial, PolynomialBinding},
18+
multilinear_polynomial::{
19+
BindingOrder, MultilinearPolynomial, PolynomialBinding, PolynomialEvaluation,
20+
},
2021
opening_proof::{
2122
OpeningAccumulator, OpeningPoint, ProverOpeningAccumulator, VerifierOpeningAccumulator,
2223
BIG_ENDIAN, LITTLE_ENDIAN,
@@ -33,7 +34,69 @@ use joltworks::{
3334
utils::errors::ProofVerifyError,
3435
};
3536

36-
impl_standard_sumcheck_proof_api!(Add, AddParams, AddProver, AddVerifier);
37+
/// No-sumcheck proof for element-wise addition.
38+
///
39+
/// `Add` is linear: `output(x) = left(x) + right(x)` for all `x` on the shared
40+
/// boolean hypercube, so the multilinear identity `output ≡ left + right` holds
41+
/// at the (random) output opening point `r` by Schwartz–Zippel. We therefore
42+
/// skip the sumcheck entirely (cf. [`super::identity`]): the prover opens both
43+
/// operands at the *same* `r` the output is opened at, and the verifier checks
44+
/// `left(r) + right(r) == output(r)` directly. The batched HyperKZG opening
45+
/// proof discharges the operand openings, exactly as for `Identity`.
46+
///
47+
/// The sumcheck structs below ([`AddParams`]/[`AddProver`]/[`AddVerifier`]) are
48+
/// retained because the ZK (BlindFold) path in `onnx_proof::zk` still proves
49+
/// `Add` via a batched sumcheck; converting that path requires a dedicated
50+
/// BlindFold algebraic-identity constraint and is tracked separately.
51+
impl<F: JoltField, T: Transcript> OperatorProofTrait<F, T> for Add {
52+
#[tracing::instrument(skip_all, name = "Add::prove")]
53+
fn prove(
54+
&self,
55+
node: &ComputationNode,
56+
prover: &mut Prover<F, T>,
57+
) -> Vec<(ProofId, SumcheckInstanceProof<F, T>)> {
58+
let (opening_point, _claim) =
59+
AccOpeningAccessor::new(&prover.accumulator, node).get_reduced_opening();
60+
61+
let LayerData { operands, .. } = Trace::layer_data(&prover.trace, node);
62+
let [left, right] = operands[..] else {
63+
panic!("Expected two operands for Add operation")
64+
};
65+
let left_claim =
66+
MultilinearPolynomial::from(left.padded_next_power_of_two()).evaluate(&opening_point.r);
67+
let right_claim = MultilinearPolynomial::from(right.padded_next_power_of_two())
68+
.evaluate(&opening_point.r);
69+
70+
let mut provider = AccOpeningAccessor::new(&mut prover.accumulator, node)
71+
.into_provider(&mut prover.transcript, opening_point);
72+
provider.append_nodeio(Target::Input(0), left_claim);
73+
provider.append_nodeio(Target::Input(1), right_claim);
74+
vec![]
75+
}
76+
77+
#[tracing::instrument(skip_all, name = "Add::verify")]
78+
fn verify(
79+
&self,
80+
node: &ComputationNode,
81+
verifier: &mut Verifier<'_, F, T>,
82+
) -> Result<(), ProofVerifyError> {
83+
let (opening_point, claim) =
84+
AccOpeningAccessor::new(&verifier.accumulator, node).get_reduced_opening();
85+
let mut provider = AccOpeningAccessor::new(&mut verifier.accumulator, node)
86+
.into_provider(&mut verifier.transcript, opening_point);
87+
provider.append_nodeio(Target::Input(0));
88+
provider.append_nodeio(Target::Input(1));
89+
let left_claim = provider.get_nodeio(Target::Input(0)).1;
90+
let right_claim = provider.get_nodeio(Target::Input(1)).1;
91+
92+
if left_claim + right_claim != claim {
93+
return Err(ProofVerifyError::InvalidOpeningProof(
94+
"Add output claim should equal the sum of operand claims".to_string(),
95+
));
96+
}
97+
Ok(())
98+
}
99+
}
37100

38101
/// Shared parameter block for the element-wise addition sumcheck proof.
39102
#[derive(Clone)]

jolt-atlas-core/src/onnx_proof/ops/eval_reduction.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -287,7 +287,6 @@ mod tests {
287287
&mut eval_reduction_proofs,
288288
);
289289

290-
assert!(!proofs.is_empty());
291290
assert!(eval_reduction_proofs.contains_key(&output_idx));
292291
}
293292

Lines changed: 32 additions & 141 deletions
Original file line numberDiff line numberDiff line change
@@ -1,162 +1,53 @@
11
//! Malicious variant of the Sub prover for soundness testing.
22
//!
3-
//! This module contains a Sub prover that forges virtual operand claims
4-
//! while keeping the underlying sumcheck honest. It is used exclusively
5-
//! by [`MaliciousONNXProof`] to test that the verifier correctly handles
6-
//! (and rejects) such attacks.
3+
//! The honest `Sub` proof is a no-sumcheck op (see [`crate::onnx_proof::ops::sub`]):
4+
//! it opens both operands at the node's reduced output point `r` and the verifier
5+
//! checks `left(r) - right(r) == output(r)` directly. This module forges the left
6+
//! operand opening (off by one) so that `left(r) - right(r) != output(r)`, exercising
7+
//! that the verifier's direct difference check rejects the attack. Used exclusively by
8+
//! [`MaliciousONNXProof`](crate::onnx_proof::malicious_prover).
79
810
use crate::{
9-
onnx_proof::{malicious_prover::malicious_sumcheck_prove, ProofId, ProofType, Prover},
11+
onnx_proof::{ProofId, Prover},
1012
utils::opening_access::{AccOpeningAccessor, Target},
1113
};
1214
use atlas_onnx_tracer::{
1315
model::trace::{LayerData, Trace},
1416
node::ComputationNode,
1517
};
1618
use joltworks::{
17-
field::{IntoOpening, JoltField},
18-
poly::{
19-
eq_poly::EqPolynomial,
20-
multilinear_polynomial::{BindingOrder, MultilinearPolynomial, PolynomialBinding},
21-
opening_proof::ProverOpeningAccumulator,
22-
split_eq_poly::GruenSplitEqPolynomial,
23-
unipoly::UniPoly,
24-
},
25-
subprotocols::{
26-
sumcheck::SumcheckInstanceProof, sumcheck_prover::SumcheckInstanceProver,
27-
sumcheck_verifier::SumcheckInstanceParams,
28-
},
19+
field::JoltField,
20+
poly::multilinear_polynomial::{MultilinearPolynomial, PolynomialEvaluation},
21+
subprotocols::sumcheck::SumcheckInstanceProof,
2922
transcripts::Transcript,
3023
};
3124

32-
use crate::onnx_proof::ops::sub::SubParams;
33-
3425
/// Run the malicious Sub prover for a single node.
3526
///
36-
/// Returns the proof entry suitable for insertion into the proof map.
27+
/// Mirrors the honest no-sumcheck `Sub::prove` but forges the left operand
28+
/// opening as `left(r) + 1`, leaving the right operand honest. The verifier's
29+
/// `left - right == output` check then necessarily fails. Returns `vec![]` (no
30+
/// execution proof), exactly like the honest op.
3731
pub fn malicious_sub_prove<F: JoltField, T: Transcript>(
3832
node: &ComputationNode,
3933
prover: &mut Prover<F, T>,
4034
) -> Vec<(ProofId, SumcheckInstanceProof<F, T>)> {
41-
let params = SubParams::new(node.clone(), &prover.accumulator);
42-
let mut prover_sumcheck = MaliciousSubProver::initialize(&prover.trace, params);
43-
44-
let (proof, r_sumcheck, final_claim) = malicious_sumcheck_prove(
45-
&mut prover_sumcheck,
46-
&mut prover.accumulator,
47-
&mut prover.transcript,
48-
);
49-
prover_sumcheck.final_claim = Some(final_claim);
50-
prover_sumcheck.cache_openings(&mut prover.accumulator, &mut prover.transcript, &r_sumcheck);
51-
vec![(ProofId(node.idx, ProofType::Execution), proof)]
52-
}
53-
54-
/// Malicious prover state for element-wise subtraction sumcheck protocol.
55-
///
56-
/// Identical to the honest SubProver in sumcheck computation, but forges
57-
/// operand claims in `cache_openings` to demonstrate the attack vector.
58-
struct MaliciousSubProver<F: JoltField> {
59-
params: SubParams<F>,
60-
eq_r_node_output: GruenSplitEqPolynomial<F>,
61-
left_operand: MultilinearPolynomial<F>,
62-
right_operand: MultilinearPolynomial<F>,
63-
final_claim: Option<F>,
64-
}
65-
66-
impl<F: JoltField> MaliciousSubProver<F> {
67-
/// Initialize the prover with trace data and parameters.
68-
fn initialize(trace: &Trace, params: SubParams<F>) -> Self {
69-
let eq_r_node_output =
70-
GruenSplitEqPolynomial::new(&params.r_node_output.r, BindingOrder::LowToHigh);
71-
let LayerData {
72-
operands,
73-
output: _,
74-
} = Trace::layer_data(trace, &params.computation_node);
75-
let [left_operand, right_operand] = operands[..] else {
76-
panic!("Expected two operands for Sub operation")
77-
};
78-
let left_operand = MultilinearPolynomial::from(left_operand.clone());
79-
let right_operand = MultilinearPolynomial::from(right_operand.clone());
80-
Self {
81-
params,
82-
eq_r_node_output,
83-
left_operand,
84-
right_operand,
85-
final_claim: None,
86-
}
87-
}
88-
}
89-
90-
impl<F: JoltField, T: Transcript> SumcheckInstanceProver<F, T> for MaliciousSubProver<F> {
91-
fn get_params(&self) -> &dyn SumcheckInstanceParams<F> {
92-
&self.params
93-
}
94-
95-
fn compute_message(&mut self, _round: usize, previous_claim: F) -> UniPoly<F> {
96-
let Self {
97-
eq_r_node_output,
98-
left_operand,
99-
right_operand,
100-
..
101-
} = self;
102-
let [q_constant] = eq_r_node_output.par_fold_out_in_unreduced::<9, 1>(&|g| {
103-
let lo0 = left_operand.get_bound_coeff(2 * g);
104-
let ro0 = right_operand.get_bound_coeff(2 * g);
105-
[lo0 - ro0]
106-
});
107-
eq_r_node_output.gruen_poly_deg_2(q_constant, previous_claim)
108-
}
109-
110-
fn ingest_challenge(&mut self, r_j: F::Challenge, _round: usize) {
111-
self.eq_r_node_output.bind(r_j);
112-
self.left_operand
113-
.bind_parallel(r_j, BindingOrder::LowToHigh);
114-
self.right_operand
115-
.bind_parallel(r_j, BindingOrder::LowToHigh);
116-
}
117-
118-
fn cache_openings(
119-
&self,
120-
accumulator: &mut ProverOpeningAccumulator<F>,
121-
transcript: &mut T,
122-
sumcheck_challenges: &[F::Challenge],
123-
) {
124-
let opening_point = self
125-
.params
126-
.normalize_opening_point(&sumcheck_challenges.into_opening());
127-
128-
// Malicious behavior: forge virtual operand claims while preserving the
129-
// same subtraction difference, so expected_output_claim remains unchanged.
130-
let left_claim = self.left_operand.final_claim();
131-
let right_claim = self.right_operand.final_claim();
132-
let final_claim = self
133-
.final_claim
134-
.expect("final_claim must be set before cache_openings");
135-
let r_node_output_prime = self
136-
.params
137-
.normalize_opening_point(&sumcheck_challenges.into_opening())
138-
.r;
139-
let eq_eval = EqPolynomial::mle(&self.params.r_node_output.r, &r_node_output_prime);
140-
141-
// Choose forged claims so that:
142-
// final_claim == eq_eval * (forged_left - forged_right)
143-
let forged_left = left_claim + F::one();
144-
let forged_right = if eq_eval.is_zero() {
145-
// If eq_eval == 0, the only valid final claim is 0.
146-
debug_assert!(final_claim.is_zero());
147-
right_claim
148-
} else {
149-
let inv = eq_eval
150-
.inverse()
151-
.expect("non-zero eq_eval must be invertible");
152-
forged_left - final_claim * inv
153-
};
154-
debug_assert_eq!(final_claim, eq_eval * (forged_left - forged_right));
155-
let mut provider = AccOpeningAccessor::new(accumulator, &self.params.computation_node)
156-
.into_provider(transcript, opening_point);
157-
158-
// Insert forged claims while keeping transcript/opener state consistent.
159-
provider.append_nodeio(Target::Input(0), forged_left);
160-
provider.append_nodeio(Target::Input(1), forged_right);
161-
}
35+
let (opening_point, _claim) =
36+
AccOpeningAccessor::new(&prover.accumulator, node).get_reduced_opening();
37+
38+
let LayerData { operands, .. } = Trace::layer_data(&prover.trace, node);
39+
let [left, right] = operands[..] else {
40+
panic!("Expected two operands for Sub operation")
41+
};
42+
let forged_left = MultilinearPolynomial::from(left.padded_next_power_of_two())
43+
.evaluate(&opening_point.r)
44+
+ F::one();
45+
let right_claim =
46+
MultilinearPolynomial::from(right.padded_next_power_of_two()).evaluate(&opening_point.r);
47+
48+
let mut provider = AccOpeningAccessor::new(&mut prover.accumulator, node)
49+
.into_provider(&mut prover.transcript, opening_point);
50+
provider.append_nodeio(Target::Input(0), forged_left);
51+
provider.append_nodeio(Target::Input(1), right_claim);
52+
vec![]
16253
}

jolt-atlas-core/src/onnx_proof/ops/mod.rs

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,6 @@ use atlas_onnx_tracer::{node::ComputationNode, ops::Operator};
109109
// Re-export handler types for convenient access
110110

111111
// Re-export operator param/prover/verifier types for the handler macro
112-
pub use add::{AddParams, AddProver, AddVerifier};
113112
use common::CommittedPoly;
114113
pub use cube::{CubeParams, CubeProver, CubeVerifier};
115114
pub use div::{DivParams, DivProver, DivVerifier};
@@ -122,11 +121,9 @@ use joltworks::{
122121
utils::errors::ProofVerifyError,
123122
};
124123
pub use mul::{MulParams, MulProver, MulVerifier};
125-
pub use neg::{NegParams, NegProver, NegVerifier};
126124
pub use rsqrt::{RsqrtParams, RsqrtProver, RsqrtVerifier};
127125
pub use square::{SquareParams, SquareProver, SquareVerifier};
128126
use std::collections::BTreeMap;
129-
pub use sub::{SubParams, SubProver, SubVerifier};
130127

131128
use crate::onnx_proof::{ProofId, Prover, Verifier};
132129

0 commit comments

Comments
 (0)