1 ReLU networks and piecewise linear functions
This chapter fixes the model of computation and records the standard closure properties of the classes \(\mathrm{ReLU}_{n,k}\), the classical representation theorem of Wang and Sun, and the resulting \(\lceil \log _2(n+1)\rceil \) upper bound of Arora, Basu, Mianjy, and Mukherjee.
The rectified linear unit is the function \(\mathrm{ReLU}\colon \mathbb {R} \to \mathbb {R}\) defined by \(\mathrm{ReLU}(x) = \max \{ 0, x\} \). Applied to a vector, \(\mathrm{ReLU}\) acts component-wise.
A ReLU network with \(k\) hidden layers (equivalently, of depth \(k+1\)) on \(\mathbb {R}^n\) is a tuple of affine transformations \(T^{(i)}\colon \mathbb {R}^{n_{i-1}} \to \mathbb {R}^{n_i}\) for \(i = 1, \dots , k+1\), where \(n_0 = n\) and \(n_{k+1} = 1\). The network computes the function
the alternating composition of the affine maps with component-wise applications of \(\mathrm{ReLU}\). The entries of the matrices and translation vectors of the \(T^{(i)}\) are the weights and biases of the network.
For integers \(n \ge 1\) and \(k \ge 0\), denote by \(\mathrm{ReLU}_{n,k}\) the set of all functions \(f\colon \mathbb {R}^n \to \mathbb {R}\) computed by some ReLU network with at most \(k\) hidden layers. (For \(k = 0\) these are exactly the affine functions.)
A function \(f\colon \mathbb {R}^n \to \mathbb {R}\) is continuous piecewise linear (CPWL) if \(f\) is continuous and there exist finitely many affine functions \(g_1, \dots , g_p \colon \mathbb {R}^n \to \mathbb {R}\) such that for every \(x \in \mathbb {R}^n\) there is an index \(i\) with \(f(x) = g_i(x)\). We write \(\mathrm{CPWL}_n\) for the set of all CPWL functions on \(\mathbb {R}^n\); it is a real vector space under pointwise operations.
For \(n \ge 1\), the function \(\mathrm{MAX}_n \colon \mathbb {R}^n \to \mathbb {R}\) is defined by
It is a member of \(\mathrm{CPWL}_n\).
For all \(n \ge 1\) and \(k \ge 0\) we have \(\mathrm{ReLU}_{n,k} \subseteq \mathrm{CPWL}_n\).
We first record a composition principle: if \(F\colon \mathbb {R}^n \to \mathbb {R}^m\) has all coordinate functions CPWL, witnessed by finite affine families \(G_1, \dots , G_m\), and \(f\colon \mathbb {R}^m \to \mathbb {R}\) is CPWL with affine family \(\{ g_i\} \), then \(f \circ F\) is CPWL: it is continuous as a composition of continuous maps, and at every point \(x\) its value equals \(g_i(F(x))\) for some \(i\), which is an affine combination of values of members of the \(G_j\) at \(x\); hence \(f\circ F\) agrees pointwise with a member of the finite affine family \(\{ \, c_0 + \sum _j c_j h_j : h_j \in G_j,\ (c_j)\ \text{a coefficient row of some } g_i\, \} \). Now argue by induction on \(k\). For \(k = 0\), affine functions are CPWL. For the step, every coordinate of \(\mathrm{ReLU}\circ T^{(k)} \circ \cdots \circ \mathrm{ReLU}\circ T^{(1)}\) is CPWL: coordinates of the inner composition are CPWL by induction (applied coordinate-wise, with output dimension absorbed into the affine maps), and \(\mathrm{ReLU} = \max \{ 0, x\} \) is CPWL, so the composition principle applies. A final affine map preserves CPWL.
The identity function \(x \mapsto x\) on \(\mathbb {R}\) lies in \(\mathrm{ReLU}_{1,1}\).
For every \(x \in \mathbb {R}\) we have \(x = \mathrm{ReLU}(x) - \mathrm{ReLU}(-x)\): if \(x \ge 0\) the right-hand side is \(x - 0\), otherwise it is \(0 - (-x)\). The right-hand side is a network with one hidden layer of two neurons.
For all \(n \ge 1\) and \(k \ge 0\) we have \(\mathrm{ReLU}_{n,k} \subseteq \mathrm{ReLU}_{n,k+1}\).
Let \(f\) be computed by \(T^{(k+1)} \circ \mathrm{ReLU} \circ \cdots \circ T^{(1)}\). Insert one further hidden layer after \(T^{(k+1)}\) that computes the pair \((\mathrm{ReLU}(y), \mathrm{ReLU}(-y))\) of the scalar output \(y\), followed by the affine map \((u, v) \mapsto u - v\). By Lemma 7 this leaves the computed function unchanged and uses \(k+1\) hidden layers.
Let \(n \ge 1\), \(k \ge 0\), let \(f_1, \dots , f_s \in \mathrm{ReLU}_{n,k}\) and \(c_1, \dots , c_s \in \mathbb {R}\). Then \(c_1 f_1 + \dots + c_s f_s \in \mathrm{ReLU}_{n,k}\).
Run networks for \(f_1, \dots , f_s\) in parallel: the \(i\)-th hidden layer of the combined network is the concatenation of the \(i\)-th hidden layers of the individual networks (a block-diagonal affine map followed by component-wise \(\mathrm{ReLU}\)), and the final affine map outputs \(\sum _j c_j y_j\) where \(y_j\) is the affine output of the \(j\)-th network, which is again affine in the last hidden layer.
Let \(f \in \mathrm{ReLU}_{m,k}\) and let \(A\colon \mathbb {R}^n \to \mathbb {R}^m\) be affine. Then \(f \circ A \in \mathrm{ReLU}_{n,k}\).
If \(f\) is computed with first affine layer \(T^{(1)}\), replace \(T^{(1)}\) by the affine map \(T^{(1)} \circ A\) and keep the remaining layers. For \(k = 0\) the composition of affine maps is affine.
Let \(g_1, \dots , g_m \in \mathrm{ReLU}_{n,k}\) and \(f \in \mathrm{ReLU}_{m,j}\). Then the function \(x \mapsto f(g_1(x), \dots , g_m(x))\) lies in \(\mathrm{ReLU}_{n,k+j}\).
Stack networks for \(g_1, \dots , g_m\) in parallel as in Lemma 9; after \(k\) hidden layers their outputs are an affine function \(G\) of the last hidden layer, with values \((g_1(x), \dots , g_m(x))\). Compose the first affine layer \(T^{(1)}\) of a network for \(f\) with \(G\): the merged map \(T^{(1)} \circ G\) is affine, so the total number of hidden layers is \(k + j\).
\(\mathrm{MAX}_2 \in \mathrm{ReLU}_{2,1}\).
For all \(x_1, x_2\),
the last two terms sum to \(x_1\), and \(x_1 + \mathrm{ReLU}(x_2 - x_1)\) equals \(x_2\) if \(x_2 \ge x_1\) and \(x_1\) otherwise. The right-hand side is a one-hidden-layer network with three neurons.
For every \(n \ge 1\) we have \(\mathrm{MAX}_n \in \mathrm{ReLU}_{n,\lceil \log _2 n \rceil }\).
Induction on \(n\). For \(n = 1\) the identity is affine, so \(\mathrm{MAX}_1 \in \mathrm{ReLU}_{1,0}\). For \(n \ge 2\) write
By induction and Lemma 10 (precomposing with coordinate projections), both inner maxima lie in \(\mathrm{ReLU}_{n, \lceil \log _2 n\rceil - 1}\), using \(\lceil \log _2 \lceil n/2 \rceil \rceil \le \lceil \log _2 n\rceil - 1\) and Lemma 8 to equalize depths. Composing with \(\mathrm{MAX}_2 \in \mathrm{ReLU}_{2,1}\) (Lemma 12) via Lemma 11 adds one hidden layer.
Let \(1 \le m \le n\) and \(k \ge 0\). If \(\mathrm{MAX}_n \in \mathrm{ReLU}_{n,k}\) then \(\mathrm{MAX}_m \in \mathrm{ReLU}_{m,k}\).
Let \(D\colon \mathbb {R}^m \to \mathbb {R}^n\) be the linear map \(D(x_1, \dots , x_m) = (x_1, \dots , x_m, x_m, \dots , x_m)\) repeating the last coordinate. Then \(\mathrm{MAX}_m = \mathrm{MAX}_n \circ D\), and Lemma 10 applies.
For every \(f \in \mathrm{CPWL}_n\) there exist \(s \in \mathbb {N}\), affine maps \(A_1, \dots , A_s \colon \mathbb {R}^n \to \mathbb {R}^{n+1}\) and signs \(\sigma _1, \dots , \sigma _s \in \{ \pm 1\} \) such that
For every \(n \ge 1\) we have \(\mathrm{CPWL}_n \subseteq \mathrm{ReLU}_{n, \lceil \log _2 (n+1) \rceil }\).
Let \(f \in \mathrm{CPWL}_n\) and write \(f = \sum _i \sigma _i\, \mathrm{MAX}_{n+1}\circ A_i\) as in Theorem 15. By Lemma 13, \(\mathrm{MAX}_{n+1} \in \mathrm{ReLU}_{n+1, \lceil \log _2(n+1)\rceil }\); by Lemma 10 each \(\mathrm{MAX}_{n+1}\circ A_i \in \mathrm{ReLU}_{n, \lceil \log _2(n+1)\rceil }\); by Lemma 9 the signed sum stays in this class.
\(\mathrm{MAX}_3 \notin \mathrm{ReLU}_{3,1}\); that is, two hidden layers are necessary to compute the maximum of three numbers.