Documentation

LeanMlir.Proofs.Foundation.HeadLayers

The classifier head as CertLayers — batched GAP and the dense classifier #

Every conv net in the suite ends in global average pooling and a dense classifier, and both are globally certified (GAP is linear, dense is affine), so ok := True. Written once here and shared by MobileNetV4, ResNet-34 and ResNet-50.

⭐ Both backward graphs tie by rfl. den of .gapBackBatched is definitionally the row-wise GAP VJP, and den of .denseRowBack is rowDenseBackFlat, which is what batchMap_has_vjp reduces to. ⚠ GAP's VJP does not depend on its input, and den .gapBackBatched uses that by evaluating the backward at fun _ => 0. That is sound because GAP is linear, and it is why the tie holds at every x.

source
noncomputable def Proofs.StableHLO.gapLayer (N : ℕ) {c h w : ℕ} :
CertLayer (N * (c * h * w)) (N * c)

Batched global average pool as a CertLayer. Globally certified — GAP is linear.

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    source
    theorem Proofs.StableHLO.gapLayer_fwd_apply (N : ℕ) {c h w : ℕ} (v : Vec (N * (c * h * w))) :
    (gapLayer N).fwd v = batchMap N (globalAvgPoolFlat c h w) v

    The GAP layer's forward is the batched global average pool.

    source
    noncomputable def Proofs.StableHLO.denseLayer (N : ℕ) {a nC : ℕ} (W : Mat a nC) (b : Vec nC) :
    CertLayer (N * a) (N * nC)

    Batched dense classifier as a CertLayer. Globally certified — dense is affine.

    Equations
    • One or more equations did not get rendered due to their size.
    Instances For
      source
      theorem Proofs.StableHLO.denseLayer_fwd_apply (N : ℕ) {a nC : ℕ} (W : Mat a nC) (b : Vec nC) (v : Vec (N * a)) :
      (denseLayer N W b).fwd v = batchMap N (dense W b) v

      The classifier layer's forward is the batched dense.