Bridging Scalar ReLU MLPs to Tensor Inputs

This file is a first “bridge step” between:

• the constructive scalar ReLU approximation theorem in UniversalApproximation.lean, and
• tensor inputs Tensor ℝ [n] used throughout TorchLean.

What is proved here (fully proved):

1. Exact representability of affine maps $x\mapsto w\mathbin{\cdot}x+b$ by a two-layer ReLU MLP of width $2$, using $\operatorname{ReLU}(u)-\operatorname{ReLU}(-u)=u$.
2. Ridge lifting: any scalar-input 2-layer ReLU MLP can be lifted to a tensor input via $u=w\mathbin{\cdot}x+c$, by scaling each first-layer weight by $w$ and adjusting biases accordingly.

What is not proved here: the full classical multivariate universal approximation theorem for ReLU MLPs. That requires substantially more formalization (e.g. piecewise-linear approximation machinery or a functional-analytic Cybenko/Leshno style proof).

theorem NN.MLTheory.Proofs.ReLUMlpBridge.getScalar_dim_getScalar {n : ℕ} (x : TorchLean.Tensor ℝ [n]) : (fun (i : Fin n) => (TorchLean.Tensor.dim fun (j : Fin n) => TorchLean.Tensor.scalar (x.getScalar j)).getScalar i) = fun (i : Fin n) => x.getScalar i

Rewrapping a rank-one tensor by Tensor.dim preserves every coordinate.

noncomputable def NN.MLTheory.Proofs.ReLUMlpBridge.dot {n : ℕ} (w : Fin n → ℝ) (x : TorchLean.Tensor ℝ [n]) : ℝ

Dot product $w\mathbin{\cdot}x$ for coordinate weights w and a rank-one tensor x.

noncomputable def NN.MLTheory.Proofs.ReLUMlpBridge.mlpEval {n hidDim : ℕ} (l1 : Spec.LinearSpec ℝ n hidDim) (l2 : Spec.LinearSpec ℝ hidDim 1) (x : TorchLean.Tensor ℝ [n]) : ℝ

Evaluate a single-hidden-layer ReLU MLP on a tensor input and return the scalar output.

theorem NN.MLTheory.Proofs.ReLUMlpBridge.relu_sub_relu_neg (u : ℝ) : UniversalApproximation.relu u - UniversalApproximation.relu (-u) = u

The identity $\operatorname{ReLU}(u)-\operatorname{ReLU}(-u)=u$, used to represent affine maps exactly with ReLU.

theorem NN.MLTheory.Proofs.ReLUMlpBridge.mlp_forward_eq_linear_relu_linear {n hidDim : ℕ} (l1 : Spec.LinearSpec ℝ n hidDim) (l2 : Spec.LinearSpec ℝ hidDim 1) (x : TorchLean.Tensor ℝ [n]) : Examples.mlpForward l1 l2 x = have z1 := Spec.linearSpec l1 x; have a1 := Activation.reluSpec z1; Spec.linearSpec l2 a1

Unfold mlpForward as $\operatorname{linear}\circ\operatorname{ReLU}\circ\operatorname{linear}$.

This lemma is used as the standard normalization step in “network algebra” proofs.

noncomputable def NN.MLTheory.Proofs.ReLUMlpBridge.matDim1Get {m : ℕ} (A : TorchLean.Tensor ℝ [m, 1]) (i : Fin m) : ℝ

Extract the unique entry from row i of an (m×1) tensor interpreted as a matrix.

theorem NN.MLTheory.Proofs.ReLUMlpBridge.mat_vec_mul_spec_dim1 {m : ℕ} (A : TorchLean.Tensor ℝ [m, 1]) (x : ℝ) : Spec.matVecMulSpec A (Spec.singleton x) = TorchLean.Tensor.dim fun (i : Fin m) => TorchLean.Tensor.scalar (matDim1Get A i * x)

Specialized matrix-vector multiplication when the input is a scalar (dimension 1).

theorem NN.MLTheory.Proofs.ReLUMlpBridge.mat_vec_mul_spec_matrix_vector (m n : ℕ) (c : Fin m → Fin n → ℝ) (v : Fin n → ℝ) : Spec.matVecMulSpec (TorchLean.Tensor.matrix c) (TorchLean.Tensor.dim fun (j : Fin n) => TorchLean.Tensor.scalar (v j)) = TorchLean.Tensor.dim fun (i : Fin m) => TorchLean.Tensor.scalar (∑ j : Fin n, c i j * v j)

General matrix-vector multiplication for Tensor.matrix and a vector written as Tensor.dim.

This generalizes the one-row dot-product lemma from UniversalApproximation.lean to arbitrary m.

noncomputable def NN.MLTheory.Proofs.ReLUMlpBridge.affineIdLayer1 {n : ℕ} (w : Fin n → ℝ) (b : ℝ) : Spec.LinearSpec ℝ n 2

First layer for exact affine representability.

Given an affine form $u(x)=w\mathbin{\cdot}x+b$, this layer outputs $[u(x),-u(x)]$.

noncomputable def NN.MLTheory.Proofs.ReLUMlpBridge.affineIdLayer2 : Spec.LinearSpec ℝ 2 1

Second layer for exact affine representability.

With hidden activations $[\operatorname{ReLU}(u),\operatorname{ReLU}(-u)]$, this output layer computes $\operatorname{ReLU}(u)-\operatorname{ReLU}(-u)=u$.

noncomputable def NN.MLTheory.Proofs.ReLUMlpBridge.stdBasis {n : ℕ} (i : Fin n) : Fin n → ℝ

Standard basis vector $e_i\in\mathbb{R}^n$.

theorem NN.MLTheory.Proofs.ReLUMlpBridge.dot_stdBasis {n : ℕ} (i : Fin n) (x : TorchLean.Tensor ℝ [n]) : dot (stdBasis i) x = x.getScalar i

$\operatorname{dot}(e_i,x)=x_i$ for the standard basis stdBasis.

theorem NN.MLTheory.Proofs.ReLUMlpBridge.mlp_eval_affine_id {n : ℕ} (w : Fin n → ℝ) (b : ℝ) (x : TorchLean.Tensor ℝ [n]) : mlpEval (affineIdLayer1 w b) affineIdLayer2 x = dot w x + b

Exact representability of affine maps by a 2-layer ReLU MLP (width 2).

This is the core bridge lemma that turns scalar affine forms $w\mathbin{\cdot}x+b$ into MLP evaluations.

theorem NN.MLTheory.Proofs.ReLUMlpBridge.mlp_eval_coord {n : ℕ} (i : Fin n) (x : TorchLean.Tensor ℝ [n]) : mlpEval (affineIdLayer1 (stdBasis i) 0) affineIdLayer2 x = x.getScalar i

Exact representability of coordinate projections $x\mapsto x_i$ by a width-$2$ ReLU MLP.

Ridge lifting

Given a scalar-input MLP (l1,l2) and an affine scalar map $u=w\mathbin{\cdot}x+c$, we build a tensor-input MLP whose pre-activations match the scalar-input pre-activations at $u$. This lets you reuse any scalar approximation result for functions of one affine form (“ridge functions”).

noncomputable def NN.MLTheory.Proofs.ReLUMlpBridge.liftScalarLayer1 {n hidDim : ℕ} (l1 : Spec.LinearSpec ℝ 1 hidDim) (w : Fin n → ℝ) (c : ℝ) : Spec.LinearSpec ℝ n hidDim

Lift a scalar-input first layer to a tensor-input first layer along a ridge direction.

Given a first layer that expects a scalar $u\in\mathbb{R}$, this constructs a tensor-input layer that feeds it $u=w\mathbin{\cdot}x+c$.

theorem NN.MLTheory.Proofs.ReLUMlpBridge.mlp_eval_lift_from_scalar {n hidDim : ℕ} (l1 : Spec.LinearSpec ℝ 1 hidDim) (l2 : Spec.LinearSpec ℝ hidDim 1) (w : Fin n → ℝ) (c : ℝ) (x : TorchLean.Tensor ℝ [n]) : mlpEval (liftScalarLayer1 l1 w c) l2 x = UniversalApproximation.mlpEvalScalar hidDim l1 l2 (dot w x + c)

Lifting lemma: the lifted tensor-input MLP agrees with the scalar-input MLP evaluated at $w\mathbin{\cdot}x+c$.