Documentation

LeanMlir.Proofs.Nets.Small.MlpTrainStep

The MLP train step: per-layer parameter-gradient assembly #

The MLP (dense → relu → dense → relu → dense) train step updates six parameters W₀,b₀,W₁,b₁,W₂,b₂. Each layer's gradient is the assembly of backprop: ∂L/∂Wᵢ = (layer-i input) ⊗ (cotangent the backward chain delivers at layer i), ∂L/∂bᵢ = (that cotangent).

The generic per-layer bridges already certify this for any backward subgraph e : Back (weight_grad_bridge/bias_grad_bridge, IR.lean) — emit*Grad equals the certified Jacobian of that dense layer contracted with e.denote dy. So the assembly is choosing the right cotangent subgraph per layer:

This file supplies mlpCotOut0 and its weight/bias bridges — completing the three-layer assembly, the multi-layer step the one-layer linear net has no analogue of — and the total-loss folds (StableHLO.lossWeightGrad_eq_sum, and at smooth points mlp_hidden_total_loss_grad, mlp_input_total_loss_grad, mlp_whole_net_weight_grads). The SGD wrapping θ − lr·∇ on top is identical to the linear case (StableHLO.sgdW).

source
def Proofs.IR.mlpCotOut0 {d₁ d₂ d₃ : ℕ} (W₁ : Mat d₁ d₂) (W₂ : Mat d₂ d₃) (p₀ : Vec d₁) (p₁ : Vec d₂) :
Back d₃ d₁

Layer-0 cotangent subgraph — the cotangent the backward chain delivers at the layer-0 dense output p₀: relu'(p₀) ⊙ (W₁ · mlpCotOut1). Prepends one more relu-back ∘ dense-back to mlpCotOut1, exactly as mlpCotOut1 extends Back.cotangent.

Equations
Instances For
    source
    theorem Proofs.IR.mlp_layer0_weight_grad_bridge {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (W₂ : Mat d₂ d₃) (x₀ : Vec d₀) (p₀ : Vec d₁) (p₁ : Vec d₂) (dy : Vec d₃) (i : Fin d₀) (j : Fin d₁) :
    emitWeightGrad x₀ (mlpCotOut0 W₁ W₂ p₀ p₁) dy i j = ∑ k : Fin d₁, pdiv (fun (v : Vec (d₀ * d₁)) => dense (Mat.unflatten v) b₀ x₀) W₀.flatten (finProdFinEquiv (i, j)) k * (mlpCotOut0 W₁ W₂ p₀ p₁).denote dy k

    MLP layer-0 weight-gradient bridge. The emitted layer-0 weight gradient (x₀ ⊗ mlpCotOut0) equals the certified Jacobian of the layer-0 dense wrt W₀, contracted with the cotangent the backward chain delivers there — the deepest chain (relu'(p₀) ⊙ W₁ · relu'(p₁) ⊙ W₂ · dy). The layer-0 peer of mlp_layer1_weight_grad_bridge.

    source
    theorem Proofs.IR.mlp_layer0_bias_grad_bridge {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (W₂ : Mat d₂ d₃) (x₀ : Vec d₀) (p₀ : Vec d₁) (p₁ : Vec d₂) (dy : Vec d₃) (i : Fin d₁) :
    emitBiasGrad (mlpCotOut0 W₁ W₂ p₀ p₁) dy i = ∑ j : Fin d₁, pdiv (fun (b' : Vec d₁) => dense W₀ b' x₀) b₀ i j * (mlpCotOut0 W₁ W₂ p₀ p₁).denote dy j

    MLP layer-0 bias-gradient bridge. Likewise the layer-0 bias gradient is the certified ∂/∂b₀ Jacobian contracted with the same deepest cotangent.

    source
    theorem Proofs.IR.mlp_layer1_bias_grad_bridge {d₁ d₂ d₃ : ℕ} (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (x₁ : Vec d₁) (p₁ : Vec d₂) (dy : Vec d₃) (i : Fin d₂) :
    emitBiasGrad (mlpCotOut1 W₂ p₁) dy i = ∑ j : Fin d₂, pdiv (fun (b' : Vec d₂) => dense W₁ b' x₁) b₁ i j * (mlpCotOut1 W₂ p₁).denote dy j

    MLP layer-1 bias-gradient bridge (the bias peer of the existing mlp_layer1_weight_grad_bridge).

    source
    theorem Proofs.IR.mlp_layer2_weight_grad_bridge {d₂ d₃ : ℕ} (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) (x₂ : Vec d₂) (dy : Vec d₃) (i : Fin d₂) (j : Fin d₃) :
    emitWeightGrad x₂ Back.cotangent dy i j = ∑ k : Fin d₃, pdiv (fun (v : Vec (d₂ * d₃)) => dense (Mat.unflatten v) b₂ x₂) W₂.flatten (finProdFinEquiv (i, j)) k * dy k

    MLP layer-2 (output) weight-gradient bridge — the output layer's cotangent is the loss cotangent dy itself (Back.cotangent, denote dy = dy); the bridge is the generic one specialized there.

    source
    theorem Proofs.IR.mlpCotOut1_denote {d₂ d₃ : ℕ} (W₂ : Mat d₂ d₃) (p₁ : Vec d₂) (g : Vec d₃) :
    (mlpCotOut1 W₂ p₁).denote g = fun (i : Fin d₂) => if p₁ i > 0 then W₂.mulVec g i else 0

    Layer-1 cotangent, explicit. mlpCotOut1.denote g = relu'(p₁) ⊙ (W₂ · g).

    source
    theorem Proofs.IR.mlpCotOut0_denote {d₁ d₂ d₃ : ℕ} (W₁ : Mat d₁ d₂) (W₂ : Mat d₂ d₃) (p₀ : Vec d₁) (p₁ : Vec d₂) (g : Vec d₃) :
    (mlpCotOut0 W₁ W₂ p₀ p₁).denote g = fun (i : Fin d₁) => if p₀ i > 0 then W₁.mulVec (fun (k : Fin d₂) => if p₁ k > 0 then W₂.mulVec g k else 0) i else 0

    Layer-0 cotangent, explicit — the deepest chain relu'(p₀) ⊙ (W₁ · (relu'(p₁) ⊙ (W₂ · g))).

    source
    theorem Proofs.IR.mlp_hidden_total_loss_grad {d₁ d₂ d₃ : ℕ} (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) (a₀ : Vec d₁) (label : Fin d₃) (h_smooth : ∀ (k : Fin d₂), dense W₁ b₁ a₀ k ≠ 0) (i : Fin d₁) (j : Fin d₂) :
    pdiv (fun (v : Vec (d₁ * d₂)) (x : Fin 1) => crossEntropy d₃ (dense W₂ b₂ (relu d₂ (dense (Mat.unflatten v) b₁ a₀))) label) W₁.flatten (finProdFinEquiv (i, j)) 0 = ∑ k : Fin d₂, pdiv (fun (v : Vec (d₁ * d₂)) => dense (Mat.unflatten v) b₁ a₀) W₁.flatten (finProdFinEquiv (i, j)) k * pdiv (fun (z : Vec d₂) (x : Fin 1) => crossEntropy d₃ (dense W₂ b₂ (relu d₂ z)) label) (dense W₁ b₁ a₀) k 0

    Hidden-layer total-loss fold (conditional). At a smooth point — the hidden pre-activation p₁ = dense W₁ b₁ a₀ off the ReLU kinks — the single gradient of the whole softmax-CE loss wrt the hidden weights W₁ folds, by the chain rule (pdiv_comp), into the certified ∂p₁/∂W₁ contracted with the loss gradient at the hidden pre-activation, ∂L/∂p₁, left as a pdiv. (Off the kinks that factor is the cotangent the backward chain delivers at layer 1, relu'(p₁) ⊙ (W₂ · (softmax−onehot)), cf. mlpCotOut1_denote; that identification is not stated here.) Conditionality is intrinsic: the chain runs back through the ReLU kink, so — unlike the linear / output-layer fold — this needs the smoothness hypothesis. The hidden-layer analogue of lossWeightGrad_eq_sum.

    source
    theorem Proofs.IR.mlp_input_total_loss_grad {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) (x : Vec d₀) (label : Fin d₃) (h_smooth_0 : ∀ (k : Fin d₁), dense W₀ b₀ x k ≠ 0) (h_smooth_1 : ∀ (k : Fin d₂), dense W₁ b₁ (relu d₁ (dense W₀ b₀ x)) k ≠ 0) (i : Fin d₀) (j : Fin d₁) :
    pdiv (fun (v : Vec (d₀ * d₁)) (x_1 : Fin 1) => crossEntropy d₃ (dense W₂ b₂ (relu d₂ (dense W₁ b₁ (relu d₁ (dense (Mat.unflatten v) b₀ x))))) label) W₀.flatten (finProdFinEquiv (i, j)) 0 = ∑ k : Fin d₁, pdiv (fun (v : Vec (d₀ * d₁)) => dense (Mat.unflatten v) b₀ x) W₀.flatten (finProdFinEquiv (i, j)) k * pdiv (fun (z : Vec d₁) (x : Fin 1) => crossEntropy d₃ (dense W₂ b₂ (relu d₂ (dense W₁ b₁ (relu d₁ z)))) label) (dense W₀ b₀ x) k 0

    Input-layer total-loss fold (conditional, deepest). The same fold for the first layer W₀, whose chain runs back through both ReLUs — so it carries both smoothness hypotheses (the same pair as mlpHasVJPAt). The total loss gradient wrt W₀ = certified ∂p₀/∂W₀ contracted with the loss gradient at p₀, left as a pdiv (its identification with the chain cotangent mlpCotOut0_denote computes is not stated here).

    source
    theorem Proofs.IR.mlp_whole_net_weight_grads {d₀ d₁ d₂ d₃ : ℕ} (W₀ : Mat d₀ d₁) (b₀ : Vec d₁) (W₁ : Mat d₁ d₂) (b₁ : Vec d₂) (W₂ : Mat d₂ d₃) (b₂ : Vec d₃) (x : Vec d₀) (label : Fin d₃) (h_smooth_0 : ∀ (k : Fin d₁), dense W₀ b₀ x k ≠ 0) (h_smooth_1 : ∀ (k : Fin d₂), dense W₁ b₁ (relu d₁ (dense W₀ b₀ x)) k ≠ 0) :
    (∀ (i : Fin d₀) (j : Fin d₁), pdiv (fun (v : Vec (d₀ * d₁)) (x_1 : Fin 1) => crossEntropy d₃ (dense W₂ b₂ (relu d₂ (dense W₁ b₁ (relu d₁ (dense (Mat.unflatten v) b₀ x))))) label) W₀.flatten (finProdFinEquiv (i, j)) 0 = ∑ k : Fin d₁, pdiv (fun (v : Vec (d₀ * d₁)) => dense (Mat.unflatten v) b₀ x) W₀.flatten (finProdFinEquiv (i, j)) k * pdiv (fun (z : Vec d₁) (x : Fin 1) => crossEntropy d₃ (dense W₂ b₂ (relu d₂ (dense W₁ b₁ (relu d₁ z)))) label) (dense W₀ b₀ x) k 0) ∧ (∀ (i : Fin d₁) (j : Fin d₂), pdiv (fun (v : Vec (d₁ * d₂)) (x_1 : Fin 1) => crossEntropy d₃ (dense W₂ b₂ (relu d₂ (dense (Mat.unflatten v) b₁ (relu d₁ (dense W₀ b₀ x))))) label) W₁.flatten (finProdFinEquiv (i, j)) 0 = ∑ k : Fin d₂, pdiv (fun (v : Vec (d₁ * d₂)) => dense (Mat.unflatten v) b₁ (relu d₁ (dense W₀ b₀ x))) W₁.flatten (finProdFinEquiv (i, j)) k * pdiv (fun (z : Vec d₂) (x : Fin 1) => crossEntropy d₃ (dense W₂ b₂ (relu d₂ z)) label) (dense W₁ b₁ (relu d₁ (dense W₀ b₀ x))) k 0) ∧ ∀ (i : Fin d₂) (j : Fin d₃), pdiv (fun (v : Vec (d₂ * d₃)) (x_1 : Fin 1) => crossEntropy d₃ (dense (Mat.unflatten v) b₂ (relu d₂ (dense W₁ b₁ (relu d₁ (dense W₀ b₀ x))))) label) W₂.flatten (finProdFinEquiv (i, j)) 0 = ∑ k : Fin d₃, pdiv (fun (v : Vec (d₂ * d₃)) => dense (Mat.unflatten v) b₂ (relu d₂ (dense W₁ b₁ (relu d₁ (dense W₀ b₀ x))))) W₂.flatten (finProdFinEquiv (i, j)) k * (softmax d₃ (mnistLinear W₂ b₂ (relu d₂ (dense W₁ b₁ (relu d₁ (dense W₀ b₀ x))))) k - oneHot d₃ label k)

    Whole-network MLP weight-gradient capstone. The three weight layers' total-loss gradients, jointly, under the two smoothness hypotheses.