LayerNorm

GPT-2

$$Y[i,j] = \frac{X[i,j]-\mu_i}{\sqrt{\frac{1}{d}\sum_j (X[i,j]-\mu_i)^2 + \epsilon}} \\[12pt] \text{paper's identity:}\quad \mathrm{LayerNorm}(X[i,:]) = \mathrm{RMSNorm}\big(X[i,:] - \mu_i\big)\ \text{— but the code does }\textbf{not}\text{ reduce this way}$$

The algorithm

Centre each row, then rescale to unit variance.

Why it resists a proof

The same square root and reciprocal as RMSNorm, plus a mean that is itself a division by $d$. LayerNorm is also a documented source of the outlier activations that break quantization elsewhere in the model.

Cryptographic toolbox

  • A dedicated LayerNorm layer (not a reduction)
  • Committed mean + std-dev witnesses
  • Multiply by std, subtract mean, shift, clamping lookup
  • Keeps ε, unlike RMSNorm

How each scheme proves it

The paper reduces LayerNorm to RMSNorm(X−μ); the code does not. LayerNorm is a fully separate layer with its own proof — it commits both mean and std-dev witnesses, keeps ε, and calls no RMSNorm code.

  • Correction: our old page said "reduce to the RMSNorm protocol on X−μ." The implementation is a standalone LayerNorm layer with its own LayerNormProvingData and LayerNormLookupVerifier. It never invokes RMSNorm.
  • It commits two witnesses (the per-row mean and std-dev), and its apply multiplies by the std-dev, subtracts the committed mean, then shifts and looks up — a clamping table lookup, not a bare rescale.
  • Two things distinguish it from RMSNorm in code: it adds a sum/mean well-formedness check that RMSNorm lacks, and — unlike RMSNorm — it keeps ε (.add_scalar(self.eps)).
zkml/src/layers/transformer/normalisation/layernorm/mod.rs · 307-319
    fn apply(
        &self,
        input: WrappedTensor<Element>,
        table: &Table,
    ) -> Result<WrappedTensor<Element>> {
        let rescaled_input = input.mul(self.std_dev.clone())?;
        let zero_mean_input = rescaled_input.sub(self.mean.clone())?;

        let to_lookup = zero_mean_input
            .add_scalar(self.rounding_constant())
            .bitwise_right_shift_scalar(self.right_shift() as Element);
        table.lookup_tensor(to_lookup)
    }
LayerNorm's own apply: multiply by std-dev, subtract the committed mean, shift, then a clamping table lookup — it does not call RMSNorm.

Assembled from ONNX primitives: ReduceMean is decomposed internally into ReduceSum + Div.

Reduction Operations: ReduceSum along arbitrary axes; ReduceMean (decomposed internally into ReduceSum + Div). Math Functions: Rsqrt (reciprocal square root).

Jolt Atlas · §5.2
  • ReduceMean is explicitly "decomposed internally into ReduceSum + Div", and ReduceSum works along arbitrary axes.
  • The Div carries a non-deterministic advice value; the Rsqrt is a lookup. Composition gives LayerNorm with no bespoke protocol.
  • Elegant in the graph-compiler sense — you get LayerNorm for free once the primitives exist — but it inherits exactly the accuracy concern DeepProve raises about quantizing inverses and square roots.
  • In code it bottoms out in the same fused MeanOfSquares + Rsqrt chain, so the same "honesty rests on the remainder range checks" caveat applies. Note the repo ships only GPT-2 / nanoGPT models, so the norm operators are proven-out by unit tests but not exercised by a real RMSNorm/LayerNorm LLM end-to-end.

Normalization is a named target, covering batch and layer norm together: $y_i = \gamma\,(x_i - \mu)/\sigma + \beta$, with the reciprocal square root discharged as a first-class lookup-based building block.

  • The protocol is built from the reciprocal-square-root building block, which -- like division -- is a digit-decomposed table lookup refined by a fixed number of iterations ($I = 1$ by default, at a 6-bit lookup).
  • This is the 'standard approach' DeepProve explicitly rejects: quantize the inverse square root and eat a rounding error per row, per layer. DeepProve's objection is that these errors accumulate through the variance sum and crater accuracy. Hao et al. take the rejected path and report no accuracy measurement, so the objection is neither confirmed nor answered.
  • It also directly attacks Mystique's worst operator: Batch Normalization consumed the majority of Mystique's ResNet-101 prover precisely because of the float/fixed conversions this paper deletes.
  • The reported shapes are small -- normalization over dimension 16, a CIFAR-class shape. An LLM's LayerNorm runs over the model dimension, and the digit-decomposition cost model is sensitive to input range, so these figures should not be extrapolated to a transformer.