Proving inference · Operator atlas · Model transforms & fusion

Model transforms & fusion

the graph rewrite before proving

$$\text{a transform is one method: } \mathrm{apply}(\text{Model}) \to \text{Model} \\[8pt] \text{GPT2RMSNorm: } \mathrm{LayerNorm} \to \mathrm{RMSNorm},\ \text{pushing mean-subtraction and } \gamma,\beta\ \text{into neighbours} \\[6pt] \text{numerics preserved (tests assert pre/post traces are close)}$$

The algorithm

Before proving, DeepProve runs a graph-rewrite pass over the model — the concrete mechanism behind the "operator fusion" our docs only mentioned in the abstract. Rewrites fold constants into neighbouring layers so an expensive operator becomes a cheaper one, without changing the numbers.

Why it resists a proof

This page also corrects a framing error: DeepProve is not LLM-only. The Layer enum carries CNN/MLP operators (Convolution, Pooling, Activation, Flatten) beside the transformer ops, all behind the same Evaluate/ProveInfo/ProvableOp traits, so the transform, quantization and proving passes are uniform across architectures. LLM models are loaded from GGUF and safetensors, not only ONNX.

Cryptographic toolbox

  • ModelTransform trait + apply_transformations driver
  • GPT2RMSNorm: LayerNorm → RMSNorm, absorb γ/β into neighbours
  • Gemma3Transform: absorb the RMSNorm α scale
  • SoftmaxMaskTransform: set the quantized −∞ sentinel

How each scheme proves it

DeepProve

Three named graph rewrites fold arithmetic into neighbours: LayerNorm becomes RMSNorm, RMSNorm's scale is absorbed, and the attention mask's quantized −∞ is fixed up post-quantization.

  • GPT2RMSNorm replaces every LayerNorm with a cheaper RMSNorm. It right-multiplies the preceding matrix by a "mean-subtraction" matrix so every output row already has mean 0, and absorbs $\gamma,\beta$ into the following EinSum. It pattern-matches the GPT-2 graph shape (and detects the final projection by the magic number vocab == 50257).
  • Gemma3Transform absorbs an RMSNorm's $\alpha$ scale into the next layer, inserting a new scaling EinSum node when the successor can't absorb it. SoftmaxMaskTransform is a post-quantization fixup that writes the correct quantized $-\infty$ sentinel onto the attention mask feeding a softmax.
  • Correctness is guarded by trace-equivalence tests: each rewrite runs the model before and after and asserts the inputs to the final logits node match within tolerance; the GPT-2 rewrite is additionally proved end-to-end.
zkml/src/model/transform/mod.rs · 12-33
pub trait ModelTransform<N>: Debug
where
    N: TensorTypeParam,
{
    /// This applies the transformation returning a new model with the changes applied.
    fn apply(&self, model: Model<N>) -> Result<Model<N>>;
}

pub fn apply_transformations<N>(
    mut model: Model<N>,
    transforms: Vec<Box<dyn ModelTransform<N>>>,
) -> Result<Model<N>>
where
    N: TensorTypeParam,
{
    for transform in transforms {
        model = transform.apply(model)?;
    }
    Ok(model)
}
A transform is one method apply(Model)→Model; a driver folds a list of them over the model in order.
zkml/src/layers/mod.rs · 78-98
Convolution(Convolution<T>),
Activation(Activation<T>),
Requant(Requant),
Pooling(Pooling),
Flatten(Flatten),
EinSum(EinSum<T>),
LayerNorm(LayerNorm<T>),
RMSNorm(RMSNorm<T>),
Softmax(Softmax<T>),
Add(Add<T>),
Reshape(Reshape),
Embeddings(Embeddings<T>),
Positional(Positional<T>),
AttentionMask(AttentionMask<T>),
Logits(Logits),
Split(SplitLayer),
Recombination(RecombinationLayer),
The model-generic Layer enum: CNN/MLP ops (Convolution, Pooling, Flatten) sit beside the transformer ops.
← previousClaims & the backward pass