Better Neural Network Expressivity: Subdividing the Simplex

1 Preliminaries and Expressivity Setup

Definition 1 Binary maximum on real numbers
#

For real numbers \(a,b \in \mathbb {R}\), their binary maximum is \(\max (a,b)\), characterized by \(\max (a,b) \ge a\), \(\max (a,b) \ge b\), and if \(c \ge a\) and \(c \ge b\) then \(c \ge \max (a,b)\).

Definition 2 Affine maps
#

An affine map \(A : \mathbb {R}^m \to \mathbb {R}^n\) has the form \(A(x)=L(x)+b\) for a linear map \(L : \mathbb {R}^m \to \mathbb {R}^n\) and a vector \(b \in \mathbb {R}^n\).

Definition 3 The ReLU function

The rectified linear unit is the function \(\operatorname {ReLU}: \mathbb {R}\to \mathbb {R}\) defined by \(\operatorname {ReLU}(t)=\max \{ 0,t\} \). Applied to a vector, \(\operatorname {ReLU}\) acts coordinatewise.

Definition 4 ReLU network with \(k\) hidden layers

A ReLU network on \(\mathbb {R}^n\) with \(k\) hidden layers is specified by affine maps \(T^{(i)}:\mathbb {R}^{n_{i-1}}\to \mathbb {R}^{n_i}\) for \(i=1,\ldots ,k+1\), with \(n_0=n\) and \(n_{k+1}=1\). It computes

\[ T^{(k+1)}\circ \operatorname {ReLU}\circ T^{(k)}\circ \cdots \circ \operatorname {ReLU}\circ T^{(2)}\circ \operatorname {ReLU}\circ T^{(1)}. \]
Definition 5 Continuous piecewise linear functions
#

A function \(f:\mathbb {R}^n\to \mathbb {R}\) is continuous piecewise linear, abbreviated CPWL, if \(\mathbb {R}^n\) admits a finite polyhedral subdivision such that \(f\) restricts to an affine function on every cell of the subdivision and the restrictions agree continuously on common faces.

Lemma 6 ReLU networks compute CPWL functions

Every function computed by a ReLU network is CPWL.

Proof

Affine maps are CPWL. The scalar function \(t\mapsto \max \{ 0,t\} \) is affine on the two half-lines \(t\le 0\) and \(t\ge 0\), and the two affine pieces agree at \(0\), so it is CPWL. Coordinatewise application preserves finite polyhedral subdivisions by intersecting the preimages of the coordinate break hyperplanes. Finally, a composition of CPWL maps is CPWL after pulling back and refining the finitely many subdivisions. Applying these observations to each affine layer and ReLU layer in Definition 4 gives the claim.

Definition 7 \(\operatorname {CPWL}_n\) and \(\operatorname {ReLU}_{n,k}\)

Let \(\operatorname {CPWL}_n\) be the set of CPWL functions \(f:\mathbb {R}^n\to \mathbb {R}\). Let \(\operatorname {ReLU}_{n,k}\) be the subset of \(\operatorname {CPWL}_n\) consisting of functions representable by ReLU networks with \(k\) hidden layers.

Definition 8 The maximum function \(\operatorname {MAX}_n\)

For \(n\ge 1\), define

\[ \operatorname {MAX}_n(x_1,\ldots ,x_n)=\max \{ x_1,\ldots ,x_n\} . \]
Lemma 9 Binary maximum has one hidden layer

The function \(\operatorname {MAX}_2(x_1,x_2)=\max \{ x_1,x_2\} \) is representable by a ReLU network with one hidden layer.

Proof

For all \(x_1,x_2\in \mathbb {R}\),

\[ \max \{ x_1,x_2\} =x_1+\operatorname {ReLU}(x_2-x_1). \]

If \(x_2\le x_1\), the ReLU term is \(0\) and the expression equals \(x_1\). If \(x_2\ge x_1\), the ReLU term is \(x_2-x_1\) and the expression equals \(x_2\). This is exactly one affine input to a ReLU neuron, followed by one affine output map.

Lemma 10 Binary tree computation of \(\operatorname {MAX}_m\)

For every \(m\ge 1\), \(\operatorname {MAX}_m\) can be computed with \(\lceil \log _2 m\rceil \) hidden layers.

Proof

Pair the \(m\) inputs and compute the maximum of each pair using Lemma 9. If \(m\) is odd, pass the unpaired input forward by an affine coordinate. One layer reduces the number of live quantities from \(m\) to at most \(\lceil m/2\rceil \) while preserving their maximum. Repeating this reduction for \(\lceil \log _2 m\rceil \) rounds leaves one quantity, and induction over the rounds shows that the remaining quantity is \(\operatorname {MAX}_m\).

Theorem 11 Wang–Sun max-affine representation

For every \(f\in \operatorname {CPWL}_n\), there exist affine maps \(A_1,\ldots ,A_s:\mathbb {R}^n\to \mathbb {R}^{n+1}\) and signs \(\sigma _1,\ldots ,\sigma _s\in \{ \pm 1\} \) such that

\[ f(x)=\sum _{i=1}^s \sigma _i \operatorname {MAX}_{n+1}(A_i(x)). \]
Corollary 12 Classical logarithmic CPWL upper bound

Every function in \(\operatorname {CPWL}_n\) can be represented by a ReLU network with \(\lceil \log _2(n+1)\rceil \) hidden layers; equivalently \(\operatorname {CPWL}_n\subseteq \operatorname {ReLU}_{n,\lceil \log _2(n+1)\rceil }\).

Proof

Write \(f\) as the signed sum of \(\operatorname {MAX}_{n+1}\) applied to affine maps, as in Theorem 11. By Lemma 10, each \(\operatorname {MAX}_{n+1}\) term is computable with \(\lceil \log _2(n+1)\rceil \) hidden layers after the affine input map. Running the finitely many subnetworks in parallel and taking the signed affine sum in the output layer preserves the same depth. Thus \(f\) has the stated depth.

Conjecture 13 Hertrich–Basu–Di Summa–Skutella depth conjecture

The previously conjectured sharp bound was that \(\operatorname {MAX}_n\) requires \(\lceil \log _2 n\rceil \) hidden layers, and consequently that the bound in Corollary 12 is optimal for general CPWL functions.

Lemma 14 Known lower bound: \(\operatorname {MAX}_3\) is not one-hidden-layer

The function \(\operatorname {MAX}_3\) cannot be represented by a ReLU network with one hidden layer; it needs at least two hidden layers.