Fixing the NTK: From Neural Network Linearizations to Exact Convex Programs

TL;DR

This paper bridges neural tangent kernels and convex reformulations of ReLU networks, showing how to optimize the NTK via multiple kernel learning to achieve exact convex solutions with improved predictive performance.

Contribution

It establishes a connection between NTK and MKL for gated ReLU networks, enabling exact convex optimization and improved kernel-based training.

Findings

NTK is equivalent to a specific MKL kernel with fixed mask weights.

Iterative reweighting yields the optimal MKL kernel matching the convex reformulation.

Numerical simulations support the theoretical connection and improvements.

Abstract

Recently, theoretical analyses of deep neural networks have broadly focused on two directions: 1) Providing insight into neural network training by SGD in the limit of infinite hidden-layer width and infinitesimally small learning rate (also known as gradient flow) via the Neural Tangent Kernel (NTK), and 2) Globally optimizing the regularized training objective via cone-constrained convex reformulations of ReLU networks. The latter research direction also yielded an alternative formulation of the ReLU network, called a gated ReLU network, that is globally optimizable via efficient unconstrained convex programs. In this work, we interpret the convex program for this gated ReLU network as a Multiple Kernel Learning (MKL) model with a weighted data masking feature map and establish a connection to the NTK. Specifically, we show that for a particular choice of mask weights that do not depend on the learning targets, this kernel is equivalent to the NTK of the gated ReLU network on the training data. A consequence of this lack of dependence on the targets is that the NTK cannot perform better than the optimal MKL kernel on the training set. By using iterative reweighting, we improve the weights induced by the NTK to obtain the optimal MKL kernel which is equivalent to the solution of the exact convex reformulation of the gated ReLU network. We also provide several numerical simulations corroborating our theory. Additionally, we provide an analysis of the prediction error of the resulting optimal kernel via consistency results for the group lasso.