Linear regression (spec model)

Defines linear regression as a dot product plus bias (one output):

y = wᵀ x + b

The corresponding PyTorch operations are:

• torch.nn.Linear(in_features, out_features=1) for the forward pass,
• torch.nn.functional.mse_loss(..., reduction="mean") for the MSE objective,
• an SGD-style parameter update step (as in torch.optim.SGD) for training.

This file is a spec: it states the math (forward + VJPs) with shapes tracked by the type system. It prioritizes clarity and explicit derivatives over performance, and it does not include the closed-form normal-equations solution.

Implementation status

No API builder implements this model as a unit (nn.linear is the layer, not this model with its loss and training step). No theorem relates it to runtime code.

structure Spec.LinearRegressionSpec (α : Type) [TorchLean.Storage α] (inDim : ℕ) : Type

Parameters for a single-output linear regression model.

PyTorch analogy: the weights and bias fields correspond to nn.Linear(inDim, 1).weight and nn.Linear(inDim, 1).bias, but with shapes tracked in the tensor type.

• weights : TorchLean.Tensor α [inDim]

  Regression coefficients, one per input feature.

• bias : TorchLean.Tensor α Shape.scalar

  Scalar intercept term.

def Spec.linearRegressionForwardSpec {α : Type} [TorchLean.Storage α] [Context α] {inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [inDim]) : TorchLean.Tensor α Shape.scalar

Forward pass for linear regression: y = wᵀ x + b.

def Spec.linearRegressionBatchedForwardSpec {α : Type} [TorchLean.Storage α] [Context α] {batch inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [batch, inDim]) : TorchLean.Tensor α [batch]

Batched forward pass, applied independently to each input row.

def Spec.linearRegressionWeightsDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {inDim : ℕ} (input : TorchLean.Tensor α [inDim]) (gradOutput : TorchLean.Tensor α Shape.scalar) : TorchLean.Tensor α [inDim]

VJP contribution for weights: dL/dw = x * (dL/dy) (scalar-times-vector scaling).

def Spec.linearRegressionBiasDerivSpec {α : Type} [TorchLean.Storage α] {inDim : ℕ} (_weights : TorchLean.Tensor α [inDim]) (gradOutput : TorchLean.Tensor α Shape.scalar) (_input : TorchLean.Tensor α [inDim]) : TorchLean.Tensor α Shape.scalar

VJP contribution for bias: dL/db = dL/dy.

def Spec.linearRegressionInputDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {inDim : ℕ} (weights : TorchLean.Tensor α [inDim]) (gradOutput : TorchLean.Tensor α Shape.scalar) : TorchLean.Tensor α [inDim]

VJP contribution for input: dL/dx = w * (dL/dy).

structure Spec.LinearRegressionGradients (α : Type) [TorchLean.Storage α] (inDim : ℕ) (inputShape : Shape) : Type

Gradients for a linear regression model.

inputShape is a parameter so that the unbatched and batched backward passes return the same record: the parameter gradients have the same shape either way, and only the input gradient grows a leading batch axis.

• weightGradient : TorchLean.Tensor α [inDim]

  Gradient with respect to the weight vector.

• biasGradient : TorchLean.Tensor α Shape.scalar

  Gradient with respect to the scalar bias.

• inputGradient : TorchLean.Tensor α inputShape

  Gradient with respect to the input.

def Spec.linearRegressionBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [inDim]) (gradOutput : TorchLean.Tensor α Shape.scalar) : LinearRegressionGradients α inDim [inDim]

Full backward pass for one example.

def Spec.linearRegressionBatchedBackwardSpec {α : Type} [TorchLean.Storage α] [Context α] {batch inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [batch, inDim]) (gradOutput : TorchLean.Tensor α [batch]) (h : batch ≠ 0) : LinearRegressionGradients α inDim [batch, inDim]

Batched backward pass.

This aggregates parameter gradients across the batch (a sum over batch), matching PyTorch's default behavior for loss reductions like "mean" when you subsequently scale appropriately.

def Spec.mseLossSpec {α : Type} [TorchLean.Storage α] [Context α] {batch inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [batch, inDim]) (target : TorchLean.Tensor α [batch]) (h : batch ≠ 0) : TorchLean.Tensor α Shape.scalar

Mean Squared Error loss (MSE).

PyTorch analogy: F.mse_loss(predictions, target, reduction="mean").

Note: the batch ≠ 0 hypothesis avoids dividing by zero.

def Spec.mseLossGradSpec {α : Type} [TorchLean.Storage α] [Context α] {batch : ℕ} (predictions target : TorchLean.Tensor α [batch]) : TorchLean.Tensor α [batch]

Gradient of MSE w.r.t. predictions: d/dy (mean (y - t)^2) = (2/batch) * (y - t).

This is only meaningful when batch > 0 (callers typically already carry batch ≠ 0).

def Spec.linearRegressionTrainStepSpec {α : Type} [TorchLean.Storage α] [Context α] {batch inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [batch, inDim]) (target : TorchLean.Tensor α [batch]) (learningRate : α) (h : batch ≠ 0) : TorchLean.Tensor α Shape.scalar × LinearRegressionSpec α inDim

One gradient-descent training step for linear regression.

def Spec.linearRegressionOpSpec {α : Type} [TorchLean.Storage α] [Context α] {inDim : ℕ} (model : LinearRegressionSpec α inDim) : OpSpec α [inDim] Shape.scalar

OpSpec wrapper for linear regression.

This is useful when composing the op in a spec-level AD development.

def Spec.rSquaredSpec {α : Type} [TorchLean.Storage α] [Context α] {batch inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [batch, inDim]) (target : TorchLean.Tensor α [batch]) (h : batch ≠ 0) : TorchLean.Tensor α Shape.scalar

R-squared (coefficient of determination) for model evaluation.

PyTorch analogy: there is no single built-in for R² in core PyTorch; this matches the standard definition 1 - SS_res / SS_tot.

Note: if SS_tot = 0 (targets are constant), this divides by zero. Many libraries treat that as a special case; this spec keeps the plain formula.

def Spec.ridgeLossSpec {α : Type} [TorchLean.Storage α] [Context α] {batch inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [batch, inDim]) (target : TorchLean.Tensor α [batch]) (lambda : α) (h : batch ≠ 0) : TorchLean.Tensor α Shape.scalar

Ridge loss: MSE plus lambda * ||w||_2^2.

Regularization changes the training objective. For fixed parameters, predictions use linearRegressionForwardSpec and do not take a regularization coefficient.

Reference: Hoerl and Kennard, "Ridge Regression: Biased Estimation for Nonorthogonal Problems" (1970). https://doi.org/10.1080/00401706.1970.10488634

def Spec.ridgeWeightsDerivSpec {α : Type} [TorchLean.Storage α] [Context α] {batch inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [batch, inDim]) (gradOutput : TorchLean.Tensor α [batch]) (lambda : α) (h : batch ≠ 0) : TorchLean.Tensor α [inDim]

Ridge gradient w.r.t. weights.

This is the usual batched gradient plus the derivative of lambda * ||w||_2^2, which contributes 2 * lambda * w.

def Spec.lassoSoftThresholdSpec {α : Type} [TorchLean.Storage α] [Context α] {inDim : ℕ} (weights : TorchLean.Tensor α [inDim]) (threshold : α) : TorchLean.Tensor α [inDim]

Soft-thresholding operator (often written S_λ), used in proximal-gradient updates for L1.

Reference: Tibshirani, "Regression Shrinkage and Selection via the Lasso" (1996). https://doi.org/10.1111/j.2517-6161.1996.tb02080.x

def Spec.lassoLossSpec {α : Type} [TorchLean.Storage α] [Context α] {batch inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [batch, inDim]) (target : TorchLean.Tensor α [batch]) (lambda : α) (h : batch ≠ 0) : TorchLean.Tensor α Shape.scalar

Lasso loss: MSE plus lambda * ||w||_1.

Regularization changes the training objective. For fixed parameters, predictions use linearRegressionForwardSpec and do not take a regularization coefficient.

def Spec.elasticNetLossSpec {α : Type} [TorchLean.Storage α] [Context α] {batch inDim : ℕ} (model : LinearRegressionSpec α inDim) (input : TorchLean.Tensor α [batch, inDim]) (target : TorchLean.Tensor α [batch]) (l1_ratio alpha : α) (h : batch ≠ 0) : TorchLean.Tensor α Shape.scalar

Elastic net loss: a convex combination of L1 and L2 penalties.

Reference: Zou and Hastie, "Regularization and Variable Selection via the Elastic Net" (2005). https://doi.org/10.1111/j.1467-9868.2005.00503.x

Polynomial features

Polynomial regression can be expressed as linear regression on a fixed feature expansion φ(x) = [x, x^2, ..., x^degree] (per input coordinate). We keep this as a named helper, then reuse linearRegressionForwardSpec on the expanded input.

def Spec.polynomialFeaturesSpec {α : Type} [TorchLean.Storage α] [Context α] {inDim : ℕ} (degree : ℕ) (input : TorchLean.Tensor α [inDim]) : TorchLean.Tensor α [inDim * degree]

Expand a length-inDim input vector into polynomial features up to degree.

This expansion does not include a constant feature (the model bias already plays that role). Features are ordered by degree, with all input coordinates at power one followed by all at power two, and so on. Degree zero or an empty input gives an empty output.

def Spec.polynomialRegressionForwardSpec {α : Type} [TorchLean.Storage α] [Context α] {inDim degree : ℕ} (model : LinearRegressionSpec α (inDim * degree)) (input : TorchLean.Tensor α [inDim]) : TorchLean.Tensor α Shape.scalar

Forward pass for polynomial regression: expand features, then run linear regression.