Documentation

LeanMlir.Proofs.Float.ConvMixedFloatBridge

Mixed-precision convolution forward error #

The accuracy half of the bf16 conv ops in Proofs.StableHLO (BatchableOp.convBf16, flatConvFBf16, and the dgrad/wgrad twins).

FloatBridge stops at dense: it has dot_close_mixed / dense_close_mixed and no conv peer, and every net where bf16 pays is conv-dominated.

It is not a new hard theorem. A convolution output is a dot product over its flattened receptive field (conv2d_eq_flat_dot), so the whole thing is dot_close_mixed_uniform instantiated at fan-in ic·kH·kW, plus one leaf rounding for the bf16 STORE and one accumulate rounding for the bias add.

Note: The store term is the part that has no analogue in dense_close_mixed, and it is forced by the hardware rather than chosen: a conv with bf16 operands and an f32-typed result has its casts deleted by XLA and runs entirely in fp32 (measured — see BatchableOp.convBf16). The only emit shape that reaches the tensor cores gives the convolution a bf16-typed result, so the accumulator is rounded on store and the error model has to say so.

Non-vacuous, and the fan-in is not what costs. convBr at fp32 accumulate (M.u = 2⁻²⁴) and bf16 leaf (L.u = 2⁻⁸), evaluated at ResNet-34's own layers — arithmetic outside Lean, quoted as illustration, not proved here:

layer                  fan-in n   fan-in term   leaf term   convBr
stem 7×7, ic=3              147     8.9e-06     7.83e-03    0.0078
stage-1 3×3, ic=64          576     3.5e-05     7.83e-03    0.0079
stage-4 3×3, ic=512        4608     2.8e-04     7.83e-03    0.0081
1×1 projection, ic=512      512     3.1e-05     7.83e-03    0.0079

Under 1% everywhere, and the LEAF term dominates the fan-in term by ~28× even at the widest layer. Set the accumulate to bf16 as well and ((1+u)^(n+1) − 1) at n=4608 is 6.4e7 — a bound that says nothing at all. That contrast is the whole argument for bf16-mixed over bf16: the 1/u fan-in wall sits at the ACCUMULATE precision, which stays fp32.

source
theorem Proofs.Tensor3.sum_flatten {c h w : ℕ} (T : Tensor3 c h w) :
∑ k : Fin (c * h * w), T.flatten k = ∑ i : Fin c, ∑ j : Fin h, ∑ l : Fin w, T i j l

Summing a flattened Tensor3 is the triple sum.

source
theorem Proofs.Tensor3.flatten_mul {c h w : ℕ} (A B : Tensor3 c h w) (k : Fin (c * h * w)) :
A.flatten k * B.flatten k = flatten (fun (i : Fin c) (j : Fin h) (l : Fin w) => A i j l * B i j l) k

Flattening commutes with a pointwise product: both sides look up the same index.

source
noncomputable def Proofs.convWindow3 {ic h w : ℕ} (kH kW : ℕ) (x : Tensor3 ic h w) (hi : Fin h) (wi : Fin w) :
Tensor3 ic kH kW

The kH × kW receptive field conv2d reads at output pixel (hi, wi), zero outside: convPad (ConvIndex) as a Tensor3, so the conv's fan-in can be summed with Tensor3.sum_flatten (conv2d_eq_flat_dot). convWindow (ConvFloat) is the same read at the flat type Vec (ic*kH*kW).

Equations
Instances For
    source
    noncomputable def Proofs.convSlice {ic oc kH kW : ℕ} (W : Kernel4 oc ic kH kW) (o : Fin oc) :
    Tensor3 ic kH kW

    The output channel's kernel slice, as a Tensor3.

    Equations
    Instances For
      source
      theorem Proofs.conv2d_eq_flat_dot {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (x : Tensor3 ic h w) (o : Fin oc) (hi : Fin h) (wi : Fin w) :
      conv2d W b x o hi wi = b o + ∑ k : Fin (ic * kH * kW), (convSlice W o).flatten k * (convWindow3 kH kW x hi wi).flatten k

      A convolution output is a DOT PRODUCT of length ic·kH·kW over the flattened receptive field. This is the whole reason conv_close_mixed is not a new hard theorem: it lets the conv reuse dot_close_mixed_uniform at that fan-in.

      source
      noncomputable def Proofs.convBrR (uacc uleaf : ℝ) (n : ℕ) :
      ℝ

      The Higham-style bracket dot_close_mixed_uniform produces at fan-in n, as a function of the two roundoffs alone (so a concrete instance evaluates by norm_num): fan-in amplification rides the ACCUMULATE roundoff uacc, the leaf roundoff contributes a flat per-leaf term.

      Equations
      Instances For
        source
        theorem Proofs.convBrR_nonneg {uacc uleaf : ℝ} (hacc : 0 ≤ uacc) (hleaf : 0 ≤ uleaf) (n : ℕ) :
        0 ≤ convBrR uacc uleaf n
        source
        noncomputable def Proofs.convBr (M L : FloatModel) (n : ℕ) :
        ℝ

        The bracket at an accumulate model M and a leaf model L.

        Equations
        Instances For
          source
          theorem Proofs.convBr_eq_convBrR (M L : FloatModel) (n : ℕ) :
          convBr M L n = convBrR M.u L.u n
          source
          noncomputable def Proofs.convFanS {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (x : Tensor3 ic h w) (o : Fin oc) (hi : Fin h) (wi : Fin w) :
          ℝ

          Σ|kernel·window| over the receptive field — the magnitude the bound scales.

          Equations
          Instances For
            source
            theorem Proofs.FloatModel.storeBias_close (M L : FloatModel) {n : ℕ} (x y : Vec n) (β : ℝ) :
            |M.add (L.rnd (M.dotMixed L x y)) β - (∑ k : Fin n, x k * y k + β)| ≤ M.u * ((1 + L.u) * (1 + convBr M L n) * ∑ k : Fin n, |x k * y k| + |β|) + L.u * (1 + convBr M L n) * ∑ k : Fin n, |x k * y k| + convBr M L n * ∑ k : Fin n, |x k * y k|

            The mixed-precision store-then-bias step, once. A dot accumulated at M over leaves rounded to L (dotMixed), rounded to L again on store, then a bias added at M: three terms, one per rounding — the dot (convBr, fan-in n), the store (L.u) and the bias add (M.u). conv_close_mixed and depthwise_close_mixed are this at their receptive fields.

            source
            noncomputable def Proofs.FloatModel.convMixed (M L : FloatModel) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (x : Tensor3 ic h w) :
            Tensor3 oc h w

            The mixed-precision convolution, as the emitted graph computes it. Operands rounded to the leaf precision L and accumulated at M (dotMixed), the accumulator then rounded to L again — the bf16-typed result, i.e. the store — and only then the bias added at M.

            Note: The second L.rnd is what distinguishes this from denseMixed, and it is not optional: BatchableOp.convBf16 must give the convolution a bf16-typed result or XLA deletes the casts and runs the whole conv in f32. The store is a consequence of the only emit shape that reaches the tensor cores, so the error model has to carry it.

            Equations
            Instances For
              source
              theorem Proofs.FloatModel.conv_close_mixed (M L : FloatModel) {ic oc h w kH kW : ℕ} (W : Kernel4 oc ic kH kW) (b : Vec oc) (x : Tensor3 ic h w) (o : Fin oc) (hi : Fin h) (wi : Fin w) :
              |M.convMixed L W b x o hi wi - conv2d W b x o hi wi| ≤ M.u * ((1 + L.u) * (1 + convBr M L (ic * kH * kW)) * convFanS W x o hi wi + |b o|) + L.u * (1 + convBr M L (ic * kH * kW)) * convFanS W x o hi wi + convBr M L (ic * kH * kW) * convFanS W x o hi wi

              Mixed-precision convolution forward error. Three terms, one per rounding the emitted graph performs: the dot (convBr, fan-in ic·kH·kW), the bf16 STORE of the accumulator (L.u), and the f32 bias add (M.u).

              It is dot_close_mixed_uniform instantiated at the conv's fan-in, because a convolution output IS a dot product over its flattened receptive field (conv2d_eq_flat_dot). The fan-in wall therefore still sits at 1/M.u = 2²⁴ and not at the leaf precision — the same reason bf16-mixed is non-vacuous for dense.