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 ℝ (.dim 2 .scalar).

@[reducible, inline]

abbrev NN.MLTheory.Proofs.ReLUMulApprox.PlaneTensorVec : Type

TensorVec specialized to the 2D (rank-2) tensor-vector shape.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.firstCoordinate (x : PlaneTensorVec) : ℝ

First coordinate projection for PlaneTensorVec.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.secondCoordinate (x : PlaneTensorVec) : ℝ

Second coordinate projection for PlaneTensorVec.

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.box (M : ℝ) : Set PlaneTensorVec

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

noncomputable def NN.MLTheory.Proofs.ReLUMulApprox.mulFun (x : PlaneTensorVec) : ℝ

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 : PlaneTensorVec) : 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 : PlaneTensorVec) : 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 : PlaneTensorVec} (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 : PlaneTensorVec} (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 : PlaneTensorVec} (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} {m n : ℕ} {s : Spec.Shape} (a : Spec.Tensor α (Spec.Shape.dim m s)) (b : Spec.Tensor α (Spec.Shape.dim n s)) : Spec.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 : Spec.Tensor ℝ (Spec.Shape.dim 1 (Spec.Shape.dim n Spec.Shape.scalar))) (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_matrixMN {n : ℕ} (f : Fin 1 → Fin n → ℝ) (j : Fin n) : mat1Get (Spec.matrixMN 1 n fun (i : Fin 1) (j : Fin n) => f i j) j = f 0 j

mat1_get agrees with the matrixMN 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.vec_get_append_left {m n : ℕ} (a : Spec.Tensor ℝ (Spec.Shape.dim m Spec.Shape.scalar)) (b : Spec.Tensor ℝ (Spec.Shape.dim n Spec.Shape.scalar)) (i : Fin m) : ReLUMlpBridge.vecGet (appendDim a b) (Fin.castAdd n i) = ReLUMlpBridge.vecGet a i

Reading the left component from an appended hidden vector.

theorem NN.MLTheory.Proofs.ReLUMulApprox.vec_get_append_right {m n : ℕ} (a : Spec.Tensor ℝ (Spec.Shape.dim m Spec.Shape.scalar)) (b : Spec.Tensor ℝ (Spec.Shape.dim n Spec.Shape.scalar)) (i : Fin n) : ReLUMlpBridge.vecGet (appendDim a b) (Fin.natAdd m i) = ReLUMlpBridge.vecGet b i

Reading the right component from an appended hidden vector.

theorem NN.MLTheory.Proofs.ReLUMulApprox.vec_get_relu {n : ℕ} (z : Spec.Tensor ℝ (Spec.Shape.dim n Spec.Shape.scalar)) (i : Fin n) : ReLUMlpBridge.vecGet (Activation.reluSpec z) i = UniversalApproximation.relu (ReLUMlpBridge.vecGet z i)

Pointwise behavior of the ReLU activation on tensor-vectors.

theorem NN.MLTheory.Proofs.ReLUMulApprox.mat_vec_mul_spec_oneRow {n : ℕ} (A : Spec.Tensor ℝ (Spec.Shape.dim 1 (Spec.Shape.dim n Spec.Shape.scalar))) (v : Spec.Tensor ℝ (Spec.Shape.dim n Spec.Shape.scalar)) : Spec.matVecMulSpec A v = Spec.Tensor.dim fun (x : Fin 1) => Spec.Tensor.scalar (∑ j : Fin n, mat1Get A j * ReLUMlpBridge.vecGet v 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 : Spec.Tensor ℝ (Spec.Shape.dim inDim Spec.Shape.scalar)) : ReLUMlpBridge.mlpEvalNd l1 l2 x = UniversalApproximation.extractScalarOutput l2.bias + ∑ j : Fin hidDim, mat1Get l2.weights j * UniversalApproximation.relu (ReLUMlpBridge.vecGet (Spec.linearSpec l1 x) 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.vec_get_linear_spec_append_left {inDim m n : ℕ} (l1a : Spec.LinearSpec ℝ inDim m) (l1b : Spec.LinearSpec ℝ inDim n) (x : Spec.Tensor ℝ (Spec.Shape.dim inDim Spec.Shape.scalar)) (i : Fin m) : ReLUMlpBridge.vecGet (Spec.linearSpec (appendLinearSpec l1a l1b) x) (Fin.castAdd n i) = ReLUMlpBridge.vecGet (Spec.linearSpec l1a x) i

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

theorem NN.MLTheory.Proofs.ReLUMulApprox.vec_get_linear_spec_append_right {inDim m n : ℕ} (l1a : Spec.LinearSpec ℝ inDim m) (l1b : Spec.LinearSpec ℝ inDim n) (x : Spec.Tensor ℝ (Spec.Shape.dim inDim Spec.Shape.scalar)) (i : Fin n) : ReLUMlpBridge.vecGet (Spec.linearSpec (appendLinearSpec l1a l1b) x) (Fin.natAdd m i) = ReLUMlpBridge.vecGet (Spec.linearSpec l1b x) 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 : Spec.Tensor ℝ (Spec.Shape.dim inDim Spec.Shape.scalar)) : ReLUMlpBridge.mlpEvalNd (appendLinearSpec l1a l1b) (combineOutput α β γ l2a l2b) x = γ + α * ReLUMlpBridge.mlpEvalNd l1a l2a x + β * ReLUMlpBridge.mlpEvalNd 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.mlpEvalNd 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.