Documentation

LeanMlir.Proofs.Nets.ViT.ViTFold

ViT-Tiny fold — each emitted param-SGD op denotes the certified SGD step at its cotangent #

The ViT peer of ConvNeXtFold/EfficientNetFold: for every param-SGD op the vitTrainStepRenderV renderer emits, prove den(op) = θ − lr·(certified Jacobian · cotangent). The vector-LN γ / β nodes (veclnGammaSgd_den, rowDenseBiasSgd_den_lnbeta and their tie clauses) are in SgdNodes, under this namespace, because ConvNeXt's head uses them too. Each is a one-or-few-line delegation to the already-proven render certs in ViTVecLN (vector-[D] LN γ/β) and TokenParamGrad (rowwise dense W/b, patch conv W/b, cls, pos); the classifier head reuses Cifar8Tie.dense{W,B}_den. Together these cover EVERY parameter family of the depth-12 ViT-Tiny train step (200 params), so the step tie (ViTStepTie) can thread them at the real backward chain cotangents.

The op dens and the cert LHSs line up by construction (the core ops were built to denote exactly these grads): veclnGammaSgd→layerNormVec_gamma_sgd_certified, rowDenseWeightSgd→ rowDense_weight_sgd_certified, rowDenseBiasSgd→rowDense_bias_sgd_certified (dense bias) and layerNormVec_beta_sgd_certified (LN β — same op, different forward in the pdiv), patchEmbedWeightSgd →patchEmbed_weight_sgd_certified, patchEmbedBiasSgd→patchEmbed_bias_sgd_certified, posEmbedSgd→ posEmbed_sgd_certified, cls (clsSliceF→denseBiasSgdB)→clsToken_sgd_certified.

source
theorem Proofs.ViTFold.rowDenseWeightSgd_den {N a c : ℕ} (xN wN lrStr cotN : String) (bb : Vec c) (x : Vec (N * a)) (W : Mat a c) (dy : Vec (N * c)) (lr : ℝ) (i : Fin a) (j : Fin c) :
StableHLO.den (StableHLO.SHlo.rowDenseWeightSgd xN wN lrStr x W lr (StableHLO.SHlo.operand cotN dy)) (finProdFinEquiv (i, j)) = W i j - lr * ∑ o : Fin (N * c), pdiv (fun (v : Vec (a * c)) => Mat.flatten fun (r : Fin N) => dense (Mat.unflatten v) bb (Mat.unflatten x r)) W.flatten (finProdFinEquiv (i, j)) o * dy o

Per-token dense weight op denotes the certified step. den(rowDenseWeightSgd) (flat (i,j)) = W_ij − lr·(Σ_tokens x·dy), the certified ∂(rowwise dense)/∂W contraction. Covers Wq/Wk/Wv/Wo/ Wfc1/Wfc2 (all 6 per-block denses). Delegation to rowDense_weight_sgd_certified.

source
theorem Proofs.ViTFold.rowDenseBiasSgd_den {N a c : ℕ} (bN lrStr cotN : String) (W : Mat a c) (X : Mat N a) (b : Vec c) (dy : Vec (N * c)) (lr : ℝ) (i : Fin c) :
StableHLO.den (StableHLO.SHlo.rowDenseBiasSgd bN lrStr b lr (StableHLO.SHlo.operand cotN dy)) i = b i - lr * ∑ o : Fin (N * c), pdiv (fun (b' : Vec c) => Mat.flatten fun (r : Fin N) => dense W b' (X r)) b i o * dy o

Per-token dense bias op denotes the certified step (dense-bias forward). den(rowDenseBiasSgd) = b − lr·(Σ_tokens dy). Covers bq/bk/bv/bo/bfc1/bfc2. Delegation to rowDense_bias_sgd_certified.

source
theorem Proofs.ViTFold.patchEmbedWeightSgd_den {ic H W P N D : ℕ} (wN xN lrStr cotN : String) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (Wp : Kernel4 D ic P P) (dy : Vec ((N + 1) * D)) (lr : ℝ) (d : Fin D) (c : Fin ic) (kh kw : Fin P) :
StableHLO.den (StableHLO.SHlo.patchEmbedWeightSgd wN xN lrStr img Wp lr (StableHLO.SHlo.operand cotN dy)) (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (d, c), kh), kw)) = Wp d c kh kw - lr * ∑ o : Fin ((N + 1) * D), pdiv (fun (v : Vec (D * ic * P * P)) => patchEmbedFlat ic H W P N D (Kernel4.unflatten v) bc cls pos img) Wp.flatten (finProdFinEquiv (finProdFinEquiv (finProdFinEquiv (d, c), kh), kw)) o * dy o

Patch-embed conv weight op denotes the certified step. den(patchEmbedWeightSgd) (flat (d,c,kh,kw)) = W − lr·(certified patchify-conv weight grad). The ViT analogue of ConvNeXt's stem 4×4/s4 weight — but here a VJP cert EXISTS, so it ties (vit has no even-kernel weight gap). Delegation to patchEmbed_weight_sgd_certified.

source
theorem Proofs.ViTFold.patchEmbedBiasSgd_den {ic H W P N D : ℕ} (bN lrStr cotN : String) (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (dy : Vec ((N + 1) * D)) (lr : ℝ) (i : Fin D) :
StableHLO.den (StableHLO.SHlo.patchEmbedBiasSgd bN lrStr bc lr (StableHLO.SHlo.operand cotN dy)) i = bc i - lr * ∑ o : Fin ((N + 1) * D), pdiv (fun (b' : Vec D) => patchEmbedFlat ic H W P N D Wc b' cls pos img) bc i o * dy o

Patch-embed conv bias op denotes the certified step. den(patchEmbedBiasSgd) = b − lr·(Σ_patches dy) (CLS row 0 excluded). Delegation to patchEmbed_bias_sgd_certified.

source
theorem Proofs.ViTFold.posEmbedSgd_den {ic H W P N D : ℕ} (pN lrStr cotN : String) (Wc : Kernel4 D ic P P) (bc cls : Vec D) (pos : Mat (N + 1) D) (img : Vec (ic * H * W)) (dy : Vec ((N + 1) * D)) (lr : ℝ) (i : Fin ((N + 1) * D)) :
StableHLO.den (StableHLO.SHlo.posEmbedSgd pN lrStr pos lr (StableHLO.SHlo.operand cotN dy)) i = pos.flatten i - lr * ∑ j : Fin ((N + 1) * D), pdiv (fun (p : Vec ((N + 1) * D)) => patchEmbedFlat ic H W P N D Wc bc cls (Mat.unflatten p) img) pos.flatten i j * dy j

Positional-embed op denotes the certified step. den(posEmbedSgd) = pos − lr·dy (the pos Jacobian is the identity — pos is added to every token). Delegation to posEmbed_sgd_certified.

source
def Proofs.ViTFold.RowDenseWSgdTied (N : ℕ) {a c : ℕ} (xN wN lrStr cotN : String) (bb : Vec c) (x : Vec (N * a)) (W : Mat a c) (dy : Vec (N * c)) (lr : ℝ) :
Prop

A per-token dense weight SGD node, tied (rowDenseWeightSgd_den).

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    source
    def Proofs.ViTFold.RowDenseBSgdTied (N : ℕ) {a c : ℕ} (bN lrStr cotN : String) (W : Mat a c) (x : Vec (N * a)) (b : Vec c) (dy : Vec (N * c)) (lr : ℝ) :
    Prop

    A per-token dense bias SGD node, tied (rowDenseBiasSgd_den).

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

      Each clause holds, every argument implicit (read off the goal by a step tie's constructor).

      source
      theorem Proofs.ViTFold.rowDenseWSgdTied_holds {N a c : ℕ} {xN wN lrStr cotN : String} {bb : Vec c} {x : Vec (N * a)} {W : Mat a c} {dy : Vec (N * c)} {lr : ℝ} :
      RowDenseWSgdTied N xN wN lrStr cotN bb x W dy lr
      source
      theorem Proofs.ViTFold.rowDenseBSgdTied_holds {N a c : ℕ} {bN lrStr cotN : String} {W : Mat a c} {x : Vec (N * a)} {b : Vec c} {dy : Vec (N * c)} {lr : ℝ} :
      RowDenseBSgdTied N bN lrStr cotN W x b dy lr