Approximating multiplication with a 2-layer ReLU MLP (2D box)

This file gives a constructive, fully proved approximation result: on $[-M,M]^2$, the function $(x_0,x_1)\mapsto x_0x_1$ can be uniformly approximated by a single-hidden-layer ReLU MLP on Tensor ℝ [2].

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.firstCoordinate (x : TorchLean.Tensor ℝ [2]) : ℝ

First coordinate projection from a rank-one tensor with two entries.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.secondCoordinate (x : TorchLean.Tensor ℝ [2]) : ℝ

Second coordinate projection from a rank-one tensor with two entries.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.box (M : ℝ) : Set (TorchLean.Tensor ℝ [2])

The closed box domain $[-M,M]\times[-M,M]$.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.mulFun (x : TorchLean.Tensor ℝ [2]) : ℝ

The target multiplication map: multiply the two coordinates.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.wPlus : Fin 2 → ℝ

Ridge direction with dot wPlus x equal to the sum of the two coordinates.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.wMinus : Fin 2 → ℝ

Ridge direction with dot wMinus x equal to the first coordinate minus the second.

theorem NN.MLTheory.Proofs.ReLUMulApprox.dot_wPlus (x : TorchLean.Tensor ℝ [2]) : ReLUMlpBridge.dot wPlus x = firstCoordinate x + secondCoordinate x

Evaluate the ridge wPlus: it sums the two coordinates.

theorem NN.MLTheory.Proofs.ReLUMulApprox.dot_wMinus (x : TorchLean.Tensor ℝ [2]) : ReLUMlpBridge.dot wMinus x = firstCoordinate x - secondCoordinate x

Evaluate the ridge wMinus: $\operatorname{dot}(w_-,x)=x_0-x_1$.

theorem NN.MLTheory.Proofs.ReLUMulApprox.mul_identity (x y : ℝ) : x * y = ((x + y) * (x + y) - (x - y) * (x - y)) / 4

Algebraic identity expressing multiplication via a difference of squares.

theorem NN.MLTheory.Proofs.ReLUMulApprox.box_bounds {M : ℝ} (_hM : 0 ≤ M) {x : TorchLean.Tensor ℝ [2]} (hx : x ∈ box M) : firstCoordinate x ∈ Set.Icc (-M) M ∧ secondCoordinate x ∈ Set.Icc (-M) M

Unpack the defining bounds of membership in box M.

theorem NN.MLTheory.Proofs.ReLUMulApprox.sum_mem_Icc {M : ℝ} (_hM : 0 ≤ M) {x : TorchLean.Tensor ℝ [2]} (hx : x ∈ box M) : ReLUMlpBridge.dot wPlus x ∈ Set.Icc (-2 * M) (2 * M)

If $x\in\operatorname{box}(M)$, then $x_0+x_1\in[-2M,2M]$.

theorem NN.MLTheory.Proofs.ReLUMulApprox.diff_mem_Icc {M : ℝ} (_hM : 0 ≤ M) {x : TorchLean.Tensor ℝ [2]} (hx : x ∈ box M) : ReLUMlpBridge.dot wMinus x ∈ Set.Icc (-2 * M) (2 * M)

If $x\in\operatorname{box}(M)$, then $x_0-x_1\in[-2M,2M]$.

theorem NN.MLTheory.Proofs.ReLUMulApprox.square_lipschitz_Icc {R : ℝ} (_hR : 0 ≤ R) (x : ℝ) : x ∈ Set.Icc (-R) R → ∀ y ∈ Set.Icc (-R) R, |x * x - y * y| ≤ 2 * R * |x - y|

Lipschitz bound for square on $[-R,R]$: $|x^2-y^2|\leq 2R|x-y|$.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.appendDim {α : Type} [TorchLean.Storage α] {m n : ℕ} {s : Spec.Shape} (a : TorchLean.Tensor α (Spec.Shape.dim m s)) (b : TorchLean.Tensor α (Spec.Shape.dim n s)) : TorchLean.Tensor α (Spec.Shape.dim (m + n) s)

Concatenate tensors along the leading dimension.

In this file, this is used to append the hidden-unit vectors of two subnetworks.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.appendLinearSpec {inDim m n : ℕ} (a : Spec.LinearSpec ℝ inDim m) (b : Spec.LinearSpec ℝ inDim n) : Spec.LinearSpec ℝ inDim (m + n)

Append two first-layer linear specs by appending their weight and bias tensors.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.mat1Get {n : ℕ} (A : TorchLean.Tensor ℝ [1, n]) (j : Fin n) : ℝ

Extract the j-th entry from a 1 × n tensor interpreted as a row matrix.

theorem NN.MLTheory.Proofs.ReLUMulApprox.singleRowMatrix_get_matrix {n : ℕ} (f : Fin 1 → Fin n → ℝ) (j : Fin n) : mat1Get (TorchLean.Tensor.matrix f) j = f 0 j

mat1Get agrees with the Tensor.matrix constructor.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.combineOutput {m n : ℕ} (α β γ : ℝ) (a : Spec.LinearSpec ℝ m 1) (b : Spec.LinearSpec ℝ n 1) : Spec.LinearSpec ℝ (m + n) 1

Combine two scalar-output linear specs into one scalar-output spec on an appended hidden layer.

If the appended hidden vector is $[z_a;z_b]$, the resulting output layer computes $\gamma+\alpha\,\mathrm{out}_a(z_a)+\beta\,\mathrm{out}_b(z_b)$.

theorem NN.MLTheory.Proofs.ReLUMulApprox.getScalar_append_left {m n : ℕ} (a : TorchLean.Tensor ℝ [m]) (b : TorchLean.Tensor ℝ [n]) (i : Fin m) : (appendDim a b).getScalar (Fin.castAdd n i) = a.getScalar i

Reading the left component from an appended hidden vector.

theorem NN.MLTheory.Proofs.ReLUMulApprox.getScalar_append_right {m n : ℕ} (a : TorchLean.Tensor ℝ [m]) (b : TorchLean.Tensor ℝ [n]) (i : Fin n) : (appendDim a b).getScalar (Fin.natAdd m i) = b.getScalar i

Reading the right component from an appended hidden vector.

theorem NN.MLTheory.Proofs.ReLUMulApprox.getScalar_relu {n : ℕ} (z : TorchLean.Tensor ℝ [n]) (i : Fin n) : (Activation.reluSpec z).getScalar i = UniversalApproximation.relu (z.getScalar i)

Pointwise behavior of the ReLU activation on tensor-vectors.

theorem NN.MLTheory.Proofs.ReLUMulApprox.mat_vec_mul_spec_oneRow {n : ℕ} (A : TorchLean.Tensor ℝ [1, n]) (v : TorchLean.Tensor ℝ [n]) : Spec.matVecMulSpec A v = TorchLean.Tensor.dim fun (x : Fin 1) => TorchLean.Tensor.scalar (∑ j : Fin n, mat1Get A j * v.getScalar j)

Matrix-vector multiplication for a 1 × n matrix produces a single scalar coordinate.

theorem NN.MLTheory.Proofs.ReLUMulApprox.mlp_eval_nd_eq_bias_sum {inDim hidDim : ℕ} (l1 : Spec.LinearSpec ℝ inDim hidDim) (l2 : Spec.LinearSpec ℝ hidDim 1) (x : TorchLean.Tensor ℝ [inDim]) : ReLUMlpBridge.mlpEval l1 l2 x = UniversalApproximation.extractScalarOutput l2.bias + ∑ j : Fin hidDim, mat1Get l2.weights j * UniversalApproximation.relu ((Spec.linearSpec l1 x).getScalar j)

Expand mlp_eval_nd into “bias + sum over hidden units” form.

This is the main normalization lemma used to prove that appendLinearSpec together with combineOutput implements affine combinations of subnetworks.

theorem NN.MLTheory.Proofs.ReLUMulApprox.getScalar_linear_spec_append_left {inDim m n : ℕ} (l1a : Spec.LinearSpec ℝ inDim m) (l1b : Spec.LinearSpec ℝ inDim n) (x : TorchLean.Tensor ℝ [inDim]) (i : Fin m) : (Spec.linearSpec (appendLinearSpec l1a l1b) x).getScalar (Fin.castAdd n i) = (Spec.linearSpec l1a x).getScalar i

Selecting the left block of a linear spec appended via appendLinearSpec.

theorem NN.MLTheory.Proofs.ReLUMulApprox.getScalar_linear_spec_append_right {inDim m n : ℕ} (l1a : Spec.LinearSpec ℝ inDim m) (l1b : Spec.LinearSpec ℝ inDim n) (x : TorchLean.Tensor ℝ [inDim]) (i : Fin n) : (Spec.linearSpec (appendLinearSpec l1a l1b) x).getScalar (Fin.natAdd m i) = (Spec.linearSpec l1b x).getScalar i

Selecting the right block of a linear spec appended via appendLinearSpec.

theorem NN.MLTheory.Proofs.ReLUMulApprox.mlp_eval_append_linear {inDim m n : ℕ} (l1a : Spec.LinearSpec ℝ inDim m) (l1b : Spec.LinearSpec ℝ inDim n) (l2a : Spec.LinearSpec ℝ m 1) (l2b : Spec.LinearSpec ℝ n 1) (α β γ : ℝ) (x : TorchLean.Tensor ℝ [inDim]) : ReLUMlpBridge.mlpEval (appendLinearSpec l1a l1b) (combineOutput α β γ l2a l2b) x = γ + α * ReLUMlpBridge.mlpEval l1a l2a x + β * ReLUMlpBridge.mlpEval l1b l2b x

Appending hidden units and wiring the output with combineOutput yields an affine combination.

Concretely, the combined network computes: $\gamma+\alpha\,\mathrm{net}_a(x)+\beta\,\mathrm{net}_b(x)$.

theorem NN.MLTheory.Proofs.ReLUMulApprox.relu_mul_universal_approximation_box {M : ℝ} (hM : 0 < M) (ε : ℝ) : ε > 0 → ∃ (hidDim : ℕ) (l1 : Spec.LinearSpec ℝ 2 hidDim) (l2 : Spec.LinearSpec ℝ hidDim 1), ∀ x ∈ box M, |mulFun x - ReLUMlpBridge.mlpEval l1 l2 x| < ε

Uniform approximation of multiplication on $[-M,M]^2$ by a single-hidden-layer ReLU MLP.

The construction follows the classical reduction $xy=((x+y)^2-(x-y)^2)/4$, combined with a one-dimensional ReLU approximator for square on $[-2M,2M]$ that is lifted along the ridge directions wPlus and wMinus.