Fixing the NTK: From Neural Network Linearizations to Exact Convex
  Programs
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.
Rajat Vadiraj Dwaraknath, Tolga Ergen, Mert Pilanci
# Fixing the NTK: From Neural Network Linearizations to Exact Convex Programs

###### 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.

## 1 Introduction

Neural Networks (NNs) have become popular in various machine learning applications due to their remarkable modeling capabilities and generalization performance. However, their highly nonlinear and non-convex structure precludes an effective theoretical analysis. Therefore, developing theoretical tools to understand the fundamental mechanisms behind neural networks is still an active research topic. To tackle this problem, [1] studied the training dynamics of neural networks trained with Stochastic Gradient Descent (SGD) in a regime where each layer has infinitely many neurons and SGD uses an infinitesimally small learning rate, i.e., gradient flow. Thus, they related the training dynamics of neural networks to the training dynamics of a fixed kernel called the Neural Tangent Kernel (NTK). However, [2] showed that neurons barely move from their initial values in this regime so that neural networks fail to learn useful features from the training data. This is in contrast to their finite width counterparts, which are able to learn predictive features in practice [3]. Moreover, [4; 5] provided further theoretical and empirical evidence to show that existing kernel approaches are not able to explain the remarkable performance of finite width networks. Therefore, although NTK and
similar kernel based approaches enable theoretical analysis unlike standard finite width networks, they fail to explain the effectiveness of finite width neural networks that are employed in practice.

Recently a series of papers [6; 7; 8; 9; 10; 11; 12; 13; 14; 15] introduced an analytic framework to analyze finite width neural networks by leveraging certain convex duality arguments. Particularly, they showed that the standard regularized non-convex training problem can be equivalently cast as a finite dimensional convex program. This convex approach has two major advantages over standard non-convex training: **(1)** Since the training objective is convex, one can find globally optimal parameters of the network efficiently and reliably unlike standard nonconvex training which can get stuck at a local minimum, and **(2)** As we show in this work, a class of convex reformulations can be interpreted as an instance of Multiple Kernel Learning (MKL) [16] which allows us to characterize the corresponding finite width networks by a learned data-dependent kernel that can be iteratively computed. This is in contrast to the infinite-width kernel characterization in which the NTK stays constant throughout training.

**Notation and Preliminaries.** In the paper, we use lowercase and uppercase bold letters to denote vectors and matrices respectively. We also use subscripts to denote a certain column or element. We denote the identity matrix of size \(k\times k\) as \(\mathbf{I}_{k}\). To denote the set \(\{1,2,\ldots,n\}\), we use \([n]\). We also use \(\left\lVert\cdot\right\rVert_{p}\) and \(\left\lVert\cdot\right\rVert_{F}\) to represent the standard \(\ell_{p}\) and Frobenius norms. Additionally, we denote 0-1 valued indicator function and ReLU activation as \(\mathbb{1}\)\(\{x\geq 0\}\) and \(\left(x\right)_{+}:=\max\{x,0\}\), respectively.

In this paper, we focus on analyzing the regularized training problem of ReLU networks. Particularly, we consider a two-layer ReLU network with \(m\) neurons whose output function is defined as follows

\[f(\mathbf{x},\boldsymbol{\theta})\!:=\sum_{j=1}^{m}\left(\mathbf{x}^{T}\mathbf{ w}_{j}^{(1)}\right)_{+}w_{j}^{(2)}=\sum_{j=1}^{m}\left(\mathbb{1}\left\{ \mathbf{x}^{T}\mathbf{w}_{j}^{(1)}\geq 0\right\}\mathbf{x}^{T}\mathbf{w}_{j}^{(1)} \right)w_{j}^{(2)}. \tag{1}\]

where \(\mathbf{w}_{j}^{(1)}\in\mathbb{R}^{d}\) and \(w_{j}^{(2)}\) are the \(j^{th}\) hidden and output layer weights, respectively and \(\boldsymbol{\theta}:=\{(\mathbf{w}_{j}^{(1)},w_{j}^{(2)})\}_{j=1}^{m}\) represents all trainable parameters. Given a training data matrix \(\mathbf{X}\in\mathbb{R}^{n\times d}\) and a target vector \(\mathbf{y}\in\mathbb{R}^{n}\), we minimize the following weight decay regularized training objective

\[\min_{\mathbf{W}^{(1)},\mathbf{w}^{(2)}}\left\lVert\sum_{j=1}^{m}\left( \mathbf{X}\mathbf{w}_{j}^{(1)}\right)_{+}w_{j}^{(2)}-\mathbf{y}\right\rVert_{ 2}^{2}+\lambda\sum_{j=1}^{m}\left(\left\lVert\mathbf{w}_{j}^{(1)}\right\rVert_ {2}^{2}+|w_{j}^{(2)}|^{2}\right), \tag{2}\]

where \(\lambda>0\) is the regularization coefficient. We also use \(f(\mathbf{X},\boldsymbol{\theta})=\sum_{j=1}^{m}\left(\mathbf{X}\mathbf{w}_{j} ^{(1)}\right)_{+}w_{j}^{(2)}\) for convenience. We discuss extensions to generic loss and deeper architectures in appendix.

**Our Contributions.**

* In section 4, we show that the convex formulation of the _Gated ReLU_ network is equivalent to _Multiple Kernel Learning_ with a specific set of _Masking Kernels_. (Theorem 4.3).
* In section 5, we connect this formulation to the Neural Tangent Kernel by showing that the NTK is a specific weighted combination of masking kernels (Theorem 5.1).
* In Corollary 5.2, we show that on the training set, the **NTK is suboptimal when compared to the optimal kernel learned by MKL**, which is equivalent to the model learnt by our convex Gated ReLU program.
* In section 6, we also derive bounds on the prediction error of this optimal kernel and specify how to choose the regularization parameter \(\lambda\) (Theorem 6.1).

## 2 Convex Optimization and the NTK

Here, we briefly review the literature on convex training and the NTK theory of neural networks.

### Convex Programs for ReLU Networks

Even though the network in (1) has only two layers, previous studies show that (2) is a challenging optimization problem due to the non-convexity of the objective function. Thus, local search heuristics such as SGD might fail to globally optimize the training objective [17; 18; 19; 20]. To eliminate the issues
associated with the inherent non-convexity, [7] introduced an exact convex reformulation of (2) as the following constrained optimization problem

\[\min_{\mathbf{w}_{i},\mathbf{w}_{i}^{\prime}}\left\|\sum_{i=1}^{p}\mathbf{D}_{i} \mathbf{X}(\mathbf{w}_{i}-\mathbf{w}_{i}^{\prime})-\mathbf{y}\right\|_{2}^{2}+ \lambda\sum_{i=1}^{p}\left(\left\|\mathbf{w}_{i}\right\|_{2}+\left\|\mathbf{w }_{i}^{\prime}\right\|_{2}\right)\text{ s.t. }\begin{array}{ll}&(2\mathbf{D}_{i}-\mathbf{I})\mathbf{X}\mathbf{w}_{i} \geq\mathbf{0}\\ &(2\mathbf{D}_{i}-\mathbf{I})\mathbf{X}\mathbf{w}_{i}^{\prime}\geq\mathbf{0} \end{array},\forall i, \tag{3}\]

where \(\mathbf{D}_{i}\in\mathcal{D}_{\mathbf{X}}\) are \(n\times n\) binary masking diagonal matrices given by \(p\!:=\left|\mathcal{D}_{\mathbf{X}}\right|\) and

\[\mathcal{D}_{\mathbf{X}}\!:=\left\{\mathrm{diag}\left(\mathbb{1}\left\{ \mathbf{X}\mathbf{u}\geq\mathbf{0}\right\}\right):\mathbf{u}\in\mathbb{R}^{d }\right\}. \tag{4}\]

The mask set \(\mathcal{D}_{\mathbf{X}}\) can be interpreted as the set of all possible ways to separate the training data \(\mathbf{X}\) by a hyperplane passing through the origin. With these masks, we can characterize a single ReLU activated neuron on the training data as follows: \(\left(\mathbf{X}\mathbf{w}_{i}\right)_{+}=\mathrm{diag}\left(\mathbb{1}\left\{ \mathbf{X}\mathbf{w}_{i}\geq\mathbf{0}\right\}\right)\mathbf{X}\mathbf{w}_{i} =\mathbf{D}_{i}\mathbf{X}\mathbf{w}_{i}\) provided that \((2\mathbf{D}_{i}-\mathbf{I}_{n})\mathbf{X}\mathbf{w}_{i}\geq 0\). Therefore, by enforcing these cone constraints in (3), we maintain the masking property of the ReLU activation and parameterize the neural network as a linear function of the weights, thus make the learning problem convex. We refer the reader to [6] for more details.

Although (3) is convex and therefore eliminates the drawbacks associated with the non-convexity of (2), it might still be computationally complex to solve. Precisely, a worst-case upper-bound on the number of variables is \(\mathcal{O}(n^{r})\), where \(r=\mathrm{rank}(\mathbf{X})\leq\min\{n,d\}\). Although this is still significantly better than brute-force search over \(2^{mn}\) ReLU patterns, it could still be exponential in the dimension \(d\). To mitigate this issue, [8] proposed a relaxation of (1) called the _gated ReLU_ network

\[f_{\mathcal{G}}\left(\mathbf{x},\boldsymbol{\theta}\right)\!:=\sum_{j=1}^{m} \left(\mathbb{1}\left\{\mathbf{x}^{T}\mathbf{g}_{j}\geq\mathbf{0}\right\} \mathbf{x}^{T}\mathbf{w}_{j}^{(1)}\right)w_{j}^{(2)}, \tag{5}\]

where \(\mathcal{G}\!:=\{\mathbf{g}_{j}\}_{j=1}^{m}\) is a set of _gate_ vectors that are also optimized throughout training. Then, the corresponding non-convex learning problem is as follows

\[\min_{\mathbf{W}^{(1)},\mathbf{w}^{(2)},\mathcal{G}}\left\|\sum_{j=1}^{m} \mathrm{diag}\left(\mathbb{1}\left\{\mathbf{X}\mathbf{g}_{j}\geq\mathbf{0} \right\}\right)\mathbf{X}\mathbf{w}_{j}^{(1)}w_{j}^{(2)}-\mathbf{y}\right\|_{2 }^{2}+\lambda\sum_{j=1}^{m}\left(\left\|\mathbf{w}_{j}^{(1)}\right\|_{2}^{2}+ |w_{j}^{(2)}|^{2}\right). \tag{6}\]

By performing this relaxation, we decouple the dependence between the indicator function and the linear term in the exact ReLU network (1). To express the equivalent convex optimization problem corresponding to the gated ReLU network, we introduce the notion of complete gate sets.

**Definition 2.1**.: _A gate set \(\mathcal{G}\) is complete with respect to a dataset \(\mathbf{X}\) if the corresponding set of hyperplane arrangements covers all possible arrangement patterns for \(\mathbf{X}\) defined in (4), i.e.,_

\[\left\{\mathrm{diag}\left(\mathbb{1}\left\{\mathbf{X}\mathbf{g}_{j}\geq \mathbf{0}\right\}\right):\mathbf{g}_{j}\in\mathcal{G}\right\}=\mathcal{D}_{ \mathbf{X}}.\]

_Additionally, \(\mathcal{G}\) is minimally complete if \(\left|\mathcal{G}\right|=\left|\mathcal{D}_{\mathbf{X}}\right|=p\)._

Now, with this relaxation, [8] showed that the optimal value for (6) can always be achieved by choosing \(\mathcal{G}\) to be a complete gate set. Therefore, by working only with complete gate sets, we can modify (6) to only optimize over the network parameters \(\mathbf{W}^{(1)}\) and \(\mathbf{w}^{(2)}\)

\[\min_{\mathbf{W}^{(1)},\mathbf{w}^{(2)}}\left\|\sum_{j=1}^{m}\mathrm{diag} \left(\mathbb{1}\left\{\mathbf{X}\mathbf{g}_{j}\geq\mathbf{0}\right\}\right) \mathbf{X}\mathbf{w}_{j}^{(1)}w_{j}^{(2)}-\mathbf{y}\right\|_{2}^{2}+\lambda \sum_{j=1}^{m}\left(\left\|\mathbf{w}_{j}^{(1)}\right\|_{2}^{2}+|w_{j}^{(2)}|^ {2}\right). \tag{7}\]

Additionally, [21] also showed that we can set \(m=p\) without loss of generality. Then, [8] showed that the equivalent convex optimization problem for (6) and also (7) in the complete gate set setting is

\[\min_{\mathbf{w}_{i}}\left\|\sum_{i=1}^{p}\mathbf{D}_{i}\mathbf{X}\mathbf{w}_{i }-\mathbf{y}\right\|_{2}^{2}+\lambda\sum_{i=1}^{p}\left\|\mathbf{w}_{i}\right\| _{2}. \tag{8}\]

Notice that (8) is a least squares problem with _group lasso_ regularization [22]. Therefore, this relaxation for (3) can be efficiently optimized via convex optimization solvers. Furthermore, [8] proved that after solving the relaxed problem (8), one can construct an equivalent ReLU network from a gated ReLU network via a convex optimization based procedure called _cone decomposition_. We discuss the computational complexity of this approach in Section E of the supplementary material.
### The Neural Tangent Kernel

Previous works [1; 2; 5; 23; 24] characterized the training dynamics of SGD with infinitesimally small learning rates on neural networks in the _infinite-width_ limit, i.e., as \(m\rightarrow\infty\), via the NTK. [25; 26] also present analyses in this regime via a mean-field approach. In this section, we provide a brief overview of this theory and refer the reader to [1; 27; 28] for more details. The main idea behind the NTK theory is to approximate the neural network model function \(f\left(\mathbf{x},\boldsymbol{\theta}\right)\) by _linearizing it_ with respect to the parameters \(\boldsymbol{\theta}\) around its initialization \(\boldsymbol{\theta}_{0}\):

\[\hat{f}\left(\mathbf{x},\boldsymbol{\theta}\right)\approx f\left(\mathbf{x},\boldsymbol{\theta}_{0}\right)+\nabla_{\boldsymbol{\theta}}f\left(\mathbf{x },\boldsymbol{\theta}_{0}\right)^{T}\left(\boldsymbol{\theta}-\boldsymbol{ \theta}_{0}\right).\]

The authors of [1] show that if \(f\) is a neural network (with appropriately scaled output), and parameters \(\boldsymbol{\theta}\) initialized as i.i.d standard Gaussians, the linearization \(\hat{f}\) better approximates \(f\) in the infinite-width limit. We can interpret the linearized model \(\hat{f}\) as a kernel method with a feature map given by \(\boldsymbol{\phi}\left(\mathbf{x}\right)=\nabla_{\boldsymbol{\theta}}f\left( \mathbf{x},\boldsymbol{\theta}_{0}\right)\). The corresponding kernel induced by this feature map is termed as the NTK. Note that this is a random kernel since it depends on the random initialization of the parameters denoted as \(\boldsymbol{\theta}_{0}\). The main result of [1] in the simplified case of two-layer neural networks is that, in the infinite-width limit this kernel approaches a fixed deterministic limit given by \(H\big{(}\mathbf{x},\mathbf{x}^{\prime}\big{)}:=\mathbb{E}\big{[}\nabla_{ \boldsymbol{\theta}}f\big{(}\mathbf{x},\hat{\boldsymbol{\theta}}_{0}\big{)}^ {T}\nabla_{\boldsymbol{\theta}}f\big{(}\mathbf{x}^{\prime},\hat{\boldsymbol{ \theta}}_{0}\big{)}\big{]}\), where \(\hat{\boldsymbol{\theta}}\) corresponds to the parameters of a single neuron. Furthermore, [1] show that in this infinite limit, SGD with an infinitesimally small learning rate is equivalent to performing kernel regression with the fixed NTK.

To link the convex formulation (8) with NTK theory, we first present a scaled version of the gated ReLU network in (5) as follows

\[\tilde{f}_{\psi}\left(\mathbf{x},\boldsymbol{\theta}\right) :=\frac{1}{\sqrt{2m}}\sum_{j=1}^{m}\left(\mathbb{I}\left\{\mathbf{x}^{T} \mathbf{g}_{j}\geq\boldsymbol{0}\right\}\mathbf{x}^{T}\mathbf{w}_{j}^{(1)} \right)w_{j}^{(2)}. \tag{9}\]

In the next lemma, we provide the infinite width NTK of the scaled gated ReLU network in (9).

**Lemma 2.2**.: 2 _The infinite width NTK of the gated ReLU network (9) with i.i.d gates sampled as \(\mathbf{g}_{j}\sim\mathcal{N}\left(\mathbf{0},\mathbf{I_{d}}\right)\) and randomly initialized parameters as \(\mathbf{w}_{j}^{(1)}\sim\mathcal{N}\left(\mathbf{0},\mathbf{I_{d}}\right)\) and \(w_{j}^{(2)}\sim\mathcal{N}\left(0,1\right)\) is_

Footnote 2: All the proofs and derivations are presented in the supplementary material.

\[H\left(\mathbf{x},\mathbf{x}^{\prime}\right) :=\frac{1}{2\pi}\left(\pi-\arccos\left(\frac{\mathbf{x}^{T} \mathbf{x}^{\prime}}{\left\|\mathbf{x}\right\|_{2}\left\|\mathbf{x}^{\prime} \right\|_{2}}\right)\right)\mathbf{x}^{T}\mathbf{x}^{\prime}. \tag{10}\]

Additionally, we introduce a reparameterization of the standard ReLU network (1) with \(\boldsymbol{\theta}:=\left\{\left(\mathbf{w}_{j}^{(+)},\mathbf{w}_{j}^{(-)} \right)\right\}_{j=1}^{m}\), with \(\mathbf{w}_{j}^{(+)}\), \(\mathbf{w}_{j}^{(-)}\in\mathbb{R}^{d}\) which can still represent all the functions that (1) can

\[f_{r}(\mathbf{x},\boldsymbol{\theta}) :=\frac{1}{\sqrt{2m}}\sum_{j=1}^{m}\left(\mathbf{x}^{T}\mathbf{w} _{j}^{(+)}\right)_{+}-\left(\mathbf{x}^{T}\mathbf{w}_{j}^{(-)}\right)_{+}. \tag{11}\]

**Lemma 2.3**.: _The gated ReLU network (9) and the reparameterized ReLU network (11) have the same infinite width NTK given by Lemma 2.2._

Next, we present an equivalence between the gated ReLU network and the MKL model [16; 29].

## 3 Multiple Kernel Learning and Group Lasso

The Multiple Kernel Learning (MKL) model [16; 29] is an extension of the standard kernel method that learns an optimal data-dependent kernel as a convex combination of a set of fixed kernels and then performs regression with this learned kernel. We provide a brief overview of the MKL setting based on the exposition in [30]. Consider a set of \(p\) kernels given by corresponding feature maps \(\boldsymbol{\phi}_{i}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d_{i}}\). Given \(n\) training samples \(\mathbf{X}\in\mathbb{R}^{n\times d}\) with targets \(\mathbf{y}\in\mathbb{R}^{n}\), we define the feature matrices on this data by stacking the feature vectors as \(\boldsymbol{\Phi}_{i}:=[\boldsymbol{\phi}_{i}(\mathbf{x}_{1})^{T};\ldots; \boldsymbol{\phi}_{i}(\mathbf{x}_{n})^{T}]\in\mathbb{R}^{n\times d_{i}}\). Then, the corresponding \(n\times n\) kernel matrices are given by \(\mathbf{K}_{i}:=\boldsymbol{\Phi}_{i}\boldsymbol{\Phi}_{i}^{T}\). A convex combination of these kernels can be written as \(\mathbf{K}\left(\boldsymbol{\eta}\right):=\sum_{i=1}^{p}\eta_{i}\mathbf{K}_{i}\) where \(\boldsymbol{\eta}\in\Delta_{p}:=\{\boldsymbol{\eta}:\mathbf{1}^{T}\boldsymbol{ \eta}=1,\ \boldsymbol{\eta}\geq\boldsymbol{0}\}\) is a set of weights in the unit simplex. By noticing that the feature map corresponding to \(\mathbf{K}\left(\boldsymbol{\eta}\right)\)
is obtained by taking a weighted concatenation of \(\mathbf{\phi}_{i}\) with weights \(\sqrt{\eta_{i}}\), we can write the MKL optimization problem in terms of the feature matrices as

\[\min_{\mathbf{\eta}\in\Delta_{p},\mathbf{\mathbf{v}}_{i}\in\mathbb{R}^{d_{i}}}\left\| \sum_{i=1}^{p}\sqrt{\eta_{i}}\mathbf{\Phi}_{i}\mathbf{v}_{i}-\mathbf{y}\right\|_{2} ^{2}+\hat{\lambda}\sum_{i=1}^{p}\left\|\mathbf{v}_{i}\right\|_{2}^{2}, \tag{12}\]

where \(\hat{\lambda}>0\) is a regularization coefficient. For a set of fixed weights \(\mathbf{\eta}\), the optimal objective value of the kernel regression problem over \(\mathbf{v}\) is proportional to \(\mathbf{y}^{T}(\sum_{i=1}^{p}\eta_{i}\mathbf{K}_{i}+\hat{\lambda}\mathbf{I}_{ \mathbf{n}})^{-1}\mathbf{y}\) up to constant factors [16]. Thus, the MKL problem can be equivalently written as the following problem

\[\min_{\mathbf{\eta}\in\Delta_{p}}\quad\mathbf{y}^{T}\left(\mathbf{K}\left(\mathbf{ \eta}\right)+\hat{\lambda}\mathbf{I}_{\mathbf{n}}\right)^{-1}\mathbf{y}. \tag{13}\]

In this formulation, we can interpret MKL as finding the optimal data-dependent kernel that can be expressed as a convex combination of the fixed kernels given by \(\mathbf{K}_{i}\). In the next section, we link this kernel learning formulation with the convex group lasso problem in (8).

### Equivalence to Group Lasso

We first show that the MKL problem in (12) can be equivalently stated as a group lasso problem.

**Lemma 3.1** ([16; 29]).: _The MKL problem (12) is equivalent to the following kernel regression problem using a uniform combination of the fixed kernels with squared group lasso regularization where the groups are given by parameters corresponding to each feature map_

\[\min_{\mathbf{w}_{i}\in\mathbb{R}^{d_{i}}}\quad\left\|\sum_{i=1}^{p}\mathbf{\Phi} _{i}\mathbf{w}_{i}-\mathbf{y}\right\|_{2}^{2}+\hat{\lambda}\left(\sum_{i=1}^{ p}\left\|\mathbf{w}_{i}\right\|_{2}\right)^{2}. \tag{14}\]

We now present a short derivation of this equivalence. Using the variational formulation of the squared group \(\ell_{1}\)-norm [31]

\[\left(\sum_{i=1}^{p}\left\|\mathbf{w}_{i}\right\|_{2}\right)^{2}=\min_{\mathbf{ \eta}\in\Delta_{p}}\sum_{i=1}^{p}\frac{\left\|\mathbf{w}_{i}\right\|_{2}^{2} }{\eta_{i}},\]

we can rewrite the group lasso problem (14) as a joint minimization problem over both the parameters \(\mathbf{w}\) and regularization weights \(\mathbf{\eta}\) as follows

\[\min_{\mathbf{\eta}\in\Delta_{p}}\quad\min_{\mathbf{w}_{i}\in\mathbb{R}^{d_{i}}} \quad\left\|\sum_{i=1}^{p}\mathbf{\Phi}_{i}\mathbf{w}_{i}-\mathbf{y}\right\|_{2} ^{2}+\hat{\lambda}\sum_{i=1}^{p}\frac{\left\|\mathbf{w}_{i}\right\|_{2}^{2}}{ \eta_{i}}.\]

Finally, with a change of variables given by \(\mathbf{v}_{i}=\mathbf{w}_{i}/\sqrt{\eta_{i}}\), we recover the MKL problem (12). We note that the MKL problem (12) is also equivalent to the following standard group lasso problem

\[\min_{\mathbf{w}_{i}\in\mathbb{R}^{d_{i}}}\quad\left\|\sum_{i=1}^{p}\mathbf{\Phi} _{i}\mathbf{w}_{i}-\mathbf{y}\right\|_{2}^{2}+\lambda\sum_{i=1}^{p}\left\| \mathbf{w}_{i}\right\|_{2}. \tag{15}\]

This is due to the fact that squared and standard group lasso problems have the same regularization paths [31], so (14) and (15) are equivalent when \(\hat{\lambda}=\frac{\lambda}{\sum_{i=1}^{p}\left\|\mathbf{w}_{i}^{*}\right\|_{ 2}}\), where \(\mathbf{w}^{*}\) is the solution to (14).

### Solving Group Lasso by Iterative Reweighting

Previously, we used a variational formulation of the squared group \(\ell_{1}\)-norm to show equivalences to MKL. Now, we present the Iteratively Reweighted Least Squares (IRLS) algorithm [32; 33; 34; 35] to solve the group lasso problem (15) using the following variational formulation of the group \(\ell_{1}\)-norm [34]

\[\sum_{i=1}^{p}\left\|\mathbf{w}_{i}\right\|_{2}=\min_{\mathbf{\eta}\in\mathbb{R}_{ +}^{p}}\frac{1}{2}\sum_{i=1}^{p}\left(\frac{\left\|\mathbf{w}_{i}\right\|_{2}^{ 2}}{\eta_{i}}+\eta_{i}\right).\]
Based on this, we rewrite the group lasso problem (15) as the following minimization problem

\[\min_{\mathbf{\eta}\in\mathbb{R}_{+}^{p}}\min_{\mathbf{w}_{i}\in\mathbb{R}^{d_{i}}} \quad\left\|\sum_{i=1}^{p}\mathbf{\Phi}_{i}\mathbf{w}_{i}-\mathbf{y}\right\|_{2}^{ 2}+\frac{\lambda}{2}\sum_{i=1}^{p}\left(\frac{\left\|\mathbf{w}_{i}\right\|_{2} ^{2}}{\eta_{i}}+\eta_{i}\right).\]

Since the objective is jointly convex in \((\mathbf{\eta},\mathbf{w})\), it can be solved using alternating minimization [33]. Particularly, note that the inner minimization problem in \(\mathbf{w}_{i}\)'s is simply a \(\ell_{2}\) regularized least squares problem with different regularization strengths for each group and this can be solved in closed form

\[\min_{\mathbf{w}_{i}\in\mathbb{R}^{d_{i}}}\quad\left\|\sum_{i=1}^{p}\mathbf{\Phi} _{i}\mathbf{w}_{i}-\mathbf{y}\right\|_{2}^{2}+\lambda\sum_{i=1}^{p}\frac{ \left\|\mathbf{w}_{i}\right\|_{2}^{2}}{\eta_{i}}. \tag{16}\]

The outer problem in \(\mathbf{\eta}\) is also directly solved by setting \(\eta_{i}=\left\|\mathbf{w}_{i}\right\|_{2}\)[34]. To avoid convergence issues and instability around \(\eta_{i}=0\), we approximate the reweighting by adding a small positive constant \(\epsilon\). We use this procedure to solve the group lasso formulation of the gated ReLU network (8) by setting \(\mathbf{\Phi}_{i}=\mathbf{D}_{i}\mathbf{X}\). A detailed description is provided Algorithm 1. For further details regarding convergence, we refer the reader to [32, 33, 34].

```
1:Set iteration count \(k\gets 0\)
2:Initialize weights \(\eta_{i}^{(0)}\)
3:Set \(\mathbf{\Phi}_{i}\!:=\!\mathbf{D}_{i}\mathbf{X},\;\forall\;\mathbf{D}_{i}\in \mathcal{D}_{\mathbf{X}}\)
4:while not converged and \(k\leq\text{max}\) iteration count do
5: Solve the weighted \(\ell_{2}\) regularized least squares problem: \[\left\{\mathbf{w}_{i}^{(k)}\right\}_{i}=\operatorname*{argmin}_{\left\{ \mathbf{w}_{i}\right\}_{i}}\left\|\sum_{i=1}^{p}\mathbf{\Phi}_{i}\mathbf{w}_{i}- \mathbf{y}\right\|_{2}^{2}+\lambda\sum_{i=1}^{p}\frac{\left\|\mathbf{w}_{i} \right\|_{2}^{2}}{\eta_{i}^{(k)}}\]
6: Update the weights: \(\eta_{i}^{(k+1)}=\sqrt{\left\|\mathbf{w}_{i}^{(k)}\right\|_{2}+\epsilon}\)
7: Increment iteration count: \(k\gets k+1\)
8:endwhile
9:Optional: Convert the gated ReLU network to a ReLU network (see Section E for details)
```

**Algorithm 1** Iteratively Reweighted Least Squares (IRLS) for gated ReLU and ReLU networks

## 4 Gated ReLU as MKL with Masking Kernels

Motivated by the MKL interpretation of group lasso, we return to the convex reformulation (8) of the gated ReLU network. Notice that this problem has the same structure as the MKL equivalent group lasso problem (15) with a specific set of feature maps that we define below.

**Definition 4.1**.: _The masking feature maps \(\mathbf{\phi}_{j}:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}\) generated by a fixed set of gates \(\mathcal{G}\) are defined as \(\mathbf{\phi}_{j}(\mathbf{x})=\mathbb{1}\left\{\mathbf{x}^{T}\mathbf{g}_{j}\geq 0 \right\}\mathbf{x}\)._

These feature maps can be interpreted as simply passing the input unchanged if it lies in the positive halfspace of the corresponding gate vector \(\mathbf{g}_{j}\), i.e., \(\mathbf{x}^{T}\mathbf{g}_{j}\geq 0\), and returning zero if the input does not lie in this halfspace. Since \(\operatorname{diag}\left(\mathbb{1}\left\{\mathbf{X}\mathbf{g}_{j}\geq \mathbf{0}\right\}\right)\in\mathcal{D}_{\mathbf{X}},\;\forall\;\mathbf{g}_{j} \in\mathcal{G}\) holds for an arbitrary gate set \(\mathcal{G}\), we can conveniently express the corresponding feature matrices of these masking feature maps on the data \(\mathbf{X}\) in terms of fixed diagonal data masks as \(\mathbf{\Phi}_{j}=\mathbf{D}_{j}\mathbf{X}\), where \(\mathbf{D}_{j}=\operatorname{diag}\left(\mathbb{1}\left\{\mathbf{X}\mathbf{g} _{j}\geq\mathbf{0}\right\}\right)\). Similarly, the corresponding masking kernel matrices take the form \(\mathbf{K}_{j}=\mathbf{D}_{j}\mathbf{X}\mathbf{X}^{T}\mathbf{D}_{j}\). Note that for an arbitrary set of gates, the generated masking feature matrices on \(\mathbf{X}\) may not cover the entire set of possible masks \(\mathcal{D}_{\mathbf{X}}\). Additionally, multiple gate vectors can result in identical masks if \(\operatorname{diag}\left(\mathbb{1}\left\{\mathbf{X}\mathbf{g}_{i}\geq \mathbf{0}\right\}\right)=\operatorname{diag}\left(\mathbb{1}\left\{\mathbf{X }\mathbf{g}_{j}\geq\mathbf{0}\right\}\right)\) for \(i\neq j\) leading to degenerate feature matrices. However, if we work with minimally complete gate sets as defined in Definition 2.1, we can rectify these issues.

**Lemma 4.2**.: _For a minimally complete gate set \(\mathcal{G}\) defined in Definition 2.1, we can uniquely associate a gate vector \(\mathbf{g}_{i}\) to each data mask \(\mathbf{D}_{i}\in\mathcal{D}_{\mathbf{X}}\) such that \(\mathbf{D}_{i}=\operatorname{diag}\left(\mathbb{1}\left\{\mathbf{X}\mathbf{g} _{i}\geq\mathbf{0}\right\}\right),\forall i\in[p]\)._
Consequently, for minimally complete gate sets \(\mathcal{G}\), the generated masking feature matrices \(\forall i\in[p]\) can be expressed as \(\mathbf{\Phi}_{i}=\mathbf{D}_{i}\mathbf{X}\) and the masking kernel matrices take the form \(\mathbf{K}_{i}=\mathbf{D}_{i}\mathbf{X}\mathbf{X}^{T}\mathbf{D}_{i}\). In the context of the gated ReLU problem (7), since \(\mathcal{G}\) is complete, we can replace it with a minimally complete subset of \(\mathcal{G}\) without loss of generality since [21] showed that increasing \(m\) beyond \(p\) cannot reduce the value of the regularized training objective in (7). We are now ready to combine the MKL-group lasso equivalence with the convex reformulation of the gated ReLU network to present the following characterization of the nonconvex gated ReLU learning problem.

**Theorem 4.3**.: _The non-convex gated ReLU problem (7) with a minimally complete gate set \(\mathcal{G}\) is equivalent to performing multiple kernel learning (12) with the masking feature maps generated by \(\mathcal{G}\)_

\[\min_{\boldsymbol{\eta}\in\Delta_{p},\mathbf{v}_{i}\in\mathbb{R}^{d}}\quad \left\|\sum_{i=1}^{p}\sqrt{\eta_{i}}\mathbf{D}_{i}\mathbf{X}\mathbf{v}_{i}- \mathbf{y}\right\|_{2}^{2}+\hat{\lambda}\sum_{i=1}^{p}\left\|\mathbf{v}_{i} \right\|_{2}^{2}.\]

This theorem implies that the gated ReLU network finds the optimal combination of linear models restricted to the different masked datasets \(\mathbf{D}_{i}\mathbf{X}\) generated by the gates. By optimizing with all possible data maskings, we obtain the best possible gated ReLU network. From the kernel perspective, we have characterized the problem of finding an optimal finite width gated ReLU network as learning a data-dependent kernel and then performing kernel regression. This is in contrast to the NTK theory where the training of an infinite width network by gradient flow is characterized by regression with a constant kernel that is not learned from data. We further explore this connection below.

## 5 NTK as a Weighted Masking Kernel

We now connect the NTK of a gated ReLU network with the masking kernels generated by its gates.

**Theorem 5.1**.: _Let \(\mathbf{K}_{\mathcal{G}}\left(\boldsymbol{\tilde{\eta}}\right)\in\mathbb{R}^{n \times n}\) be the weighted masking kernel obtained by taking a convex combination of the masking feature maps generated by a minimally complete gate set \(\mathcal{G}\) with weights given by \(\tilde{\eta}_{i}=\mathbb{P}[\operatorname{diag}\left(\mathbb{1}\left\{ \mathbf{X}\mathbf{h}\geq\mathbf{0}\right\}\right)=\mathbf{D}_{i}]\) where \(\mathbf{h}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{\mathbf{d}})\) and let \(\mathbf{H}\in\mathbb{R}^{n\times n}\) be the infinite width NTK of the gated ReLU network (10) evaluated on the training data, i.e., the \(ij^{th}\) entry of \(\mathbf{H}\) is defined as \(\mathbf{H}_{ij}\!:=H(\mathbf{x}_{i},\mathbf{x}_{j})\). Then, \(\mathbf{K}_{\mathcal{G}}\left(\boldsymbol{\tilde{\eta}}\right)=\mathbf{H}\)._

A rough sketch of the proof of this theorem is to express the matrix \(\mathbf{H}\) as an expectation of indicator random variables using the definition of the NTK [1]. Then, by conditioning on the event that these indicators equal the masks \(\mathbf{D}_{i}\), we can express the NTK as a convex combination of the masking kernels \(\mathbf{K}_{i}\). The weights end up being precisely the probabilities that are described in Theorem 5.1. A detailed proof is provided in the supplementary material.

This theorem implies that the outputs of the gated ReLU network obtained via (16) with regularization weights \(\boldsymbol{\tilde{\eta}}\) on the training data is identical to that of kernel ridge regression with the NTK.

**Corollary 5.2**.: _Let \(\tilde{\mathbf{w}}\) be the solution to (16) with feature matrices \(\mathbf{\Phi}_{i}=\mathbf{D}_{i}\mathbf{X}\) and regularization weights \(\tilde{\eta}_{i}=\mathbb{P}[\operatorname{diag}\left(\mathbb{1}\left\{ \mathbf{X}\mathbf{h}\geq\mathbf{0}\right\}\right)=\mathbf{D}_{i}]\) where \(\mathbf{h}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{\mathbf{d}})\) and let \(\tilde{\mathbf{y}}=\mathbf{H}(\mathbf{H}+\lambda\mathbf{I}_{\mathbf{n}})^{-1} \mathbf{y}\) be the outputs of kernel ridge regression with the NTK on the training data. Then, \(\sum_{i=1}^{p}\mathbf{D}_{i}\mathbf{X}\tilde{\mathbf{w}}_{i}=\tilde{\mathbf{y}}\)._

Since \(\mathcal{G}\) is minimally complete, the weights in Theorem 5.1 satisfy \(\sum_{i=1}^{p}\tilde{\eta}_{i}=1\). In other words, \(\boldsymbol{\tilde{\eta}}\in\Delta_{p}\). Therefore, we can interpret Theorem 5.1 as follows - the NTK evaluated on the training data lies in the convex hull of all possible masking kernels of the training data. So \(\mathbf{H}\) lies in the feasible set of kernels for MKL using these masking kernels. By Theorem 4.3, we can find the optimal kernel in this set by solving the group lasso problem (8) of the gated ReLU network. **Therefore, we can interpret solving** (8) **as fixing the NTK by learning an improved data-dependent kernel.**

**Remark 5.3** (Suboptimality of NTK).: _Note that the weights \(\boldsymbol{\tilde{\eta}}\) do depend on the training data \(\mathbf{X}\), but do not depend on the target labels \(\mathbf{y}\). Since MKL learns the optimal kernel using both \(\mathbf{X}\) and \(\mathbf{y}\), the NTK still cannot perform better than the optimal MKL kernel on the training set. Thus, we fix the NTK._

## 6 Analysis of Prediction Error

In this section, we present an analysis of the in-sample prediction error for the gated ReLU network given in (5) along the lines of existing consistency results [36; 37] for the group lasso problem (8).
We assume that the data is generated by a noisy ReLU neural network model \(\mathbf{y}=f(\mathbf{X},\mathbf{\theta}^{*})+\epsilon\) where \(f\) is the ReLU network defined in (1) with true parameters \(\mathbf{\theta}^{*}\) and the noise is distributed as \(\epsilon\sim\mathcal{N}\big{(}0,\sigma^{2}\mathbf{I}_{n}\big{)}\). By the seminal universal approximation theorem of [38], this model is able to capture a broad class of ground-truth functions by using a ReLU network with enough neurons. We can transform \(\mathbf{\theta}^{*}\) to the weights \(\mathbf{w}^{*}\) in the convex ReLU model and write \(f(\mathbf{X},\mathbf{\theta}^{*})=\sum_{i=1}^{p}\mathbf{D}_{i}\mathbf{X}\mathbf{w} _{i}^{*}\). We denote by \(\hat{\mathbf{w}}\) the solution of the group lasso problem (8) for the gated ReLU network with an additional \(\frac{1}{n}\) factor on the loss to simplify derivations,

\[\hat{\mathbf{w}}=\operatorname*{argmin}_{\mathbf{w}_{i}\in\mathbb{R}^{d}} \quad\frac{1}{n}\left\|\sum_{i=1}^{p}\mathbf{D}_{i}\mathbf{X}\mathbf{w}_{i}- \mathbf{y}\right\|_{2}^{2}+\lambda\sum_{i=1}^{p}\left\|\mathbf{w}_{i}\right\|_ {2}. \tag{17}\]

We now present our main theorem which bounds the prediction error of the gated ReLU network obtained from the solution \(\hat{\mathbf{w}}\) below.

**Theorem 6.1** (Prediction Error of Gated ReLU).: _For some \(t>0\), let the regularization parameter in (17) be \(\lambda=t\sigma\|\mathbf{X}\|_{F}/n\). Then, with probability at least \(1-2e^{-t^{2}/8}\), we have_

\[\frac{1}{n}\left\|f_{\mathcal{G}}\big{(}\mathbf{X},\hat{\mathbf{\theta}}\big{)}-f \left(\mathbf{X},\mathbf{\theta}^{*}\right)\right\|_{2}^{2}\leq 2\lambda\sum_{i=1}^{p} \left\|\mathbf{w}_{i}^{*}\right\|_{2}\]

_where \(f_{\mathcal{G}}\big{(}\mathbf{X},\hat{\mathbf{\theta}}\big{)}=\sum_{i=1}^{p} \mathbf{D}_{i}\mathbf{X}\hat{\mathbf{w}}_{i}\) are the predictions of the gated ReLU network obtained from \(\hat{\mathbf{w}}\)._

The proof closely follows the analysis of the regular lasso problem presented in [37], but we extend it to the specific case of the group lasso problem corresponding to the gated ReLU network and leverage the masking structure of the lifted data matrix to obtain simplified bounds.

## 7 Experiments

Here, we empirically corroborate our theoretical results via experiments on several datasets.3

Figure 1: Plot of objective value of problem (8) which is solved using IRLS (algorithm 1) for a toy 1D dataset with \(n=5\). The iterates are compared to the optimal value obtained by solving (8) using CVXPY (blue). Notice that the solution to (16) with regularization weights given by the NTK weights \(\mathbf{\tilde{\eta}}\) from Theorem 5.1 (green) is sub-optimal for problem (8), and running IRLS by initializing with these weights (red) converges to the optimal objective value. We also include plots of IRLS initialized with random weights (black). The right plot shows the corresponding learned functions. The blue curve shows the output of the solution to the group lasso problem (8) after performing a cone decomposition to obtain a ReLU network. The dashed curve shows the output of running gradient descent (GD) on a ReLU network with 100 neurons. The green curve is the result of kernel ridge regression (KRR) with the NTK. **We observe that our reweighted kernel method (Group Lasso) produces an output that matches the output of the NN trained via GD. In contrast, NTK produces an erroneous smooth function due to the infinite width approximation.**
**1D datasets.** For the 1D experiments in Figure 1, we add a second data dimension with value equal to \(1\) for all data points to simulate a bias term in the first layer of the gated ReLU network. Also, in this case we can enumerate all \(2n\) data masks \(\mathbf{D}_{i}\) directly. We use the cone decomposition procedure described in [8] to obtain a ReLU network from the gated ReLU network obtained by solving (8). We also train a \(100\) neuron ReLU network using gradient descent (GD) and compare the learned output functions in the right plot in Figure 1. Exact details can be found in the supplementary material.

**Student-teacher setting.** We generate the training data by sampling \(\mathbf{X}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{\mathbf{n}})\) and computing the targets \(\mathbf{y}\) using a fixed, randomly initialized gated ReLU teacher network. In Figure 2, \(n=10\), \(d=5\), and we use a teacher network with width \(m=10\), with gates and parameters randomly drawn from independent standard multivariate Gaussians. To solve the convex formulation (8), we estimate \(\mathcal{D}_{\mathbf{X}}\) by randomly sampling unique hyperplane arrangements.

**Computing the NTK weights \(\boldsymbol{\tilde{\eta}}\).** The weights induced by the NTK are given as in Theorem 5.1 by \(\tilde{\eta}_{i}=\mathbb{P}[\operatorname{diag}\left(\mathbb{1}\left\{\mathbf{ X}\mathbf{h}\geq\mathbf{0}\right\}\right)=\mathbf{D}_{i}]\) where \(\mathbf{h}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{\mathbf{d}})\) is a standard multivariate Gaussian vector. These probabilities can be interpreted either as the orthant probabilities of the multivariate Gaussians given by \(\left(2\mathbf{D}_{i}-\mathbf{I}_{\mathbf{n}}\right)\mathbf{X}\mathbf{h}\) or as the solid angle of the cones given by \(\left\{\mathbf{u}\in\mathbb{R}^{d}:\left(2\mathbf{D}_{i}-\mathbf{I}_{\mathbf{n }}\right)\mathbf{X}\mathbf{u}\geq\mathbf{0}\right\}\). Closed form expressions exist for \(d=2,3\), and [39, 40] present approximating schemes for higher dimensions. We calculate these weights exactly for the 1D example presented in Figure 1, and estimate them using Monte Carlo sampling for the student-teacher example in Figure 2.

\begin{table}
\begin{tabular}{l c c|c c} \hline \hline
**Dataset** & \(n\) & \(d\) & **NTK** & **Ours (Alg 1)** \\ \hline acute-inflammation & \(120\) & \(6\) & \(1.000\) & \(1.000\) \\ acute-nephritis & \(120\) & \(6\) & \(1.000\) & \(1.000\) \\ balloons & \(16\) & \(4\) & \(0.75\) & \(0.75\) \\ blood & \(748\) & \(4\) & \(0.524\) & **0.583** \\ breast-cancer & \(286\) & \(9\) & \(0.417\) & **0.625** \\ breast-cancer-wisc-prog & \(699\) & \(9\) & \(9.06\) & **0.966** \\ breast-cancer-wisc-diag & \(569\) & \(30\) & **0.965** & \(0.915\) \\ breast-cancer-wisc-prog & \(198\) & \(33\) & **0.7** & \(0.66\) \\ congressional-voting & \(435\) & \(16\) & \(0.266\) & **0.266** \\ conn-bench-sonar-mines-rocks & \(208\) & \(60\) & \(6.635\) & **0.712** \\ credit-approval & \(690\) & \(15\) & \(0.838\) & **0.844** \\ cylinder-bands & \(512\) & \(35\) & \(0.773\) & **0.82** \\ echocardiogram & \(131\) & \(10\) & \(0.758\) & **0.788** \\ fertility & \(100\) & \(9\) & \(0.76\) & \(0.76\) \\ haberman-survival & \(306\) & \(3\) & \(0.481\) & **0.532** \\ heart-hungarian & \(294\) & \(12\) & \(0.743\) & **0.878** \\ hepatitis & \(155\) & \(19\) & **0.923** & \(0.897\) \\ ilpd-indian-liver & \(583\) & \(9\) & \(0.432\) & **0.555** \\ ionosphere & \(351\) & \(33\) & \(0.955\) & **0.966** \\ mammographic & \(961\) & \(5\) & \(0.783\) & **0.792** \\ molec-biol-promoter & \(106\) & \(57\) & \(0.63\) & **0.815** \\ musk-1 & \(476\) & \(166\) & \(0.782\) & **0.866** \\ oocytes\_trisoperus\_nucleus\_2T & \(912\) & \(25\) & **0.781** & 0.759 \\ parkinsons & \(195\) & \(22\) & 0.939 & **0.959** \\ pima & \(768\) & \(8\) & \(0.552\) & **0.599** \\ pittsburg-bridges-T-OR-D & \(102\) & \(7\) & \(0.731\) & **0.846** \\ planning & \(182\) & \(12\) & \(0.435\) & **0.543** \\ stalog-australian-credit & \(690\) & \(14\) & **1.0** & \(0.74\) \\ stalog-german-credit & \(1000\) & \(24\) & \(0.512\) & **0.576** \\ stalog-heart & \(270\) & \(13\) & **0.779** & \(0.765\) \\ tic-tac-toe & \(958\) & \(9\) & \(1.0\) & \(1.0\) \\ trains & \(10\) & \(29\) & \(0.667\) & \(0.667\) \\ vertebral-column-2classes & \(310\) & \(6\) & **0.821** & \(0.731\) \\ \hline \multicolumn{4}{c}{Higher (or same) accuracy} & \multicolumn{2}{c}{14/33} & \multicolumn{2}{c}{26/33} \\ \end{tabular}
\end{table}
Table 1: Test accuracies for UCI experiments with \(75\%-25\%\) training-test split. Our approach achieves either higher or the same accuracy for \(26\) out of \(33\) datasets.

Figure 2: Plot of the **group lasso objective** in the student-teacher setting with \(d=5\). Training data is generated using a teacher network with width \(m=10\). The NTK weights \(\boldsymbol{\tilde{\eta}}\) are estimated using Monte Carlo sampling, and are again sub-optimal. IRLS initialized with these weights successfully fixes the NTK and converges to the optimal weights.
**Fixing the NTK weights by IRLS.** We solve (16) with \(\mathbf{\Phi}_{i}=\mathbf{D}_{i}\mathbf{X}\) and regularization weights given by the NTK weights \(\mathbf{\tilde{\eta}}\) to obtain the solution \(\tilde{\mathbf{w}}\). We use efficient least squares solvers from [41; 42]. By Theorem 5.1, this corresponds to choosing the weights \(\mathbf{\eta}\) in the MKL problem (13) such that the resulting kernel matrix is equal to the NTK matrix \(\mathbf{H}\). We find the exact solution of the group lasso problem (8) using CVXPY. Comparing the optimal value of the group lasso problem (8) with the objective value of \(\tilde{\mathbf{w}}\) (given by the green line) in Figures 1 and 2, we observe that the NTK weighted solution is sub-optimal. This means that \(\mathbf{H}\) is not the optimal kernel that would be learnt by MKL (which is expected since \(\mathbf{H}\) has no dependence on the targets \(\mathbf{y}\)). By applying IRLS initialized with the NTK weights, we _fix_ the NTK and find the weights of the optimal MKL kernel. In Figures 1 and 2, we observe that IRLS converges to the solution of the group lasso problem (8) and fixes the NTK.

**UCI datasets.** We compare the regularised NTK with our IRLS algorithm (Algorithm 1) on the UCI ML Repository datasets. We follow the procedure described in [43] for \(n\leq 1000\) to extract and standardize the datasets. We observe that our method achieves higher (or the same) test accuracy for \(26\)**out of \(33\)** datasets (see Table 1 for details) while the NTK achieves higher (or the same) test accuracy for \(14\) datasets which empirically supports our main claim that the IRLS procedure fixes the NTK. Details of the experiments can be found in Section B of the supplementary material.

## 8 Discussion and Limitations

In this work, we explored the connection between finite-width theories of neural networks given by convex reformulations and infinite-width theories of neural networks given by the NTK. To bridge these theories, we first interpreted the group lasso convex formulation of the gated ReLU network as a multiple kernel learning model using the masking kernels generated by its gates. Then, we linked this MKL model with the NTK of the gated ReLU network evaluated on the training data. Specifically, we showed that the NTK is equivalent to the weighted masking kernel with weights that depend only on the input data \(\mathbf{X}\) and not on the targets \(\mathbf{y}\). We contrast this with the MKL interpretation of the gated ReLU network which learns the optimal data-dependent kernel using both \(\mathbf{X}\) and \(\mathbf{y}\). Therefore, the NTK cannot perform better than the optimal MKL kernel. To fix the NTK, we improve the weights induced by it using the iteratively reweighted least squares (IRLS) scheme to obtain the optimal solution of the group lasso formulation of the gated ReLU network. We corroborated our theoretical results by empirically running IRLS on toy datasets.

While our theory is able to link the optimization properties of the NTK with those of finite width networks via the MKL characterization of group lasso, we do not derive explicit generalization results on the test set. Applying existing generalization theory for kernel methods [44; 45] to the MKL interpretation of the convex reformulation could be a promising direction for future work.

Finally, although we studied fully connected networks in this paper, our approach can be directly extended to various neural network architectures, e.g., threshold/binary networks [46], convolution networks [47], generative adversarial networks [48], NNs with batch normalization [49], autoregressive models [50], and Transformers [51; 52].

## Acknowledgements

This work was supported in part by the National Science Foundation (NSF) CAREER Award under Grant CCF-2236829, Grant DMS-2134248 and Grant ECCS-2037304; in part by the U.S. Army Research Office Early Career Award under Grant W911NF-21-1-0242; in part by the Stanford Precourt Institute; and in part by the ACCESS--AI Chip Center for Emerging Smart Systems through InnoHK, Hong Kong, SAR.