Documentation

LeanMlir.Proofs.Architectures.Residual

Residual Connections — Gradient Accumulation #

The first chapter where backprop has to accumulate gradients from multiple paths into the same input. So far every layer has been a straight-line composition (chain rule), but residual blocks introduce fan-out: one input feeds two paths whose outputs are added.

The math is trivial — it's the pattern that matters. Once you see "two backwards add", you'll see it everywhere: residuals, attention, SE blocks, multi-head outputs, anywhere a tensor is consumed by more than one downstream op.

This file builds on the proved foundations in Tensor.lean:

With those in hand, the residual definitions are one-liners — no sorry's.

  1. Defines residual f x = f x + x via biPath f id and its VJP.
  2. Defines residualProj proj f x = proj x + f x via biPath proj f and its VJP.
  3. Comments on how this matches the ResNet skip connection in the MLIR (MlirCodegen.lean residual block emission).

References #

source
noncomputable def Proofs.residual {n : ℕ} (f : Vec n → Vec n) :
Vec n → Vec n

A basic residual block: output = sub-network output + identity.

residual f x = f(x) + x

The "skip connection" lets gradients flow directly from output back to input without going through f. This is why ResNets train: even if f has near-zero gradients (vanishing), the identity path keeps the signal alive.

Equations
Instances For
    source
    noncomputable def Proofs.residualHasVJP {n : ℕ} (f : Vec n → Vec n) (hf_diff : Differentiable ℝ f) (hf : HasVJP f) :
    HasVJP (residual f)

    Residual VJP: dx = f.back(x, dy) + dy.

    The skip's contribution is just dy (identity backward). The block's contribution is f.back(x, dy). They add. The skip passes dy to the input unchanged, so the gradient does not have to go through f's Jacobian to reach x (the two terms can still cancel; no lower bound on dx is claimed).

    MLIR (MlirCodegen.lean, the residual-block case of the backward walk): The "skip grad" is added to the first convBn of the block — exactly f.back(x, dy) + dy_skip, where dy_skip = dy here.

    Proof: immediate from biPathHasVJP and identityHasVJP, both proved in Tensor.lean.

    Equations
    Instances For
      source
      noncomputable def Proofs.residualProj {m n : ℕ} (proj f : Vec m → Vec n) :
      Vec m → Vec n

      Projected residual block: when input and output shapes don't match (e.g. when stride > 1 downsamples), the skip is not identity but a 1×1 projection conv.

      residualProj proj f x = proj(x) + f(x)

      Both paths now have nontrivial backwards. The gradient still adds at the input — neither path is privileged.

      Equations
      Instances For
        source
        noncomputable def Proofs.residualProjHasVJP {m n : ℕ} (proj f : Vec m → Vec n) (hproj_diff : Differentiable ℝ proj) (hf_diff : Differentiable ℝ f) (hproj : HasVJP proj) (hf : HasVJP f) :
        HasVJP (residualProj proj f)

        Projected residual VJP: dx = proj.back(x, dy) + f.back(x, dy).

        Both backwards run on the same dy and their results sum at x. This is the truly general "fan-out → backward fan-in" pattern.

        MLIR: ResNets with stride > 1 use this — see emitConvBnBackward where the projection's VJP is emitted alongside the main block's, and both gradients accumulate into the same incoming-grad SSA.

        Proof: immediate from biPathHasVJP, proved in Tensor.lean.

        Equations
        Instances For

          Why this matters beyond ResNets #

          The fan-out/backward-add pattern is the structural building block for every modern architecture:

          • ResNets — y = f(x) + x (this file). • DenseNets — y = concat(f(x), x). The concat splits dy and each half goes back through its respective path. Same pattern, different glue (split instead of add). • Squeeze-and-Excitation — y = x · gate(x). The product rule introduces a different kind of bi-path: the gate's gradient gets x ⊙ dy and the main path's gradient gets gate(x) ⊙ dy. See SE.lean for that derivation. • Multi-head attention — concatenated heads. Same structure. • Two-tower models — independent encoders → joint loss. Even more extreme fan-out.

          If you understand biPathHasVJP, you understand backprop through any DAG. Composition (chain rule) handles the "linear" part; bi-path handles the joins. Together they're enough for any computation graph.

          source
          theorem Proofs.residualHasVJP_correct {n : ℕ} (f : Vec n → Vec n) (hf_diff : Differentiable ℝ f) (hf : HasVJP f) (x dy : Vec n) (i : Fin n) :
          (residualHasVJP f hf_diff hf).backward x dy i = ∑ j : Fin n, pdiv (residual f) x i j * dy j

          Public correctness theorem for residualHasVJP: skip-connection backward equals the pdiv-contracted Jacobian of f + id.

          source
          theorem Proofs.residualProjHasVJP_correct {m n : ℕ} (proj f : Vec m → Vec n) (hproj_diff : Differentiable ℝ proj) (hf_diff : Differentiable ℝ f) (hproj : HasVJP proj) (hf : HasVJP f) (x : Vec m) (dy : Vec n) (i : Fin m) :
          (residualProjHasVJP proj f hproj_diff hf_diff hproj hf).backward x dy i = ∑ j : Fin n, pdiv (residualProj proj f) x i j * dy j

          Public correctness theorem for residualProjHasVJP: same as residualHasVJP_correct but for the projected variant where the skip isn't identity.

          ResNet residual bodies contain ReLU, which is only DifferentiableAt at smooth points (see MLP.lean), not Differentiable globally. So the end-to-end CNN VJP (cnnHasVJPAt, in CNN.lean) chains through the pointwise HasVJPAt framework — vjpCompAt + the witnesses below — exactly as mlpHasVJPAt does for the MLP. These mirror the everywhere versions above, at a fixed x; pdiv_add is already stated at-point so the proof is the everywhere one with intro x dropped.

          source
          noncomputable def Proofs.biPathHasVJPAt {m n : ℕ} (f g : Vec m → Vec n) (x : Vec m) (hf_diff : DifferentiableAt ℝ f x) (hg_diff : DifferentiableAt ℝ g x) (hf : HasVJPAt f x) (hg : HasVJPAt g x) :
          HasVJPAt (biPath f g) x

          Additive fan-in at a point — smooth-point analog of biPathHasVJP.

          Equations
          Instances For
            source
            noncomputable def Proofs.residualHasVJPAt {n : ℕ} (f : Vec n → Vec n) (x : Vec n) (hf_diff : DifferentiableAt ℝ f x) (hf : HasVJPAt f x) :
            HasVJPAt (residual f) x

            Residual VJP at a point: dx = f.back(dy) + dy. The skip (identity) is differentiable everywhere, so only f needs the smooth-point hypothesis.

            Equations
            Instances For
              source
              theorem Proofs.residual_apply {n : ℕ} (f : Vec n → Vec n) (v : Vec n) (k : Fin n) :
              residual f v k = f v k + v k
              source
              theorem Proofs.residualProj_apply {m n : ℕ} (proj f : Vec m → Vec n) (v : Vec m) (k : Fin n) :
              residualProj proj f v k = proj v k + f v k
              source
              theorem Proofs.residual_differentiableAt {n : ℕ} {f : Vec n → Vec n} {x : Vec n} (hf : DifferentiableAt ℝ f x) :
              DifferentiableAt ℝ (residual f) x
              source
              noncomputable def Proofs.residualProjHasVJPAt {m n : ℕ} (proj f : Vec m → Vec n) (x : Vec m) (hproj_diff : DifferentiableAt ℝ proj x) (hf_diff : DifferentiableAt ℝ f x) (hproj : HasVJPAt proj x) (hf : HasVJPAt f x) :
              HasVJPAt (residualProj proj f) x

              Projected residual VJP at a point: both paths carry smooth-point hypotheses (the 1×1 stride-2 projection is linear, but stated at-point for uniformity with the composition).

              Equations
              Instances For
                source
                theorem Proofs.residual_continuous {n : ℕ} (F : Vec n → Vec n) (hF : Continuous F) :
                Continuous (residual F)

                A residual branch is continuous when its body is.

                source
                theorem Proofs.residualProj_continuous {m n : ℕ} (P F : Vec m → Vec n) (hP : Continuous P) (hF : Continuous F) :
                Continuous (residualProj P F)

                A projected residual is continuous when both branches are.