Fast Convex Optimization for Two-Layer ReLU Networks: Equivalent Model Classes and Cone Decompositions

TL;DR

The work tackles the non-convex difficulty of training two-layer ReLU networks by introducing convex reformulations that recast the objective as a convex GLM with a group-$\ell_1$ penalty under polyhedral cone constraints. It builds two training routes: C-ReLU with sub-sampled activation patterns and C-GReLU, an unconstrained gated-ReLU formulation, showing exact equivalence in the zero-regularization regime and data-dependent approximation guarantees for nonzero regularization. Efficient solvers are developed: a proximal-gradient method named R-FISTA for the unconstrained form and an augmented Lagrangian scheme for the constrained reformulation, both designed to scale with first-order operations and GPU acceleration. Empirical results demonstrate faster convergence and competitive generalization on large-scale benchmarks, with an open-source implementation (scnn) to enable practical adoption of convex training for two-layer ReLU networks.

Abstract

We develop fast algorithms and robust software for convex optimization of two-layer neural networks with ReLU activation functions. Our work leverages a convex reformulation of the standard weight-decay penalized training problem as a set of group-$\ell_1$-regularized data-local models, where locality is enforced by polyhedral cone constraints. In the special case of zero-regularization, we show that this problem is exactly equivalent to unconstrained optimization of a convex "gated ReLU" network with non-singular gates. For problems with non-zero regularization, we show that convex gated ReLU models obtain data-dependent approximation bounds for the ReLU training problem. To optimize the convex reformulations, we develop an accelerated proximal gradient method and a practical augmented Lagrangian solver. We show that these approaches are faster than standard training heuristics for the non-convex problem, such as SGD, and outperform commercial interior-point solvers. Experimentally, we verify our theoretical results, explore the group-$\ell_1$ regularization path, and scale convex optimization for neural networks to image classification on MNIST and CIFAR-10.

Figures (19)

• Figure 1: Convex (solid line) and non-convex (dashed) optimization of a two-layer ReLU network for a realizable synthetic classification problem. We plot only one run of the convex solver since they are nearly identical and all reach perfect accuracy. In contrast, $4/10$ runs of SGD on the non-convex problem converge to sub-optimal stationary points.
• Figure 2: An illustration of the Cone Decomposition (CD) procedure: $u_i$ is decomposed onto the Minkowski difference $\mathcal{K}_i - \mathcal{K}_i$.
• Figure 3: Summary of equivalences between convex (blue) and non-convex (red) neural network training problems with gated ReLU (left) and ReLU (right) activations. The convex programs C-GReLU and C-ReLU are equivalent to the standard non-convex training problems NC-GReLU and NC-GReLU and are related to each other via the cone decomposition procedure.
• Figure 4: Performance profiles comparing (left) R-FISTA and MOSEK for the C-GReLU problem to Adam and SGD for NC-GReLU, and (right) the AL method and baselines for C-ReLU/NC-ReLU. A problem is solved when $\left(F(x_k) - F(x^*)\right)/F(x^*) \leq 1$, where $F(x^*)$ is the smallest objective value found by any method. This rule is method-independent as the convex and non-convex problems share the same optimal objective value. See Appendix \ref{['app:performance-profiles']} for alternative thresholds. Methods are judged by comparing time to a fixed proportion of problems solved (see dashed line at $50\%$). R-FISTA and the AL method solve a higher proportion of problems faster than the baselines.
• Figure 5: Effect of sampling activation patterns on test accuracy for networks trained using the C-ReLU and C-GReLU problems on the primary-tumor dataset. We consider a grid of regularization parameters and plot median (solid line) and first and third quartiles (shaded region) over 10 random samplings of $\tilde{\mathcal{D}}$, where $|\tilde{\mathcal{D}}|$ is limited to 10, 100, or 1000 patterns.

Theorems & Definitions (57)

• Theorem 2.1
• Theorem 2.2
• Proposition 3.0
• Proposition 3.0
• Theorem 3.1
• Proposition 3.1
• Proposition 3.1
• Proposition 3.1
• Theorem 3.2
• Theorem 4.1