When learning an energy, it is often worth building its geometry into the model. A convex objective has no spurious local minima on a convex domain, and many algorithms can exploit that structure. In quadratic optimal transport, the optimal map is the gradient of a convex potential when the source measure has a density and both measures have finite second moments. These are good reasons to use neural architectures that represent convex functions by construction [7], [8].
One architecture is especially transparent: compute several affine functions, then take their maximum. To approximate a convex target, choose these functions as supporting planes; their maximum forms a polyhedral surface beneath the graph. Replacing the maximum by a log-sum-exp gives a smooth surface. We will separate three questions: can these networks approximate every convex function, how many planes are needed, and what does smoothing cost?
A network made of supporting planes
For vectors \(a_i\in\mathbb R^d\) and offsets \(b_i\in\mathbb R\), set
\[\ell_i(x)=a_i\cdot x+b_i,\qquad M_N(x)=\max_{1\le i\le N}\ell_i(x).\]Here a two-layer max network means one affine feature layer with \(N\) outputs, followed by one scalar maximum. This is a maxout unit [3], not a sum of coordinatewise ReLUs. The weights \(a_i,b_i\) can have either sign: the final maximum guarantees convexity.
The geometric starting point is classical convex analysis: a proper lower semicontinuous convex function is the supremum of its affine minorants, the affine functions lying below it [1]. Compactness turns this infinite description into a finite approximation.
Theorem 1 (universality of max-affine networks). Let \(K\subset\mathbb R^d\) be nonempty, compact and convex, and let \(f:K\to\mathbb R\) be continuous and convex. For every \(\varepsilon>0\), there is a max-affine network \(M_N\) such that
\[0\le f(x)-M_N(x)<\varepsilon \qquad\text{for every }x\in K.\]Proof: one good plane near each point, then compactness
Extend \(f\) by \(+\infty\) outside \(K\). This is a proper lower semicontinuous convex function, so the affine-minorant representation applies. Geometrically, it comes from separating points below the graph from the closed convex epigraph. For each \(z\in K\), choose an affine function \(\ell_z\le f\) on \(K\) with
\[f(z)-\ell_z(z)<\varepsilon.\]By continuity, this inequality persists in a relative neighborhood of \(z\). These neighborhoods cover \(K\); compactness gives a finite subcover, associated with planes \(\ell_1,\ldots,\ell_N\). Their maximum is below \(f\) everywhere, and at every point at least one plane is within \(\varepsilon\) of \(f\). This proves the result.
No differentiability or interior assumption is needed. In particular, a plane need not touch the graph exactly at a boundary point: an arbitrarily small gap is enough.
How many planes are needed?
Write \(\Vert g\Vert_\infty=\sup_{x\in K}\vert g(x)\vert\). Theorem 1 gives no bound on \(N\). If the target is smooth, tangent planes have a quadratic local error, and a grid gives an explicit rate. An \(h\)-net means that every point of \(K\) is within distance \(h\) of one of the chosen points.
Theorem 2 (a constructive smooth rate). Suppose that \(f\) is convex and continuously differentiable on a neighborhood of the compact convex set \(K\), with \(\beta\)-Lipschitz gradient on \(K\). If \(z_1,\ldots,z_N\in K\) form an \(h\)-net, the tangent-plane maximum satisfies
\[0\le f(x)-M_N(x)\le \frac{\beta h^2}{2}, \qquad M_N(x)=\max_i\bigl[f(z_i)+\nabla f(z_i)\cdot(x-z_i)\bigr].\]In particular, on \(K=[0,1]^d\), use a grid with \(m+1\) equally spaced points per coordinate, where \(m\ge1\). Then \(N=(m+1)^d\), \(h\le\sqrt d/(2m)\), and
\[\|f-M_N\|_\infty\le \frac{\beta d}{8m^2} =O(N^{-2/d}).\]Proof: integrate the gradient along a segment
For \(z,x\in K\), convexity gives the lower bound, while
\[\begin{aligned} f(x)-f(z)-\nabla f(z)\cdot(x-z) &=\int_0^1 \bigl[\nabla f(z+t(x-z))-\nabla f(z)\bigr]\cdot(x-z)\,dt\\ &\le\int_0^1\beta t\|x-z\|^2\,dt =\frac{\beta}{2}\|x-z\|^2. \end{aligned}\]Choose a point of the net at distance at most \(h\), and use that \(M_N\) is at least its tangent plane.
The exponent reflects the dimension: reducing the error by a factor of four requires roughly \(2^d\) times as many planes. For a merely \(L\)-Lipschitz convex function on a neighborhood of \(K\), bounded subgradients give the coarser estimate \(2Lh\), hence \(L\sqrt d/m=O(N^{-1/d})\) on this grid.
Why bounded slopes give the Lipschitz estimate
At each \(z_i\), choose \(p_i\in\partial f(z_i)\), with \(\|p_i\|\le L\), and use \(\ell_i(x)=f(z_i)+p_i\cdot(x-z_i)\). If \(\|x-z_i\|\le h\), then
\[0\le f(x)-M_N(x) \le f(x)-f(z_i)-p_i\cdot(x-z_i)\le2Lh.\]This nonsmooth estimate is a bound for a uniform grid, not an optimal rate for the whole Lipschitz convex class. Adaptive plane placement can improve it; this distinction is important in max-affine convex regression [4].
The smooth exponent \(2/d\) cannot generally be improved by a clever choice of planes. Even a quadratic already forces it.
Proposition 1 (the smooth exponent is sharp). For \(q(x)=\|x\|^2/2\) on \([0,1]^d\), every maximum of \(N\) affine functions satisfies
\[\|q-M_N\|_\infty\ge \frac{1}{16N^{2/d}}.\]Proof: an affine patch cannot be too large
Write \(\varepsilon=\|q-M_N\|_\infty\). Intersect each affine branch’s active region with the cube. This is a convex set on which \(M_N\) is affine. For two points \(x,y\) in the same such set, the affine midpoint terms cancel, giving
\[\frac{\|x-y\|^2}{8} =\frac{q(x)+q(y)}2-q\!\left(\frac{x+y}2\right) \le 2\varepsilon.\]Thus each active region has diameter at most \(4\sqrt\varepsilon\), and is contained in a coordinate box of volume at most \((4\sqrt\varepsilon)^d\). The \(N\) regions cover the unit cube, so \(1\le N(4\sqrt\varepsilon)^d\). Rearranging proves the claim. Notice that we did not assume the approximant lies below the target.
From max to softmax
The smooth analogue is the log-sum-exp (LSE) network, at temperature \(\tau>0\):
\[S_{N,\tau}(x)= \tau\log\!\left(\sum_{i=1}^N e^{\ell_i(x)/\tau}\right).\]We use “softmax layer” for this scalar smooth maximum. Strictly speaking, the usual vector-valued softmax gives the normalized exponential weights, which appear in its gradient below. The scalar LSE is smooth and convex [2]; its universal approximation properties have been studied explicitly by Calafiore, Gaubert and Possieri [5].
Lemma 1 (smoothing error). For every \(x\),
\[M_N(x)\le S_{N,\tau}(x)\le M_N(x)+\tau\log N.\]Proof: subtract the largest exponent
Set \(a=M_N(x)\). Then
\[S_{N,\tau}(x)=a+\tau\log \sum_{i=1}^N e^{(\ell_i(x)-a)/\tau}.\]Every term is at most \(1\), and at least one equals \(1\). The sum therefore lies between \(1\) and \(N\). This shifted expression is also the numerically stable way to evaluate LSE [9].
Theorem 3 (smooth universality and rate). Allowing both \(N\) and \(\tau\) to vary, LSE networks are uniformly dense in the continuous convex functions on a nonempty compact convex set. More precisely, if affine minorants give \(0\le f-M_N\le\delta_N\), then
\[-\delta_N\le S_{N,\tau}-f\le\tau\log N, \qquad \|S_{N,\tau}-f\|_\infty \le\max\{\delta_N,\tau\log N\}.\]Thus the smooth grid rate \(O(N^{-2/d})\) is preserved by choosing \(\tau\log N\le\delta_N\), for example \(\tau=O(N^{-2/d}/\log N)\) when \(N>1\).
Proof: separate the two errors
The two-sided estimate follows by combining \(f-\delta_N\le M_N\le f\) with Lemma 1. The maximum, rather than the sum, of the two error bounds suffices because the errors have opposite signs. For universality, use Theorem 1 with error \(\varepsilon/2\), then choose \(\tau\log N<\varepsilon/2\). If \(N=1\), LSE is exactly the single affine branch, independently of \(\tau\).
The normalization convention matters: replacing the sum by its average gives \(S_{N,\tau}-\tau\log N\). This is just a constant shift, absorbable into every bias \(b_i\), so it does not change the approximation class at fixed \(N,\tau\). It does change the graph when the same tangent planes are used without refitting.
Temperature, curvature, and the size of the weights
The estimate above proves universality by decreasing temperature as width grows; the bound \(\tau\log N\) alone proves nothing at fixed temperature. A separate, quantitative obstruction appears if both the temperature and the permitted slopes are fixed: there is then not enough curvature to resolve a corner.
Define \(w_i(x)=e^{\ell_i(x)/\tau}/\sum_j e^{\ell_j(x)/\tau}\) and \(\bar a(x)=\sum_iw_i(x)a_i\). Differentiating gives
\[\nabla S_{N,\tau}=\bar a,\qquad \nabla^2 S_{N,\tau} =\frac1\tau\sum_iw_i(a_i-\bar a)(a_i-\bar a)^\top.\]If \(\|a_i\|\le A\), the covariance formula implies \(0\preceq\nabla^2S_{N,\tau}\preceq (A^2/\tau)I\). Consequently, fixed temperature and uniformly bounded slopes prevent arbitrarily accurate approximation of a corner, no matter how large \(N\) is.
A quantitative example: approximating the absolute value
Suppose \(0\le s''\le H\) on \([-1,1]\) and \(\Vert s-\vert\,\cdot\,\vert\Vert_\infty\le\varepsilon\). For \(0<h\le1\),
\[2h-4\varepsilon \le s(h)+s(-h)-2s(0)\le Hh^2.\]The first inequality uses the uniform error at three points; the second follows by integrating \(s''\). Taking \(h=1/(1+H)\) yields
\[\varepsilon\ge\frac{h}{2}-\frac{Hh^2}{4} \ge\frac{1}{4(1+H)}.\]For an LSE network with slopes bounded by \(A\), use \(H=A^2/\tau\). This is a joint restriction on temperature and slopes, not a non-universality claim about fixed temperature with unrestricted weights.
There is also a parameterization issue. Writing the network as \(\tau\,\mathrm{LSE}((Ax+b)/\tau)\), the weights and biases entering the standard LSE activation scale like \(1/\tau\). The physical slopes \(a_i\) of the planes may stay bounded even while these internal parameters grow. Any theorem imposing bounds on weights must specify which convention it uses.
Numerical experimentation
The one-dimensional panel illustrates the tangent-plane construction for \(f(x)=x^2/2\). In two dimensions, we instead learn the affine pieces of the max and softmax networks. The target is \(f(x)=x^\top Qx/2\), with
\[Q=\begin{pmatrix}1.5&0.4\\0.4&0.8\end{pmatrix}.\]Both domains are \([-1,1]^d\). The width parameter \(n\) sets \(n\) neurons in the 1D reference and \(N=n^2\) neurons in each 2D network. The 2D planes start from nonuniform, reproducible random initializations. Adam [10] then optimizes their slopes and offsets to minimize sampled mean squared error:
\[\mathcal L(\theta)=\frac1B\sum_{j=1}^B \bigl(F_\theta(z_j)-f(z_j)\bigr)^2.\]The max and LSE models are fitted independently, rather than sharing a fixed set of supporting planes. Convexity is preserved by their architectures, but the fitted functions can lie on either side of the target. The displays report sampled errors and validation losses, not certificates of global optimality.
Training details
Each model takes 2,000 Adam steps. A batch contains 256 stratified random interior points and 32 random boundary points. A separate 33 × 33 validation grid selects the best checkpoint; the 201 × 201 display grid is evaluated afterward. Fixed seeds make the fits reproducible.
For conditioning, hard-max and low-temperature fits use free plane centers and free intercept offsets. This is an invertible reparameterization of the affine coefficients, not a constraint that the planes remain tangent to the quadratic. Every slope and intercept is allowed to change.
Learning the affine pieces
2D · Trained max network
2D · Trained softmax network
Errors are maxima on a fine evaluation grid, separate from the training samples, not certified suprema. The two heatmaps share a color scale, recomputed after each fit. In error mode, blue means underestimation and red overestimation. In value mode, black contours show the target, white dashed contours the fitted network.
3D surface view
The learned approximations shown as surfaces: the max has planar faces, while softmax is smooth. Both panels use the same camera and height scale. Adjust the neuron count and temperature above, or rotate the view below.
3D · Trained max network
3D · Trained softmax network
Teal and orange surfaces show the learned networks, with studio lighting and reflective highlights to reveal facets and curvature. Both panels use the same lighting. Heights are function values, not errors; compare with the target using the 2D contours above. The max shows its planar faces exactly; softmax uses a fine 201 × 201 grid and analytic normals, without smoothing the fitted functions.
The automatic setting uses \(\tau=1/(2n)\); manual mode lets you explore other temperatures and refits the softmax network. At very low temperature, an LSE can genuinely approach a polyhedral shape even though it is smooth. The fine mesh and analytic normals reduce rendering artifacts without artificially rounding the surface. The smoothing lemma compares max and LSE built from the same planes; the two independently trained 2D models need not satisfy that comparison.
What this buys us—and what it does not
A max-affine epigraph is described by linear inequalities, which is convenient in optimization. Its smooth LSE counterpart provides an explicit gradient and a positive semidefinite Hessian. Both are universal convex models, and their rates expose the usual dependence on dimension. Differences of two LSE networks can in turn represent much richer, nonconvex functions [6].
Uniform approximation of function values should not be confused with uniform approximation of gradients. Near a genuine corner, the target does not even have a unique gradient. For transport maps or differentiable optimization, the relevant regularity and derivative estimates deserve a separate analysis.
References
- R. Tyrrell Rockafellar. Convex Analysis. Princeton University Press, 1970.
- Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, 2004.
- Ian Goodfellow, David Warde-Farley, Mehdi Mirza, Aaron Courville and Yoshua Bengio. Maxout Networks. ICML, PMLR 28(3):1319–1327, 2013.
- Gábor Balázs, András György and Csaba Szepesvári. Near-optimal max-affine estimators for convex regression. AISTATS, PMLR 38:56–64, 2015.
- Giuseppe C. Calafiore, Stéphane Gaubert and Corrado Possieri. Log-Sum-Exp Neural Networks and Posynomial Models for Convex and Log-Log-Convex Data. IEEE Transactions on Neural Networks and Learning Systems, 31(3):827–838, 2020. DOI.
- Giuseppe C. Calafiore, Stéphane Gaubert and Corrado Possieri. A Universal Approximation Result for Difference of Log-Sum-Exp Neural Networks. IEEE Transactions on Neural Networks and Learning Systems, 31(12):5603–5612, 2020.
- Brandon Amos, Lei Xu and J. Zico Kolter. Input Convex Neural Networks. ICML, PMLR 70:146–155, 2017.
- Ashok Makkuva, Amirhossein Taghvaei, Sewoong Oh and Jason Lee. Optimal transport mapping via input convex neural networks. ICML, PMLR 119:6672–6681, 2020.
- Pierre Blanchard, Desmond J. Higham and Nicholas J. Higham. Accurately computing the log-sum-exp and softmax functions. IMA Journal of Numerical Analysis, 41(4):2311–2330, 2021. DOI.
- Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. ICLR, 2015.