Documentation

LeanMlir.Proofs.Nets.EfficientNet.EfficientNetBackB0

EfficientNet-B0 block backward graphs — the residual and squeeze-excite fan-ins #

Backward StableHLO graphs for B0's MBConv blocks and the theorems that they denote the proven VJPs. B0 has two branching ops the dense-chain nets do not: the residual skip (additive fan-in) and the squeeze-excite gate (multiplicative fan-in). SHlo's backward constructors are unary, but both fan-ins are spelled with forward elementwise combinators: addV for the residual, and layerScaleF (Hadamard by a known activation) + addV for the SE gate.

Contents #

Each …_faithful states den (graph) = the certified backward.

source
def Proofs.StableHLO.seBlockBackGraph {n : ℕ} (gateBack : SHlo n) (gateVal dy : Vec n) :
SHlo n

Backward graph for an SE block x ↦ x ⊙ gate x, given a subgraph gateBack rendering the gate sub-network's input-cotangent at the SE-specific cotangent x ⊙ dy. The main (identity) path contributes gate x ⊙ dy, rendered as a Hadamard (layerScaleF) of the cotangent by the gate activation; addV sums the two paths. The renderable image of seBlockHasVJP's elemwiseProduct (bi-cotangent) backward.

Equations
Instances For
    source
    theorem Proofs.StableHLO.seBlockBackGraph_faithful {n : ℕ} (gate : Vec n → Vec n) (hg_diff : Differentiable ℝ gate) (hg : HasVJP gate) (x dy : Vec n) (gateBack : SHlo n) (hgb : den gateBack = hg.backward x fun (j : Fin n) => x j * dy j) :
    den (seBlockBackGraph gateBack (gate x) dy) = (seBlockHasVJP gate hg_diff hg).backward x dy

    SE multiplicative fan-in backward faithfulness (general). If gateBack denotes the gate path's VJP backward at the cotangent x ⊙ dy (den gateBack = hg.backward x (x ⊙ dy)), then the SE backward graph denotes the proven seBlockHasVJP backward, which is gate x ⊙ dy + gate.backward x (x ⊙ dy). Like the residual brick the proof is structural — the composition is delegated to gateBack, so the only definitional facts are layerScaleF/addV denotation and the identity main-path backward.

    source
    theorem Proofs.StableHLO.bnBack_faithful_fn {n : ℕ} (gN xN es : String) (ε γ β : ℝ) (hε : 0 < ε) (x : Vec n) (e : SHlo n) :
    den (SHlo.bnBack gN xN es ε γ x e) = (bnHasVJP n ε γ β hε).backward x (den e)

    Function-level BatchNorm backward bridge. bnBack denotes bnGradInput, which is NOT rfl-equal to (bnHasVJP …).backward (the witness is built via a rw [bnForward_eq_compose] cast). They agree through the canonical VJP sum: bnBack_faithful gives the ∑ pdiv form and bnHasVJP.correct matches it. This lemma is the one non-rfl bridge the bn-containing stages need.

    source
    noncomputable def Proofs.StableHLO.mbBodyBackBatchedGraph {N c mid h w kHd kWd r : ℕ} (We : Kernel4 mid c 1 1) (be : Vec mid) (εe : ℝ) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ℝ) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 c mid 1 1) (bp : Vec c) (εp : ℝ) (γp βp : Vec c) (x : Vec (N * (c * h * w))) (e : SHlo (N * (c * h * w))) :
    SHlo (N * (c * h * w))

    The batched MBConv body backward graph: the four stage graphs chained at their cumulative forward activations (cbsB⁻¹ ∘ dwbsB⁻¹ ∘ seB⁻¹ ∘ projB⁻¹).

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      source
      theorem Proofs.StableHLO.mbBodyBackBatchedGraph_faithful {N c mid h w kHd kWd r : ℕ} (We : Kernel4 mid c 1 1) (be : Vec mid) (εe : ℝ) (hεe : 0 < εe) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ℝ) (hεd : 0 < εd) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 c mid 1 1) (bp : Vec c) (εp : ℝ) (hεp : 0 < εp) (γp βp : Vec c) (x : Vec (N * (c * h * w))) (e : SHlo (N * (c * h * w))) :
      den (mbBodyBackBatchedGraph We be εe γe βe Wd bd εd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp γp βp x e) = (mbExpFwdBHasVJP N We be εe hεe γe βe Wd bd εd hεd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp hεp γp βp).backward x (den e)
      source
      noncomputable def Proofs.StableHLO.mbDownBodyBackBatchedGraph {N ic mid oc h w kHd kWd r : ℕ} (We : Kernel4 mid ic 1 1) (be : Vec mid) (εe : ℝ) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ℝ) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 oc mid 1 1) (bp : Vec oc) (εp : ℝ) (γp βp : Vec oc) (x : Vec (N * (ic * (2 * h) * (2 * w)))) (e : SHlo (N * (oc * h * w))) :
      SHlo (N * (ic * (2 * h) * (2 * w)))

      The batched downsample MBConv body backward graph: the four stage graphs chained at their cumulative forward activations (cbsB⁻¹ ∘ dwbsSXlaB⁻¹ ∘ seB⁻¹ ∘ projB⁻¹). Stride-2 analogue of mbBodyBackBatchedGraph (strided depthwise stage graph).

      Equations
      • One or more equations did not get rendered due to their size.
      Instances For
        source
        theorem Proofs.StableHLO.mbDownBodyBackBatchedGraph_faithful {N ic mid oc h w kHd kWd r : ℕ} (We : Kernel4 mid ic 1 1) (be : Vec mid) (εe : ℝ) (hεe : 0 < εe) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ℝ) (hεd : 0 < εd) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 oc mid 1 1) (bp : Vec oc) (εp : ℝ) (hεp : 0 < εp) (γp βp : Vec oc) (x : Vec (N * (ic * (2 * h) * (2 * w)))) (e : SHlo (N * (oc * h * w))) :
        den (mbDownBodyBackBatchedGraph We be εe γe βe Wd bd εd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp γp βp x e) = (mbStridedFwdBHasVJP N We be εe hεe γe βe Wd bd εd hεd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp hεp γp βp).backward x (den e)

        CAPSTONE — the batched EfficientNet DOWNSAMPLE MBConv body: backward graph ↔ the proven mbStridedFwdBHasVJP. The four batched stage backward graphs (cbsB/dwbsSXlaB/seB/projB) chained at their forward activations, proven equal to the downsample-body VJP. The stride-2 analogue of mbBodyBackBatchedGraph_faithful (no residual skip — the downsample block changes spatial/channels, so the body alone is the block). EfficientNet uses swish (a global VJP), so this stays in the clean global HasVJP/vjpComp form (no _at recompute, unlike r34/mnv2's relu blocks).

        source
        noncomputable def Proofs.StableHLO.mbResidBlockBackBatchedGraph {N c mid h w kHd kWd r : ℕ} (We : Kernel4 mid c 1 1) (be : Vec mid) (εe : ℝ) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ℝ) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 c mid 1 1) (bp : Vec c) (εp : ℝ) (γp βp : Vec c) (x : Vec (N * (c * h * w))) (ecot : SHlo (N * (c * h * w))) :
        SHlo (N * (c * h * w))

        The whole batched MBConv residual block backward graph (body + identity skip).

        Equations
        • One or more equations did not get rendered due to their size.
        Instances For
          source
          theorem Proofs.StableHLO.mbResidBlockBackBatchedGraph_faithful {N c mid h w kHd kWd r : ℕ} (We : Kernel4 mid c 1 1) (be : Vec mid) (εe : ℝ) (hεe : 0 < εe) (γe βe : Vec mid) (Wd : DepthwiseKernel mid kHd kWd) (bd : Vec mid) (εd : ℝ) (hεd : 0 < εd) (γd βd : Vec mid) (Wz₁ : Mat mid r) (bz₁ : Vec r) (Wz₂ : Mat r mid) (bz₂ : Vec mid) (Wp : Kernel4 c mid 1 1) (bp : Vec c) (εp : ℝ) (hεp : 0 < εp) (γp βp : Vec c) (x : Vec (N * (c * h * w))) (ecot : SHlo (N * (c * h * w))) :
          den (mbResidBlockBackBatchedGraph We be εe γe βe Wd bd εd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp γp βp x ecot) = (mbResidFwdBHasVJP N We be εe hεe γe βe Wd bd εd hεd γd βd Wz₁ bz₁ Wz₂ bz₂ Wp bp εp hεp γp βp).backward x (den ecot)

          CAPSTONE — the whole batched EfficientNet MBConv residual block: backward graph ↔ the proven mbResidFwdBHasVJP. The four batched stage backward graphs (cbsB/dwbsB/seB/projB) chained at their forward activations + the identity skip, proven equal to the repo's batched MBConv block VJP.