NetSpec-style layer DSL for the verified trainers (Tier-2)

A verified trainer should read like the reference MainResnetTrain.lean — a layer list + config + train — with the only difference being the formalization underneath. This file provides that surface:

• VLayer — the verified-vocabulary layer constructors (the ops that have proven HasVJP witnesses), mirroring the reference Layer;
• VerifiedNetSpec — a { layers := [...] } architecture, the single source of truth;
• toSpecs — folds layers into the (dims, initKind) param layout, so the layout is derived from the architecture rather than hand-listed a second time. (Kernel-check it against the audited XLayout.specs with #guard spec.toSpecs == XLayout.specs — see VerifiedNets.lean's #guard resnet34Verified.toSpecs == ResNet34Layout.specs.)

The architecture's faithfulness is the audited <net>_has_vjp theorem, which is itself a hand-unrolled foldl of the generic vjp_comp chain-rule combinator (Proofs/Tensor.lean) over these same layers — so the spec and the proof describe the same fold. Generating the verified StableHLO from layers (folding the proven op-emitters) and folding the proof via a netVjp term are the remaining Tier-2 / Tier-3 steps; for now the slug names the committed, audited render of this architecture.

inductive VLayer : Type

A verified-vocabulary layer. Restricted to ops with proven HasVJP witnesses; each carries enough to derive its slice of the param layout.

• convBn (ic oc k stride : Nat) : VLayer

  conv (oc←ic, k×k, stride) → per-channel BN → relu. Params {W,b,γ,β}.

• convBnNB (ic oc k stride : Nat) : VLayer

  conv → per-channel BN → relu with no conv bias — {W, γ, β}. BN removes a conv bias ((x+b) − mean(x+b) = x − mean(x)), so a BN-followed conv carries none in He et al.'s .convBn; ResNet-34 uses this and the nets that genuinely ship a bias use convBn (§2l step B, measured in tests/TestConvBiasZero.lean).

• maxPool (k stride : Nat) : VLayer

  max pool k×k / stride. No params.

• residualStage (ic oc nBlocks stride : Nat) : VLayer

  A basic-block residual stage: nBlocks blocks at oc channels. The first block downsamples (and projects the skip) iff stride ≠ 1 ∨ ic ≠ oc; the rest are identity.

• bottleneckStage (ic oc nBlocks stride : Nat) : VLayer

  A bottleneck residual stage (ResNet-50/101/152): nBlocks blocks, each 1×1 (oc/4) → BN → relu, 3×3 (oc/4) → BN → relu, 1×1 (oc) → BN, +skip, relu. The first block projects the skip (1×1 → BN) iff stride ≠ 1 ∨ ic ≠ oc — the same dispatch residualStage uses, and for R50 that fires on all four stages, because stage 1 changes 64→256 at stride 1 where R34's stage 1 is ic = oc.

  ⚠ This is ResNet v1.5, not He et al.'s v1: the stride sits on the 3×3 (and on the projection), with the leading 1×1 at stride 1. Measured off the reference (jax/Jax/Codegen.lean's bottleneck_block_down), not assumed — putting it on the first 1×1 compiles, trains, descends and is a different net (§2k's heavy-ball trap one layer up).

  No conv biases — every conv here is BN-followed, so a bias cannot reach the output (convBnNB's argument, four convs at a time). torchvision's R50 carries none either, which is why the derived count lands on the reference's 25,557,032 with no adjustment (§2m).

• globalAvgPool : VLayer

  global average pool. No params.

• dense (ic oc : Nat) : VLayer

  dense ic→oc. Params {W,b}.

• relu : VLayer

  ReLU activation (pointwise). No params.

• conv (ic oc k stride : Nat) : VLayer

  plain conv (oc←ic, k×k, stride) + bias, NO batch-norm. Params {W,b}.

• flatten : VLayer

  flatten [C,H,W] → vector. No params (a reshape).

• bn : VLayer

  scalar-global BatchNorm (the proven bnForward): normalize over the whole c·h·w feature map per example, scalar γ/β. Params {γ, β} (rank-0).

• bnPerChannel (oc : Nat) : VLayer

  per-channel (per-example) BatchNorm (the proven bnPerChannelFlat, m=h·w): normalize each of oc channels over its own h·w spatial map per example, per-channel γ/β [oc]. Train=eval (no running stats). Params {γ:[oc], β:[oc]}.

• invertedResidual (ic mid oc stride : Nat) : VLayer

  MobileNetV2 inverted-residual block (ic→mid→oc, depthwise stride): expand 1×1 conv→per-channel BN→relu6, depthwise 3×3→BN→relu6, project 1×1→BN (linear bottleneck),

  • residual when stride=1 ∧ ic=oc. Params {W,b,γ,β} ×3 (expand/depthwise/project).
• invertedResidualNB (ic mid oc stride : Nat) : VLayer

  MobileNetV2 inverted-residual block with no conv biases — {W,γ,β} ×3. Every conv in the block is BN-followed, so a bias cannot reach the output (the convBnNB argument, three convs at a time); the torchvision/JAX reference carries none, and ours carried 52 across the net (§2m: the +17,056-param gap to the reference, closed exactly). Kept beside invertedResidual rather than replacing it for convBnNB's reason — a net whose blocks genuinely ship biases should still be able to say so.

• mbConvSE (ic mid oc r k : Nat) : VLayer

  EfficientNet MBConv block (ic→mid=t·ic→oc, depthwise k×k, SE ratio r): expand 1×1 (skipped when mid=ic, i.e. t=1) → BN → swish, depthwise k×k → BN → swish, squeeze-excite (Ws₁[mid,r],bs₁[r],Ws₂[r,mid],bs₂[mid], sigmoid gate), project 1×1 → BN. Params: (expand{W,b,γ,β} if t≠1) ++ depthwise{W,b,γ,β} ++ SE{Ws₁,bs₁,Ws₂,bs₂} ++ project{W,b,γ,β}.

• mbConvSENB (ic mid oc r k : Nat) : VLayer

  EfficientNet MBConv with no conv biases on the BN-followed convs — expand/depthwise/ project become {W,γ,β}. ⚠ The squeeze-excite biases STAY. SE's two 1×1 convs are followed by an ACTIVATION (sigmoid gate), not BN, so nothing absorbs them and the reference carries them; only a BN-followed conv can drop its bias. That distinction is what made the +21,008 gap close exactly (§2m) — the audit's rule was "a rank-1 kind-2 param immediately after a rank-4 kernel", and SE's params are rank-2.

• uib (ic oc expand stride preDWk postDWk : Nat) : VLayer

  MobileNetV4 Universal Inverted Bottleneck (planning/archive/mnv4_verified.md): optional pre-DW (preDWk) → 1×1 expand ic→mid → optional post-DW (postDWk) → 1×1 project mid→oc, every conv BN-followed and therefore bias-free (convBnNB's argument, and what Spec.lean's baseline count already assumes). mid = ic * expand.

  ⭐ k = 0 means "omit that depthwise", which is how ONE constructor expresses all four of MNv4's block families — ExtraDW (both DWs), IB/MBConv (post only), ConvNeXt-like (pre only), FFN (neither). That is the architecture's whole "stop adding new block types" claim, and it is why this is one VLayer case rather than four.

  ⚠⚠ A wrong pre/post dispatch is INVISIBLE to this function. A pre-DW and a post-DW at the same k and the same channel count contribute identical parameter shapes, so emitting one where the table says the other yields a net that type-checks, trains, descends, and is not MobileNetV4. Same class as R50's stride-on-the-3×3 and the 2×2 stem pool. The gate is a forward tie against the reference on shared weights — not a param count, and not this.

• fusedMbConvNB (ic oc expand k stride : Nat) : VLayer

  Fused inverted bottleneck, single block, no squeeze-excite — EfficientNetV2's early-stage block, and MobileNetV4's stage 0. k×k regular conv ic→mid (stride) → BN → swish → 1×1 project mid→oc → BN, no activation after the project; skip iff stride = 1 ∧ ic = oc. "Fused" = the MBConv expand-1×1 and depthwise collapse into ONE regular k×k conv, which is why nothing here is depthwise. mid = ic * expand. Bias-free — both convs are BN-followed.

  ⚠⚠ THE ACTIVATION IS SWISH, NOT RELU, AND THAT IS A PAPER DEVIATION. MobileNetV4-Conv is a ReLU network, but both emitters that produced the 84.58% use swish here (jax/Jax/Codegen.lean:1031 — the reference — and MlirCodegen.lean:6148's emitConvBnTrainSwish), inherited from the block being shared with EfficientNetV2. Matching the REFERENCE is what lets the number be reproduced and tied; matching the PAPER would be a different net from the one with the result. Recorded rather than quietly fixed.

  ⚠ Deliberately narrower than the baseline Layer.fusedMbConv, which also carries nBlocks and useSE. MNv4 uses n = 1, useSE = false, and a layout whose render does not exist is a trap — so the constructor cannot express what this file cannot emit.

• convNextBlock (c : Nat) : VLayer

  ConvNeXt block @ c channels (expand ratio 4): depthwise 7×7 → scalar-LN → 1×1 expand c→4c → GELU → 1×1 project 4c→c → layerScale (per-channel γ). Params: depthwise{W,b}; LN{γ,β scalar}; expand{W,b}; project{W,b}; layerScale{γ:[c]}.

• convNextBlockCh (c : Nat) : VLayer

  ConvNeXt block with the real channel LayerNorm — γ,β : [c] instead of two rank-0 scalars (§2m). The normalisation axis changes with it (over c per spatial position, not over the whole c·h·w map), but that is invisible to the LAYOUT; what the layout sees is the affine going from 2 floats to 2c. Kept beside convNextBlock for convBnNB's reason.

• layerNorm (d : Nat) : VLayer

  Per-channel LayerNorm over d features (normalize ∘ [d] affine). Params {γ:[d], β:[d]} — the non-scalar form (cf. bn, which is scalar-global).

• transformerBlock (d m : Nat) : VLayer

  Pre-norm transformer block (dim d, MLP hidden m): LN1 → MHSA (Wq/Wk/Wv/Wo [d,d]) → +x → LN2 → MLP (d→m→d) → +x. Params: LN1{γ,β}; {Wq,bq,Wk,bk,Wv,bv,Wo,bo}; LN2{γ,β}; {Wfc1[d,m],bfc1, Wfc2[m,d],bfc2} (per-channel [d] LN).

• param (dims : Array Nat) (kind : Nat) : VLayer

  A bare learned parameter tensor (dims, initKind) — e.g. ViT's CLS token / positional embedding (not produced by any standard layer).

def instReprVLayer.repr : VLayer → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.
• instReprVLayer.repr VLayer.globalAvgPool prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "VLayer.globalAvgPool")).group prec✝
• instReprVLayer.repr VLayer.relu prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "VLayer.relu")).group prec✝
• instReprVLayer.repr VLayer.flatten prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "VLayer.flatten")).group prec✝
• instReprVLayer.repr VLayer.bn prec✝ = Repr.addAppParen (Std.Format.nest (if prec✝ ≥ 1024 then 1 else 2) (Std.Format.text "VLayer.bn")).group prec✝

@[implicit_reducible]

instance instReprVLayer : Repr VLayer

Equations

• instReprVLayer = { reprPrec := instReprVLayer.repr }

def VLayer.toSpecs : VLayer → Array (Array Nat × Nat)

The (dims, initKind) params this layer contributes, in func-arg order (initKind: 0 = He(fan-in), 1 = ones (γ), 2 = zeros (β / bias)).

Equations

• One or more equations did not get rendered due to their size.
• (VLayer.convBn a a_1 a_2 a_3).toSpecs = VLayer.convBnSpec✝ a a_1 a_2
• (VLayer.convBnNB a a_1 a_2 a_3).toSpecs = VLayer.convBnNBSpec✝ a a_1 a_2
• (VLayer.maxPool a a_1).toSpecs = #[]
• (VLayer.residualStage a a_1 a_2 a_3).toSpecs = VLayer.stageSpec✝ a a_1 a_2 a_3
• (VLayer.bottleneckStage a a_1 a_2 a_3).toSpecs = VLayer.bottleneckStageSpec✝ a a_1 a_2 a_3
• VLayer.globalAvgPool.toSpecs = #[]
• (VLayer.dense a a_1).toSpecs = #[(#[a, a_1], 0), (#[a_1], 2)]
• VLayer.relu.toSpecs = #[]
• (VLayer.conv a a_1 a_2 a_3).toSpecs = #[(#[a_1, a, a_2, a_2], 0), (#[a_1], 2)]
• VLayer.flatten.toSpecs = #[]
• VLayer.bn.toSpecs = #[(#[], 1), (#[], 2)]
• (VLayer.bnPerChannel a).toSpecs = #[(#[a], 1), (#[a], 2)]
• (VLayer.convNextBlock a).toSpecs = #[(#[a, 1, 7, 7], 0), (#[a], 2), (#[], 1), (#[], 2), (#[4 * a, a, 1, 1], 0), (#[4 * a], 2), (#[a, 4 * a, 1, 1], 0), (#[a], 2), (#[a], 1)]
• (VLayer.convNextBlockCh a).toSpecs = #[(#[a, 1, 7, 7], 0), (#[a], 2), (#[a], 1), (#[a], 2), (#[4 * a, a, 1, 1], 0), (#[4 * a], 2), (#[a, 4 * a, 1, 1], 0), (#[a], 2), (#[a], 1)]
• (VLayer.layerNorm a).toSpecs = #[(#[a], 1), (#[a], 2)]
• (VLayer.param a a_1).toSpecs = #[(a, a_1)]

structure VerifiedNetSpec : Type

A verified net as a NetSpec-style architecture. layers is the single source of truth; the param layout (toSpecs) and input width (d0) are derived from it.

• name : String
• slug : String

  Names the committed, audited render verified_mlir/<slug>_{train_step,fwd}.mlir.

• inC : Nat
• imageH : Nat
• imageW : Nat
• nClasses : Nat
• data : VerifiedData
• layers : List VLayer
• blurb : String
• bnChannels : Array Nat

  Per-BN-layer channel counts in forward order (empty = LayerNorm / no-BN). Drives running-stats BN threading in trainAdamSched — see VerifiedNet.bnChannels.

• dropKeeps : Array Float

  Stochastic-depth keep probabilities, one per drop site, in the render's signature order (planning/archive/stochastic_depth.md). Empty on every net without a *sd render.

  ⚠ A SECOND hand-list against the renderer's enetDropIdxs/enetDropTotal — the same two-lists shape as toSpecs == XLayout.specs, and for the same structural reason: this file sits DOWNSTREAM of VerifiedTrain, so the renderer cannot share the definition by import without inverting the dependency. tests/TestDropPathRamp.lean is the #guard that pins them, and it is what stops the ramp drifting the way §2k's α/K did.

• mlirDir : String

  Which directory this net's artifacts live in — see VerifiedNet.mlirDir. Default verified_mlir/ (the certified, pinned corpus); the width/batch SWEEP specs set .lake/build because they render from argv at run time and their output is a build product.

• dropoutKeep : Option (Float × Nat)

  ▶ CLASSIFIER DROPOUT (recipe_gaps.md gap C) — (keep_prob, per-example width), none when the net has none. See VerifiedNet.dropoutKeep for why the WIDTH is carried: this mask is per-ELEMENT where dropKeeps above is per-example, and every downstream difference (blob shape, draw count, DP shard split) falls out of that one number.

• shimScript : String

  The generated ImageNet batch shim this net streams — see VerifiedNet.shimScript for why there is no default and why an empty one refuses instead of falling back. Required on every .imagenet spec; meaningless on the others.

• lossSlot : Bool

  Does <slug>_train_step.mlir return the trailing report-only %loss scalar? Per-RENDER, and only the mlp / cnn chapter-2/3 renders carry it. See VerifiedNet.lossSlot.

def VerifiedNetSpec.toSpecs (s : VerifiedNetSpec) : Array (Array Nat × Nat)

The full (dims, initKind) param list, folded from layers (func-arg order).

Equations

• s.toSpecs = List.foldl (fun (acc : Array (Array Nat × Nat)) (L : VLayer) => acc ++ L.toSpecs) #[] s.layers

def VerifiedNetSpec.d0 (s : VerifiedNetSpec) : Nat

Per-example flattened input width.

Equations

• s.d0 = s.inC * s.imageH * s.imageW

def VerifiedNetSpec.toNet (s : VerifiedNetSpec) : VerifiedNet

Lower to the runtime VerifiedNet the driver consumes.

Equations

• One or more equations did not get rendered due to their size.

def VerifiedNetSpec.train (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

Train end-to-end (delegates to the shared VerifiedNet.train driver).

Equations

• s.train cfg dataDir = s.toNet.train cfg dataDir

def VerifiedNetSpec.trainLinear (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

Train the 2-parameter linear path (Chapter 1); see VerifiedNet.trainLinear.

Equations

• s.trainLinear cfg dataDir = s.toNet.trainLinear cfg dataDir

def VerifiedNetSpec.attackPgd (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

Phase-3 PGD adversarial attack (Chapter 1 linear); see VerifiedNet.attackPgd.

Equations

• s.attackPgd cfg dataDir = s.toNet.attackPgd cfg dataDir

def VerifiedNetSpec.attackPgdMlp (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

Phase-3 PGD attack on the MLP (Chapter 2); see VerifiedNet.attackPgdMlp.

Equations

• s.attackPgdMlp cfg dataDir = s.toNet.attackPgdMlp cfg dataDir

def VerifiedNetSpec.attackPgdCnn (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

Phase-3 PGD attack on the CNN (Chapter 3, the conv rung); see VerifiedNet.attackPgdCnn.

Equations

• s.attackPgdCnn cfg dataDir = s.toNet.attackPgdCnn cfg dataDir

def VerifiedNetSpec.attackPgdSpectralMlp (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) (caps : List Float) : IO Unit

Spectral-norm-constrained MLP training study; see VerifiedNet.attackPgdSpectralMlp.

Equations

• s.attackPgdSpectralMlp cfg dataDir caps = s.toNet.attackPgdSpectralMlp cfg dataDir caps

def VerifiedNetSpec.attackPgdSpectralCnn (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) (caps : List Float) : IO Unit

Spectral-norm-constrained CNN training study; see VerifiedNet.attackPgdSpectralCnn.

Equations

• s.attackPgdSpectralCnn cfg dataDir caps = s.toNet.attackPgdSpectralCnn cfg dataDir caps

def VerifiedNetSpec.attackPgdCifar (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

PGD attack on the CIFAR-10 CNN (the deeper conv rung); see VerifiedNet.attackPgdCifar.

Equations

• s.attackPgdCifar cfg dataDir = s.toNet.attackPgdCifar cfg dataDir

def VerifiedNetSpec.attackPgdCifarBn (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) : IO Unit

PGD attack on the CIFAR-10 CNN + per-channel BatchNorm; see VerifiedNet.attackPgdCifarBn.

Equations

• s.attackPgdCifarBn cfg dataDir = s.toNet.attackPgdCifarBn cfg dataDir

def VerifiedNetSpec.attackPgdSpectralCifar (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) (caps : List Float) : IO Unit

Spectral-norm-constrained CIFAR training study; see VerifiedNet.attackPgdSpectralCifar.

Equations

• s.attackPgdSpectralCifar cfg dataDir caps = s.toNet.attackPgdSpectralCifar cfg dataDir caps

def VerifiedNetSpec.smoothCertify (s : VerifiedNetSpec) (cfg : VerifiedConfig) (dataDir : String) (sigmas : List Float) : IO Unit

Randomized-smoothing certificate (Cohen 2019, depth-independent); see VerifiedNet.smoothCertify. Forward-only — works on any spec via its rendered fwd.

Equations

• s.smoothCertify cfg dataDir sigmas = s.toNet.smoothCertify cfg dataDir sigmas