TorchLean API

NN.Proofs.Autograd.FDeriv.MlpMse

MlpMse #

End-to-end analytic soundness for a first training target:

We prove that the usual backprop formulas (including parameter gradients) coincide with the adjoint of the Fréchet derivative (fderiv) over ℝ.

Notes:

PyTorch correspondence / citations #

theorem Proofs.Autograd.reluDerivCLM_apply {n : ℕ} (x dx : Vec n) (i : Fin n) :

Coordinate formula for reluDerivCLM: it scales each coordinate by relu'(xᵢ).

theorem Proofs.Autograd.reluDerivCLM_inner {n : ℕ} (x dx δ : Vec n) :
inner ℝ ((reluDerivCLM x) dx) δ = inner ℝ dx ((reluDerivCLM x) δ)

ReLU derivative map is self-adjoint w.r.t. the Euclidean inner product.

This is because it is a diagonal scaling map on ℝⁿ.

The adjoint of reluDerivCLM equals itself (self-adjoint operator).

The spec-level “input derivative” for a linear layer agrees with the Euclidean adjoint.

In words: the tensor expression for ∂(W x)/∂x applied to an upstream δ is Wᵀ δ, and this is exactly the adjoint of the CLM x ↦ W x.

noncomputable def Proofs.Autograd.affineMat {inDim outDim : ℕ} (W : Mat outDim inDim) (b : Vec outDim) :
Vec inDim → Vec outDim

Affine map using Mat parameters (same as affine, but with Mat instead of Matrix).

Instances For
    noncomputable def Proofs.Autograd.mlpVecMat {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) :
    Vec inDim → Vec outDim

    2-layer MLP in Euclidean Vec form, parameterized by Mat weights and Vec biases.

    This is the same computation as NN.Proofs.Autograd.FDeriv.Core’s mlpVec, but set up for parameter-gradient proofs where the parameter space is a Hilbert space.

    Instances For
      noncomputable def Proofs.Autograd.mse {n : ℕ} (t : Vec n) :
      Vec n → ℝ

      Mean-squared error loss (MSE) against a fixed target t:

      mse t y = (1/n) * ‖y - t‖².

      Instances For
        noncomputable def Proofs.Autograd.mseGrad {n : ℕ} (y t : Vec n) :
        Vec n

        Gradient of MSE with respect to y:

        ∇_y mse(t)(y) = (2/n) * (y - t).

        Instances For
          theorem Proofs.Autograd.hasFDerivAt_mse {n : ℕ} (t y : Vec n) :
          HasFDerivAt (mse t) ((2 / ↑n) • (innerSL ℝ) (y - t)) y

          Fréchet derivative of MSE, packaged as a continuous linear map Vec n →L ℝ.

          theorem Proofs.Autograd.adjoint_mseDeriv_one {n : ℕ} (t y : Vec n) :
          (ContinuousLinearMap.adjoint ((2 / ↑n) • (innerSL ℝ) (y - t))) 1 = mseGrad y t

          The adjoint of the MSE derivative, applied to the seed 1, is the gradient mseGrad y t.

          theorem Proofs.Autograd.mseGrad_eq_adjoint_fderiv {n : ℕ} (t y : Vec n) :
          (vjp (mse t) y) 1 = mseGrad y t

          The VJP of MSE at y with upstream seed 1 equals the usual gradient mseGrad y t.

          This is the scalar-loss specialization: for scalar loss ℓ, the gradient is (fderiv ℓ)† 1.

          theorem Proofs.Autograd.adjoint_fderiv_comp_apply_one {E : Type u_1} {F : Type u_2} [NormedAddCommGroup E] [InnerProductSpace ℝ E] [CompleteSpace E] [NormedAddCommGroup F] [InnerProductSpace ℝ F] [CompleteSpace F] {f : E → F} {g : F → ℝ} {f' : E →L[ℝ] F} {g' : F →L[ℝ] ℝ} {x : E} (hf : HasFDerivAt f f' x) (hg : HasFDerivAt g g' (f x)) :

          Convenience lemma: adjoint of the derivative of a scalar composition, applied to seed 1.

          For scalar loss g ∘ f, this is the reverse-mode “chain rule” in adjoint form.

          We name the intermediate activations so the gradient statements read like textbook backprop: mlpPreActivation is z₁ = W₁ x + b₁, mlpHidden is a₁ = relu z₁, and mlpOutput is the network output y = W₂ a₁ + b₂.

          noncomputable def Proofs.Autograd.mlpPreActivation {inDim hidDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (x : Vec inDim) :
          Vec hidDim

          Pre-activation W1 x + b1 (the ReLU differentiability hypothesis is stated on it).

          Instances For
            noncomputable def Proofs.Autograd.mlpHidden {inDim hidDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (x : Vec inDim) :
            Vec hidDim

            Hidden activation relu (W1 x + b1).

            Instances For
              noncomputable def Proofs.Autograd.mlpOutput {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) :
              Vec outDim

              Network output mlpVecMat W1 b1 W2 b2 x.

              Instances For
                theorem Proofs.Autograd.hasFDerivAt_mlp_wrt_W2 {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) :
                HasFDerivAt (fun (W2 : Mat outDim hidDim) => mlpVecMat W1 b1 W2 b2 x) (matApplyLin (mlpHidden W1 b1 x)) W2

                Fréchet derivative of the network output with respect to the second-layer weights W2.

                Informally: ∂y/∂W2 is the linear map dW2 ↦ dW2 a1.

                theorem Proofs.Autograd.grad_W2_mse {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) (t : Vec outDim) :
                (ContinuousLinearMap.adjoint (fderiv ℝ (fun (W2 : Mat outDim hidDim) => mse t (mlpVecMat W1 b1 W2 b2 x)) W2)) 1 = outer (mseGrad (mlpOutput W1 b1 W2 b2 x) t) (mlpHidden W1 b1 x)

                Closed-form gradient of scalar loss mse t (mlpVecMat …) with respect to W2.

                Result: ∂L/∂W2 = (∂L/∂y) ⊗ a1, i.e. outer product of the output gradient and hidden activation.

                theorem Proofs.Autograd.grad_b2_mse {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) (t : Vec outDim) :
                (ContinuousLinearMap.adjoint (fderiv ℝ (fun (b2 : Vec outDim) => mse t (mlpVecMat W1 b1 W2 b2 x)) b2)) 1 = mseGrad (mlpVecMat W1 b1 W2 b2 x) t

                Closed-form gradient of the scalar loss with respect to the second-layer bias b2.

                Result: ∂L/∂b2 = ∂L/∂y.

                theorem Proofs.Autograd.hasFDerivAt_mlp_wrt_b1 {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) (hx : ∀ (i : Fin hidDim), (mlpPreActivation W1 b1 x).ofLp i ≠ 0) :
                HasFDerivAt (fun (b1 : Vec hidDim) => mlpVecMat W1 b1 W2 b2 x) (matCLM (toMatrix W2) ∘SL reluDerivCLM (mlpPreActivation W1 b1 x)) b1

                Fréchet derivative of the network output with respect to the first-layer bias b1, under the ReLU “no kinks” hypothesis.

                Informally: ∂y/∂b1 = W2 ∘ ReLU'(z1).

                theorem Proofs.Autograd.grad_b1_mse {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) (t : Vec outDim) (hx : ∀ (i : Fin hidDim), (mlpPreActivation W1 b1 x).ofLp i ≠ 0) :
                (ContinuousLinearMap.adjoint (fderiv ℝ (fun (b1 : Vec hidDim) => mse t (mlpVecMat W1 b1 W2 b2 x)) b1)) 1 = (reluDerivCLM (mlpPreActivation W1 b1 x)) ((ContinuousLinearMap.adjoint (matCLM (toMatrix W2))) (mseGrad (mlpOutput W1 b1 W2 b2 x) t))

                Closed-form gradient of the scalar loss with respect to the first-layer bias b1 (under the ReLU “no kinks” hypothesis).

                theorem Proofs.Autograd.hasFDerivAt_mlp_wrt_x {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) (hx : ∀ (i : Fin hidDim), (mlpPreActivation W1 b1 x).ofLp i ≠ 0) :

                Fréchet derivative of the network output with respect to the input x, under the ReLU “no kinks” hypothesis.

                Informally: ∂y/∂x = W2 ∘ ReLU'(z1) ∘ W1.

                theorem Proofs.Autograd.grad_x_mse {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) (t : Vec outDim) (hx : ∀ (i : Fin hidDim), (mlpPreActivation W1 b1 x).ofLp i ≠ 0) :

                Closed-form gradient of the scalar loss with respect to the input x, under the ReLU “no kinks” hypothesis.

                theorem Proofs.Autograd.hasFDerivAt_mlp_wrt_W1 {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) (hx : ∀ (i : Fin hidDim), (mlpPreActivation W1 b1 x).ofLp i ≠ 0) :
                HasFDerivAt (fun (W1 : Mat hidDim inDim) => mlpVecMat W1 b1 W2 b2 x) (matCLM (toMatrix W2) ∘SL reluDerivCLM (mlpPreActivation W1 b1 x) ∘SL matApplyLin x) W1

                Fréchet derivative of the network output with respect to the first-layer weights W1, under the ReLU “no kinks” hypothesis.

                The derivative is linear in W1 through the slice dW1 ↦ dW1 x, then propagated through ReLU' and W2.

                theorem Proofs.Autograd.grad_W1_mse {inDim hidDim outDim : ℕ} (W1 : Mat hidDim inDim) (b1 : Vec hidDim) (W2 : Mat outDim hidDim) (b2 : Vec outDim) (x : Vec inDim) (t : Vec outDim) (hx : ∀ (i : Fin hidDim), (mlpPreActivation W1 b1 x).ofLp i ≠ 0) :
                (ContinuousLinearMap.adjoint (fderiv ℝ (fun (W1 : Mat hidDim inDim) => mse t (mlpVecMat W1 b1 W2 b2 x)) W1)) 1 = outer ((reluDerivCLM (mlpPreActivation W1 b1 x)) ((ContinuousLinearMap.adjoint (matCLM (toMatrix W2))) (mseGrad (mlpOutput W1 b1 W2 b2 x) t))) x

                Closed-form gradient of the scalar loss with respect to W1 (under the ReLU “no kinks” hypothesis).

                Result has the expected “outer product” form with backpropagated hidden gradient and input x.