Title: Fixing the NTK: From Neural Network Linearizations to Exact Convex Programs
Venue: NeurIPS 2023
Accepted: yes
Keywords: neural tangent kernel

NTK

ReLU activations

neural networks

gated ReLU

convex optimization

kernel

multiple kernel learning

MKL

group lasso

iterative reweighting

group norm
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.

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-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 × k as I k . To denote the set { 1 , 2 , . . . , n } , we use [ n ] . We also use ∥·∥ p and ∥·∥ F to represent the standard ℓ p and Frobenius norms. Additionally, we denote 0-1 valued indicator function and ReLU activation as { x ≥ 0 } and ( x ) + := 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

where w (1) j ∈ R d and w (2) j are the j th hidden and output layer weights, respectively and θ := { ( w (1) j , w (2) j ) } j m =1 represents all trainable parameters. Given a training data matrix X ∈ R n × d and a target vector y ∈ R n , we minimize the following weight decay regularized training objective

where λ > 0 is the regularization coefficient. We also use f ( X , θ ) = ∑ m j =1 ( Xw (1) j ) + w (2) j 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 λ (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.

2.1 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-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

where D i ∈ D X are n × n binary masking diagonal matrices given by p := |D X | and

D X := { diag ( ✶ { Xu ≥ 0 } ) : u ∈ R d } . (4) The mask set D X can be interpreted as the set of all possible ways to separate the training data 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: ( Xw i ) + = diag ( ✶ { Xw i ≥ 0 } ) Xw i = D i Xw i provided that (2 D i -I n ) Xw i ≥ 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 O ( n r ) , where r = rank( X ) ≤ 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

where G := { g j } j m =1 is a set of gate vectors that are also optimized throughout training. Then, the corresponding non-convex learning problem is as follows

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 G is complete with respect to a dataset X if the corresponding set of hyperplane arrangements covers all possible arrangement patterns for X defined in (4) , i.e.,

✶ Additionally, G is minimally complete if |G| = |D X | = p .

Now, with this relaxation, [8] showed that the optimal value for (6) can always be achieved by choosing 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 W (1) and w (2)

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

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.

2.2 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 → ∞ , 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 ( x , θ ) by linearizing it with respect to the parameters θ around its initialization θ 0 :

The authors of [1] show that if f is a neural network (with appropriately scaled output), and parameters θ initialized as i.i.d standard Gaussians, the linearization ˆ f better approximates f in the infinite-width limit. We can interpret the linearized model ˆ f as a kernel method with a feature map given by ϕ ( x ) = ∇ θ f ( x , θ 0 ) . 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 θ 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 ( x , x ′ ) := E [ ∇ ˆ θ f ( x , ˆ θ 0 ) T ∇ ˆ θ f ( x ′ , ˆ θ 0 )] , where ˆ θ 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

In the next lemma, we provide the infinite width NTK of the scaled gated ReLU network in (9). Lemma 2.2. 1 The infinite width NTK of the gated ReLU network (9) with i.i.d gates sampled as g j ∼ N ( 0 , I d ) and randomly initialized parameters as w (1) j ∼ N ( 0 , I d ) and w (2) j ∼ N (0 , 1) is

Additionally, we introduce a reparameterization of the standard ReLU network (1) with θ := {( w (+) j , w ( -) j )} m j =1 , with w (+) j , w ( -) j ∈ R d which can still represent all the functions that (1) can

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 ϕ i : R d → R d i . Given n training samples X ∈ R n × d with targets y ∈ R n , we define the feature matrices on this data by stacking the feature vectors as Φ i := [ ϕ i ( x 1 ) T ; . . . ; ϕ i ( x n ) T ] ∈ R n × d i . Then, the corresponding n × n kernel matrices are given by K i := Φ i Φ T i . A convex combination of these kernels can be written as K ( η ) := ∑ p i =1 η i K i where η ∈ ∆ p := { η : 1 T η = 1 , η ≥ 0 } is a set of weights in the unit simplex. By noticing that the feature map corresponding to K ( η )

is obtained by taking a weighted concatenation of ϕ i with weights √ η i , we can write the MKL optimization problem in terms of the feature matrices as

where ˆ λ > 0 is a regularization coefficient. For a set of fixed weights η , the optimal objective value of the kernel regression problem over v is proportional to y T ( ∑ p i =1 η i K i + ˆ λ I n ) -1 y up to constant factors [16]. Thus, the MKL problem can be equivalently written as the following problem

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 K i . In the next section, we link this kernel learning formulation with the convex group lasso problem in (8).

3.1 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

We now present a short derivation of this equivalence. Using the variational formulation of the squared group ℓ 1 -norm [31]

we can rewrite the group lasso problem (14) as a joint minimization problem over both the parameters w and regularization weights η as follows

Finally, with a change of variables given by v i = w i / √ η i , we recover the MKL problem (12). We note that the MKL problem (12) is also equivalent to the following standard group lasso problem

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 ˆ λ = λ ∑ p i =1 ∥ w ∗ i ∥ 2 , where w ∗ is the solution to (14).

3.2 Solving Group Lasso by Iterative Reweighting

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

Based on this, we rewrite the group lasso problem (15) as the following minimization problem

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

The outer problem in η is also directly solved by setting η i = ∥ w i ∥ 2 [34]. To avoid convergence issues and instability around η i = 0 , we approximate the reweighting by adding a small positive constant ϵ . We use this procedure to solve the group lasso formulation of the gated ReLU network (8) by setting Φ i = D i X . A detailed description is provided Algorithm 1. For further details regarding convergence, we refer the reader to [32-34].

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

1: Set iteration count k ← 0

3: Set Φ i := D i X , ∀ D i ∈ D X

2: Initialize weights η (0) i

4: while not converged and k ≤ max iteration count do

5: Solve the weighted ℓ 2 regularized least squares problem:

6: Update the weights: η ( k +1) i = √∥ ∥ ∥ w ( k ) i ∥ ∥ ∥ 2 + ϵ

7: Increment iteration count: k ← k +1

8: end while

9: Optional: Convert the gated ReLU network to a ReLU network (see Section E for details)

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.

̸

Lemma 4.2. For a minimally complete gate set G defined in Definition 2.1, we can uniquely associate a gate vector g i to each data mask D i ∈ D X such that D i = diag ( ✶ { Xg i ≥ 0 } ) , ∀ i ∈ [ p ] .

Definition 4.1. The masking feature maps ϕ j : R d → R d generated by a fixed set of gates G are defined as ϕ j ( x ) = ✶ { x T g j ≥ 0 } 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 g j , i.e., x T g j ≥ 0 , and returning zero if the input does not lie in this halfspace. Since diag ( ✶ { Xg j ≥ 0 } ) ∈ D X , ∀ g j ∈ G holds for an arbitrary gate set G , we can conveniently express the corresponding feature matrices of these masking feature maps on the data X in terms of fixed diagonal data masks as Φ j = D j X , where D j = diag ( ✶ { Xg j ≥ 0 } ) . Similarly, the corresponding masking kernel matrices take the form K j = D j XX T D j . Note that for an arbitrary set of gates, the generated masking feature matrices on X may not cover the entire set of possible masks D X . Additionally, multiple gate vectors can result in identical masks if diag ( ✶ { Xg i ≥ 0 } ) = diag ( ✶ { Xg j ≥ 0 } ) for i = 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.

Consequently, for minimally complete gate sets G , the generated masking feature matrices ∀ i ∈ [ p ] can be expressed as Φ i = D i X and the masking kernel matrices take the form K i = D i XX T D i . In the context of the gated ReLU problem (7), since G is complete, we can replace it with a minimally complete subset of 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 G is equivalent to performing multiple kernel learning (12) with the masking feature maps generated by G

This theorem implies that the gated ReLU network finds the optimal combination of linear models restricted to the different masked datasets D i 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 K G ( ˜ η ) ∈ R n × n be the weighted masking kernel obtained by taking a convex combination of the masking feature maps generated by a minimally complete gate set G with weights given by ˜ η i = P [diag ( ✶ { Xh ≥ 0 } ) = D i ] where h ∼ N ( 0 , I d ) and let H ∈ R n × n be the infinite width NTK of the gated ReLU network (10) evaluated on the training data, i.e., the ij th entry of H is defined as H ij := H ( x i , x j ) . Then, K G ( ˜ η ) = H .

A rough sketch of the proof of this theorem is to express the matrix 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 D i , we can express the NTK as a convex combination of the masking kernels 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 ˜ η on the training data is identical to that of kernel ridge regression with the NTK.

Corollary 5.2. Let ˜ w be the solution to (16) with feature matrices Φ i = D i X and regularization weights ˜ η i = P [diag ( ✶ { Xh ≥ 0 } ) = D i ] where h ∼ N ( 0 , I d ) and let ˜ y = H ( H + λ I n ) -1 y be the outputs of kernel ridge regression with the NTK on the training data. Then, ∑ p i =1 D i X ˜ w i = ˜ y .

Since G is minimally complete, the weights in Theorem 5.1 satisfy ∑ p i =1 ˜ η i = 1 . In other words, ˜ η ∈ ∆ 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 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 ˜ η do depend on the training data X , but do not depend on the target labels y . Since MKL learns the optimal kernel using both X and 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).

Grouplasso objective,入=0.01

Learnedfunctions

4 × 10-2

1.0

0.5

IRLSwithrandominits

3x10-2

IRLSwithNTKinit

0.0

GroupLasso

Optimal

NTK weights

GradientDescent

-0.5

Training data

2×10-2

-1.0

0

2

4

6

8

-2

-1

0

1

2

Iterationk

We assume that the data is generated by a noisy ReLU neural network model y = f ( X , θ ∗ ) + ϵ where f is the ReLU network defined in (1) with true parameters θ ∗ and the noise is distributed as ϵ ∼ N ( 0 , σ 2 I n ) . 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 θ ∗ to the weights w ∗ in the convex ReLU model and write f ( X , θ ∗ ) = ∑ p i =1 D i Xw ∗ i . We denote by ˆ w the solution of the group lasso problem (8) for the gated ReLU network with an additional 1 n factor on the loss to simplify derivations,

We now present our main theorem which bounds the prediction error of the gated ReLU network obtained from the solution ˆ w below.

Theorem 6.1 (Prediction Error of Gated ReLU) . For some t > 0 , let the regularization parameter in (17) be λ = tσ ∥ X ∥ F /n. Then, with probability at least 1 -2 e -t 2 / 8 , we have

where f G ( X , ˆ θ ) = ∑ p i =1 D i X ˆ w i are the predictions of the gated ReLU network obtained from ˆ 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. 2

GroupLassoobjectiveforStudent-Teachersetting

入=0.001

入=0.01

7×10-3

8×10-2

7 ×10-2

IRLSwithrandominits

6×10-3

IRLSwithNTKinit

6×10-2

Optimal

NTK weights

5×10-3

5×10-2

0

20

40

60

80

100

0

20

40

60

80

100

Iterationk

Iterationk

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 2 n data masks 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 X ∼ N ( 0 , I n ) and computing the targets 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 D X by randomly sampling unique hyperplane arrangements.

Computing the NTK weights ˜ η .

The weights induced by the NTK are given as in Theorem 5.1 by ˜ η i = P [diag ( ✶ { Xh ≥ 0 } ) = D i ] where h ∼ N ( 0 , I d ) is a standard multivariate Gaussian vector. These probabilities can be interpreted either as the orthant probabilities of the multivariate Gaussians given by (2 D i -I n ) Xh or as the solid angle of the cones given by { u ∈ R d : (2 D i -I n ) Xu ≥ 0 } . 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.

Higher (or same) accuracy

14/33

26/33

Fixing the NTK weights by IRLS. We solve (16) with Φ i = D i X and regularization weights given by the NTK weights ˜ η to obtain the solution ˜ w . We use efficient least squares solvers from [41, 42]. By Theorem 5.1, this corresponds to choosing the weights η in the MKL problem (13) such that the resulting kernel matrix is equal to the NTK matrix 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 ˜ w (given by the green line) in Figures 1 and 2, we observe that the NTK weighted solution is sub-optimal. This means that H is not the optimal kernel that would be learnt by MKL (which is expected since H has no dependence on the targets 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 regularied NTK with our IRLS algorithm (Algorithm 1) on the UCI ML Repository datasets. We follow the procedure described in [43] for n ≤ 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 X and not on the targets y . We contrast this with the MKL interpretation of the gated ReLU network which learns the optimal data-dependent kernel using both X and 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].

References

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems , 31, 2018.

Lenaic Chizat and Francis Bach. A note on lazy training in supervised differentiable programming. arXiv preprint arXiv:1812.07956 , 2018.

Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. Advances in Neural Information Processing Systems , 33:5850-5861, 2020.

Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference on Learning Representations , 2020.

Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems , pages 8139-8148, 2019.

Tolga Ergen and Mert Pilanci. Convex geometry and duality of over-parameterized neural networks. Journal of Machine Learning Research , 22(212):1-63, 2021.

Mert Pilanci and Tolga Ergen. Neural networks are convex regularizers: Exact polynomial-time convex optimization formulations for two-layer networks. In International Conference on Machine Learning , pages 7695-7705. PMLR, 2020.

Aaron Mishkin, Arda Sahiner, and Mert Pilanci. Fast convex optimization for two-layer relu networks: Equivalent model classes and cone decompositions. arXiv preprint arXiv:2202.01331 , 2022.

Arda Sahiner, Tolga Ergen, John M. Pauly, and Mert Pilanci. Vector-output relu neural network problems are copositive programs: Convex analysis of two layer networks and polynomial-time algorithms. In International Conference on Learning Representations , 2021.

Tolga Ergen and Mert Pilanci. Convex geometry of two-layer relu networks: Implicit autoencoding and interpretable models. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics , volume 108 of Proceedings of Machine Learning Research , pages 4024-4033, Online, 26-28 Aug 2020. PMLR.

Tolga Ergen and Mert Pilanci. Convex duality and cutting plane methods for over-parameterized neural networks. In OPT-ML workshop , 2019.

Tolga Ergen and Mert Pilanci. Convex programs for global optimization of convolutional neural networks in polynomial-time. In OPT-ML workshop , 2020.

Tolga Ergen and Mert Pilanci. Revealing the structure of deep neural networks via convex duality. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 3004-3014. PMLR, 18-24 Jul 2021.

Yifei Wang, Tolga Ergen, and Mert Pilanci. Parallel deep neural networks have zero duality gap. In The Eleventh International Conference on Learning Representations , 2023.

T. Ergen and M. Pilanci. Convex optimization for shallow neural networks. In 2019 57th Annual Allerton Conference on Communication, Control, and Computing (Allerton) , pages 79-83, 2019.

Alain Rakotomamonjy, Francis Bach, Stéphane Canu, and Yves Grandvalet. Simplemkl. Journal of Machine Learning Research , 9:2491-2521, 2008.

Shai Shalev-Shwartz, Ohad Shamir, and Shaked Shammah. Failures of gradient-based deep learning. arXiv preprint arXiv:1703.07950 , 2017.

Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning , volume 1. MIT press Cambridge, 2016.

Itay Safran and Ohad Shamir. Spurious local minima are common in two-layer relu neural networks. In International Conference on Machine Learning , pages 4433-4441. PMLR, 2018.

Rong Ge, Jason D. Lee, and Tengyu Ma. Learning one-hidden-layer neural networks with landscape design, 2017.

Yifei Wang, Jonathan Lacotte, and Mert Pilanci. The hidden convex optimization landscape of regularized two-layer reLU networks: an exact characterization of optimal solutions. In International Conference on Learning Representations , 2022.

Ming Yuan and Yi Lin. Model selection and estimation in regression with grouped variables. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 68(1):49-67, 2006.

Lei Tan, Shutong Wu, and Xiaolin Huang. Weighted neural tangent kernel: A generalized and improved network-induced kernel. arXiv preprint arXiv:2103.11558 , 2021.

Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. arXiv preprint arXiv:1909.12292 , 2019.

Lenaic Chizat. Sparse optimization on measures with over-parameterized gradient descent. Mathematical Programming , 194(1-2):487-532, 2022.

Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665E7671, 2018.

Rajat Vadiraj Dwaraknath. Understanding the neural tangent kernel. https:// rajatvd . github . io/NTK/ , 2019. Accessed: 2022-05-17.

Simon Du and Wei Hu. Ultra-wide deep nets and the neural tangent kernel (ntk). https://blog . ml . cmu . edu/2019/10/03/ultra-wide-deep-nets-and-theneural-tangent-kernel-ntk/ , 2019. Accessed: 2022-05-17.

Francis R Bach, Gert RG Lanckriet, and Michael I Jordan. Multiple kernel learning, conic duality, and the smo algorithm. In Proceedings of the twenty-first international conference on Machine learning , page 6, 2004.

Francis R Bach. The 'eta-trick' reloaded: multiple kernel learning, 2019. Accessed: 2022-0517.

Francis R Bach. Consistency of the group lasso and multiple kernel learning. Journal of Machine Learning Research , 9(6), 2008.

Ingrid Daubechies, Ronald DeVore, Massimo Fornasier, and C Sinan G üntürk. Iteratively reweighted least squares minimization for sparse recovery. Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences , 63(1):1-38, 2010.

Francis Bach, Rodolphe Jenatton, Julien Mairal, Guillaume Obozinski, et al. Optimization with sparsity-inducing penalties. Foundations and Trends® in Machine Learning , 4(1):1-106, 2012.

Julien Mairal, Francis Bach, and Jean Ponce. Sparse modeling for image and vision processing. arXiv preprint arXiv:1411.3230 , 2014.

Francis R Bach. The 'eta-trick' or the effectiveness of reweighted least-squares. https://francisbach . com/the-%CE%B7-trick-or-the-effectiveness-ofreweighted-least-squares/ , 2019. Accessed: 2022-05-17.

Han Liu and Jian Zhang. Estimation consistency of the group lasso and its applications. In Artificial Intelligence and Statistics , pages 376-383. PMLR, 2009.

Peter Bühlmann, Sara van de Geer, Peter Bühlmann, and Sara van de Geer. Theory for the lasso. Statistics for High-Dimensional Data: Methods, Theory and Applications , pages 99-182, 2011.

Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks , 2(5):359-366, 1989.

Dario Azzimonti and David Ginsbourger. Estimating orthant probabilities of high-dimensional gaussian vectors with an application to set estimation. Journal of Computational and Graphical Statistics , 27(2):255-267, 2018.

Peter Craig. A new reconstruction of multivariate normal orthant probabilities. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 70(1):227-243, 2008.

Christopher C Paige and Michael A Saunders. Lsqr: An algorithm for sparse linear equations and sparse least squares. ACM Transactions on Mathematical Software (TOMS) , 8(1):43-71, 1982.

David Chin-Lung Fong and Michael Saunders. Lsmr: An iterative algorithm for sparse leastsquares problems. SIAM Journal on Scientific Computing , 33(5):2950-2971, 2011.

Manuel Fernández-Delgado, Eva Cernadas, Senén Barro, and Dinani Amorim. Do we need hundreds of classifiers to solve real world classification problems? Journal of Machine Learning Research , 15(90):3133-3181, 2014.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Generalization error of random feature and kernel methods: Hypercontractivity and kernel matrix concentration. Applied and Computational Harmonic Analysis , 59:3-84, 2022.

Corinna Cortes, Mehryar Mohri, and Afshin Rostamizadeh. Generalization bounds for learning kernels. International Conference on Machine Learning , 2010.

Tolga Ergen, Halil I. Gulluk, Jonathan Lacotte, and Mert Pilanci. Globally optimal training of neural networks with threshold activation functions. In The Eleventh International Conference on Learning Representations , 2023.

Tolga Ergen and Mert Pilanci. Implicit convex regularizers of cnn architectures: Convex optimization of two- and three-layer networks in polynomial time. In International Conference on Learning Representations , 2021.

Arda Sahiner, Tolga Ergen, Batu Ozturkler, Burak Bartan, John M. Pauly, Morteza Mardani, and Mert Pilanci. Hidden convexity of wasserstein GANs: Interpretable generative models with closed-form solutions. In International Conference on Learning Representations , 2022.

Tolga Ergen, Arda Sahiner, Batu Ozturkler, John M. Pauly, Morteza Mardani, and Mert Pilanci. Demystifying batch normalization in reLU networks: Equivalent convex optimization models and implicit regularization. In International Conference on Learning Representations , 2022.

Vikul Gupta, Burak Bartan, Tolga Ergen, and Mert Pilanci. Convex neural autoregressive models: Towards tractable, expressive, and theoretically-backed models for sequential forecasting and generation. In ICASSP 2021 - 2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) , pages 3890-3894, 2021.

Tolga Ergen, Behnam Neyshabur, and Harsh Mehta. Convexifying transformers: Improving optimization and understanding of transformer networks. arXiv, 2022.

Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 19050-19088. PMLR, 17-23 Jul 2022.

Kaare Brandt Petersen, Michael Syskind Pedersen, et al. The matrix cookbook. Technical University of Denmark , 7(15):510, 2008.

Erling D Andersen and Knud D Andersen. The mosek interior point optimizer for linear programming: an implementation of the homogeneous algorithm. In High performance optimization , pages 197-232. Springer, 2000.

Dheeru Dua and Casey Graff. UCI machine learning repository, 2017.

Stephen Boyd and Lieven Vandenberghe. Convex optimization . Cambridge university press, 2004.

Tolga Ergen and Mert Pilanci. Global optimality beyond two layers: Training deep relu networks via convex programs. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning , volume 139 of Proceedings of Machine Learning Research , pages 2993-3003. PMLR, 18-24 Jul 2021.

Tolga Ergen and Mert Pilanci. Path regularization: A convexity and sparsity inducing regularization for parallel relu networks. arXiv preprint arXiv:2110.09548 , 2021.

Supplementary Material

Table of Contents

A Proofs and Theoretical Results

A.1 Proof of Lemma 2.2

The gradient feature map of (9) is

✶ The finite width random NTK given by H m, θ , G ( x , x ′ ):= ∇ θ ˜ f G ( x , θ ) T ∇ θ ˜ f G ( x ′ , θ ) is

By the Law of Large numbers since g j , w (1) j , w (2) j are iid, H m, θ , G ( x , x ′ ) converges in probability to the expectation of the quantity in the sum as m →∞ . We denote this infinite limit by H ( x , x ′ ) . Therefore,

where g ∼ N ( 0 , I d ) , w (1) ∼ N ( 0 , I d ) , w (2) ∼ N (0 , 1) . Using E [ ( w (2) ) 2 ] = 1 , E [ x T w (1) x ′ T w (1) ] = x T x ′ and the linearity of expectation, we get

To evaluate the expectation of the product of the indicator random variables, we condition on one of the indicators being equal to 1.

✶ ✶ The conditional expectation is the probability of the event that x T g ≥ 0 given that x ′ T g ≥ 0 . This can be geometrically seen as the ratio of π minus the angle between x and x ′ to π . So, we get

Combining this with (18) yields the desired result.

A.2 Proof of Lemma 2.3

The gradient feature map of (11) is

The finite width random NTK given by H r m, θ ( x , x ′ ):= ∇ θ f r ( x , θ ) T ∇ θ f r ( x ′ , θ ) is

By the Law of Large numbers since w (+) j , w ( -) j are iid, H r m, θ ( x , x ′ ) converges in probability to the expectation of the quantity in the sum as m →∞ . We denote this infinite limit by H r ( x , x ′ ) . Therefore,

where w (+) ∼ N ( 0 , I d ) and w ( -) ∼ N ( 0 , I d ) . Since w (+) and w ( -) are i.i.d., we can combine the expectations to get

Comparing to (18), we get the desired result.

A.3 Proof of Lemma 4.2

Since G is complete, ∀ i ∈ [ p ] , ∃ g i ∈ G such that diag ( ✶ { Xg i ≥ 0 } ) = D i . Additionally, these gates g i are unique since diag ( ✶ { Xg i ≥ 0 } ) = diag ( ✶ { Xg j ≥ 0 } ) = ⇒ g i = g j . Therefore, there is a subset of gates G of size p that are uniquely associated to each mask D i . Since G is also minimally complete, this subset must be G and this proves the lemma.

̸

̸

A.4 Proof of Theorem 4.3

By Lemma 3.1, the MKL problem given by

is equivalent to the group lasso problem (14). Additionally, (14) is also equivalent to the standard group lasso (15) since they have the same regularization paths. Since G is minimally complete, (7) (which is equivalent to (6) in the minimally complete gate set setting) is equivalent to the group lasso problem (15) and this proves the theorem.

A.5 Proof of Theorem 5.1

The weighted masking kernel is given by

Now, the NTK matrix H can be written as (using (18)),

H = E g ∼N ( 0 , I d ) [ diag ( ✶ { Xg ≥ 0 } ) XX T diag ( ✶ { Xg ≥ 0 } ) ] . We can condition on the events that diag ( ✶ { Xg ≥ 0 } ) = D i ∀ i ∈ [ p ] . Since D X covers all possible data maskings D i , we know that ∑ p i =1 P [diag ( { Xg ≥ 0 } ) = D i ] = 1 . Therefore,

Plugging in ˜ η i = P [diag ( ✶ { Xg ≥ 0 } ) = D i ] and comparing to (19), we get that K G ( ˜ η ) = H and this proves Theorem 5.1.

A.6 Proof of Corollary 5.2

First, we define a single combined masked data matrix ˜ X ∈ R n × pd as follows

We also define a diagonal regularization weighting matrix R ∈ R pd × pd as follows

where 1 d is a d -dimensional all-1s row vector. Now, we can rewrite (16) with Φ i = D i X and regularization weights ˜ η as

Setting the gradient of the convex objective in (20) to 0, we get the following equation for the solution ˜ w ,

Solving for ˜ w , we get

Now, we notice that

Additionally, observe that ˜ XR ˜ X T = ∑ p i =1 ˜ η i D i XXD i = H where the last equality is by Theorem 5.1. Plugging this into (21), we get

and this proves Corollory 5.2.

A.7 Proof of Theorem 6.1

We begin with the optimality of ˆ w ,

By plugging in the model y = ∑ p i =1 D i Xw ∗ i + ϵ , we simplify this to

We expand the square on the LHS and obtain the following bound on the prediction error

where we used Cauchy-Schwarz inequality in the second step. Now, since D i are diagonal matrices with either 0 or 1 on the diagonals, we can immediately say that ∥ ∥ ϵ T D i X ∥ ∥ 2 ≤ ∥ ∥ ϵ T X ∥ ∥ 2 for all i = 1 , . . . , p . As a consequence, we get

where we used the triangle inequality in the second step. Now, if λ ≥ 2 n ∥ ∥ ϵ T X ∥ ∥ 2 , we can get rid of the second term, namely ( 2 n ∥ ∥ ϵ T X ∥ ∥ 2 -λ ) ∑ p i =1 ∥ ˆ w i ∥ 2 and obtain a bound on the prediction error that does not depend on ˆ w . To this end, notice that ϵ T X is a multivariate Gaussian with mean 0 and covariance σ 2 X T X . We now present a standard concentration result for the norms of Gaussians in the lemma below,

By applying the matrix inversion lemma [53], we can rewrite this as

Lemma A.1. Let z be a d -dimensional multivariate Gaussian with mean 0 and covariance matrix Σ . We have

Applying Lemma A.1 to ϵ T X , we get

√

Plugging in z = n 2 λ where λ = tσ tr( X T X ) n as in the statement of theorem, we get

Therefore, by choosing λ = tσ √ tr( X T X ) n , we have that 2 n ∥ ∥ ϵ T X ∥ ∥ 2 ≤ λ with probability at least 1 -2 e -t 2 / 8 . Consequently, we have

with probability at least 1 -2 e -t 2 / 8 , which completes the proof.

A.7.1 Proof of Lemma A.1

Since the norm of Gaussian vector does not change by applying a rotation, we know that ∥ ∥ Λ 1 / 2 u ∥ ∥ 2 and ∥ z ∥ 2 have the same distribution where Λ is the matrix of eigenvalues from the spectral decomposition Σ = Q Λ Q T and u ∼ N ( 0 , I d ) . Now we apply a standard Chernoff bound argument. We have for some s > 0 ,

By optimizing the bound over s , we get

Finally, using P [ ∥ z ∥ 2 ≤ z ] = 1 -P [ ∥ z ∥ 2 > z ] gives the result.

B Experimental Details

Hyperplane Arrangements for 1D example. In the 1D example in Figure 1, the data matrix X is in R n × 2 , i.e. the dimension d = 2 . This is because we set the second component of each input to be exactly equal to 1 to simulate a bias in the neural network. Though the datapoints lie in R 2 , they have some structure which allows us to enumerate all hyperplane arrangements D X easily. Specifically, the datapoints all lie on the line y = 1 in the xy plane, so the hyperplanes that separate them always separate them into two contiguous segments. Therefore, by sorting the datapoints according to their first component, we can directly enumerate all 2 n hyperplane arrangements.

Nonconvex ReLU network. For the black dashed line in Figure 1, we trained a two layer ReLU network with 100 hidden neurons using the standard weight decay regularized objective with gradient descent using a learning rate of 0 . 01 for 200000 epochs.

Convex program solver specifications. For directly solving the group lasso problem (8), we used CVXPY with the MOSEK solver [54]. For the ℓ 2 regularized least squares problems (16), we used iterative methods with an efficient matrix-vector product implementation to avoid explicitly constructing the large combined masked data matrix ˜ X . Specifically, we used the LSQR and LSMR methods as described in [41, 42]. The experiments were run locally on a MacBook Air 2020 version with the Apple M1 chip.

UCI Machine Learning Repository. Here we present details for the experiments on several real datasets from the UCI Machine Learning Repository [55]. We follow the procedure described in [43] for n ≤ 1000 to extract and standardize the datasets. For these experiments, we use the 75% -25% ratio for the training and test set splits. We use the squared loss and tune the regularization parameter λ for both NTK and our approach by performing a grid search over the set { 10 -5 , 10 -4 , 10 -3 , 10 -2 , 10 -1 , 10 0 , 10 1 } . As shown in Table 1, our method achieves higher (or the same) test accuracy for 26 datasets out 33 datasets while standard NTK achieves higher (or the same) test accuracy for 14 .

C Generic Loss Functions

Here, we prove that the analysis with squared loss presented in the main paper can be extended to arbitrary convex loss function. Let us consider the following regularized two-layer ReLU network training problem with an arbitrary convex loss function L ( · , · )

Then, as shown in [7], the corresponding gated convex reformulation of (22) is

In order to realize the emergence of kernel characterization in the case of arbitrary loss function, we use Lagrange duality. Particularly, since (23) is a convex optimization problem and Slater's conditions holds as proven in [7, 8], we equivalently represent (22) as the following dual problem

where L ( z , y ) is the Fenchel conjugate of L ( · , y ) defined as [56]

We then form the Lagrangian for (24) as follows

where α i ≥ 0 is the Lagrange multiplier for the inequality constraint. Based on the representation above, the dual problem in (24) can be written as

Thus, similar to (12), the problem in (25) can be globally optimized with alternating minimization [33, 35].

D Extensions to Deeper Networks

In this section, we generalize the kernel characterization of two-layer ReLU networks presented in the main paper to deeper architectures. We first note that most prior works on convex reformulation of ReLU networks studied only two-layer ReLU networks, e.g., [6-10], since the non-convex interaction of multiple nonlinear layers hinders the convex analysis of deeper networks. However, recent followup studies [13, 57, 58] showed that the analysis can be extended to arbitrarily deep networks. Based on this recent results, below we briefly explain how one can extend our kernel characterization to three-layer networks. By following the convex reformulations in [57, 58], we have the following gated ReLU training problem for three-layer network

where the hyperplane arrangements for the first and second ReLU layers, i.e., denoted as D (1) i ∈ D (1) X and D (2) j ∈ D (2) X , are defined as follows

Then, following the derivations in Section 4, we can conveniently express the corresponding feature matrices of these masking feature maps on the data X in terms of fixed diagonal data masks as Φ ij = D (1) i D (2) j X . Similarly, the corresponding masking kernel matrices take the form K ij = D (1) i D (2) j XX T D (1) i D (2) j . The same interpretation also extends to arbitrarily deep networks and yields a more complicated feature map due to the interaction of multiple hyperplane arrangement matrices, e.g., D (1) i D (2) j D (3) k . . .

E Cone Decomposition

Here, we briefly review the cone decomposition concept introduced by [8]. We first note that (3) is the exact convex reformulation of the standard non-convex regularized training problem in (2). However, as mentioned in the main paper, the number of constraints in (3) can be prohibitively large, precisely can be O ( n d ) , and consequently prevents solving the exact convex optimization problem especially for large scale datasets. To remedy this issue, [8] proposed the unconstrained formulation in (8), but, this relaxation corresponds to another neural network architecture called gated ReLU network in (5). Although this relaxation breaks dependence between the arrangements and preactivations of a ReLU activation, [8] proved that once the gated ReLU network training problem in (8) training problem is solved, one can construct a solution for the ReLU network model in (1). This procedure is named as cone decomposition and we provide details regarding the implementation steps below.

Let us assume that we solve the gated ReLU network training problem in (8) and denote the optimal solution as { w ∗ i } p i =1 . Then, since each D i Xw ∗ i can have both negative and positive entries (due to the absence of nonnegativity constraint ), we need to decompose D i Xw ∗ i into two separate cones to maintain nonnegativity property of a ReLU activation. We formulate this decomposition step as the following convex optimization problem

After solving the convex optimization problem in (28) for all i ∈ [ p ] . We declare the optimal solutions, { v ∗ i , u ∗ i } p i =1 as a solution to the exact ReLU model in (1). Thus, instead of directly solving the convex reformulation with extremely large number of constraint in (3), we first solve the relaxed gated ReLU problem in (8) and then map the solutions to the proper cones to enforce the nonnegativity property of the ReLU activation.

E.1 Computational Complexity

Although [8] reduce the computational complexity by eliminating the constraints in the convex optimization problem in (3), the number of variables, denotes as p , can still be exponentially large in feature dimension d . However, there are multiple ways to remedy this issue.

We can change the architecture. Particularly, we can replace fully connected networks with convolutional networks. Then, since CNNs operate on the patch matrices { X b } B b =1 instead of the full data matrix X , where X b ∈ R n × h and h denotes the filter size, even when the data matrix is full rank, i.e., r = min( n, d ) , the number of hyperplane arrangements p is upperbounded as p ≤ O ( n r c ) , where r c := max b rank( X b ) ≤ h ≪ min( n, d ) . For instance, let us consider a CNN with 3 × 3 filters, then r c ≤ 9 independent of n, d . As a consequence, weight sharing structure in CNNs dramatically limits the number of possible hyperplane arrangements and avoids exponential complexity. This also explains efficiency and remarkable generalization performance of CNNs in practice.

We can also use a sampling based approach where one can randomly sample a tiny subset of all possible hyperplane arrangements and then solve the convex program with this subset. Thus, although the resulting approach is not exact, the training complexity will not be exponential in d anymore. The experimental results in Appendix C show that this approximation in fact works extremely well, specifically resulting in models that outperform the NTK in 26/33 UCI datasets as detailed in Table 1.

Peer reviews:
comments: SUMMARY:
The paper provides a connection of the convex program for gated ReLU networks to multiple kernel learning model with a weighted data masking feature map. Additionally, the paper provides a theoretical analysis of the predictive error of the proposed kernel algorithm.

STRENGTHS:
The paper provides a new framework that employs multiple kernel learning which is a convex reformulation of neural network training process, the new framework is better able to explain the remarkable performance of neural networks than the neural tangent kernel perspective, which requires the infinite width limit assumption.

WEAKNESSES:
As I am not familiar with multiple kernel learning and group lasso, I feel that I am not able to fully identify the weakness in this paper. I will give a borderline rating with a low confidence score.

QUESTIONS:
No.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work considers a convex formulation of a finite-width regularised two layer ReLU network and interpretations as multiple kernel learning. This then is related to the neural tangent kernel.  The convex formulation considers cones of parameters with fixed activation along with a bound O(n^d). The gated ReLU further decouples the activation region from the linear function. A main conclusion is that while the NTK of the gated ReLU does not take labels into account, optimal Multiple Kernel Learning can.\n\nSTRENGTHS:\n* The work deals with a relevant problem of understanding the optimization problem of neural networks beyond the infinite width setting. \n* The work gives an overview of different perspectives about ReLU networks, particularly previous work on convex programs, group lasso, multiple kernel learning, and NTK, which are interesting as they could facilitate further perspectives. \n* The work discusses convex formulations of gated ReLU networks from a perspective of multiple kernel learning.\n\nWEAKNESSES:\n* The article indicates that previous works have considered training dynamics with infinitesimally small learning rates on networks in the infinite width limit. It would be indicated to comment on works covering small width [1,2], or works that discuss NTK more independently of the width [3], or works that investigate the dynamics of finite-width networks using NTK [4].   \n    * [1] E et al. A comparative analysis of optimization and generalization properties of two-layer neural network and random feature models under gradient descent dynamics.  \n    * [2] Su and Yang. On learning over-parameterized neural networks: A functional approximation perspective.  \n    * [3] Liu et al. On the linearity of large non-linear models: when and why the tangent kernel is constant. \n    * [4] Bowman and Montufar. Spectral Bias Outside the Training Set for Deep Networks in the Kernel Regime. \n* I missed a more explicit discussion of how the presented work is significantly different from or significantly advances in relation to prior works on convex formulations of ReLU networks, particularly the sequence of works [1,2,…]. \n    * [1] Ergen and Pilanci. Convex geometry and duality of over-parametrized neural networks.\n    * [2] Pilanci and Ergen. Neural networks are convex regularizers: exact polynomial time convex optimization formulations for two layer networks. \n* A more explicit connection between the presented analysis and the result of training a neural network would have made this a stronger contribution. Particularly, the conversion of gated ReLU to ReLU should be discussed more explicitly in the main part of the document. The theoretical discussion of fixing the kernel or improving the weights seems to focus on the training error with a more explicit discussion of generalization error missing. Figure 1 and the first experiments seem to be on toy examples. The experiments on UCI data compare NTK vs iteratively reweighed least squares but seem to be missing comparison with the trained network as well as reporting training accuracy.\n\nQUESTIONS:\n* Please clarify the intended notation for w^{(+)} and w^{(-)} or why (11) should have the same expressive power as (1). Concretely, (11) is a signed sum of 2m ReLUs, whereas (1) is a sum of just m ReLUs. For d=2 and m=1, (11) can represent functions which have 4 linear regions, whereas any function represented by (1) has at most 2 linear regions.\n\nLIMITATIONS:\n* The work seems to consider minimal complete gate sets. \n* The work does not appear to include observations into overfitting or generalisation.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies gated ReLU networks with L2 regularization. The authors show that this model is equivalent to Multiple Kernel Learning with group lasso, which is a convex optimzation problem. Thus L2 regularized gated ReLU networks is equivalent to learning the NTK according to a Lasso objective and then fitting the data with the resulting learned NTK. In contrast, in the NTK limit the model fits the data with the initial NTK, which is in general not optimal w.r.t. to the optimization procedure defined previously.\n\nSTRENGTHS:\nThe article is well written and studies the important question of whether DNN can be interpreted as a kernel learning model. The theoretical results appear to be correct, and the numerical experiments illustrate the theory well.\n\nWEAKNESSES:\nThe paper studies gated ReLU networks as an approximation of ReLU networks, but a lot of either very similar or stronger results to the ones proven here are already proven for actual ReLU networks. In particular, [Francis Bach, Breaking the curse of dimensionality with convex neural networks, 2017], proves a similar convex reformulation, and though not explicitely stated, it can also be interpreted as a kernel learning objective (arguably this optimization over kernels is more obvious in [A. Jacot, E. Golikov, C. Hongler, F. Gabriel, Feature Learning in L2-regularized DNNs: Attraction/Repulsion and Sparsity, NeurIPS 2022] which generalizes Francis Bach's results  to the deep case). Francis Bach also proves generalization bounds which are missing here.\n\nThe only possible advantage of the gated ReLU setting is that the optimization might not be NP-hard as is argued in the paper (we know the NP-hardness of the equivalent reformulation of ReLU networks), though this relies on the fact that the number of required gates $p$ grows polynomially, which is not proven. Actually I think it is quite likely that optimization proposed in this paper is NP-hard as well.\n\nAnother related problem is that the gated ReLU optimization is arguably `less unique' than the traditional ReLU, since the dates $g_i$ can always be changed without changing the activations on the training data, thus changing the learned function and kernel outside the training set. It appears that the solution proposed here is to arbitrarily choose a gate direction $g_i$ for each activation pattern on the training set, which is not ideal. Also this problem is not really addressed in the paper. Note that the same problem appear in ReLU networks though it seems less severe, since there can be case where you obtain uniqueness of the solution, in contrast it appears that there is never uniqueness with gated ReLU.\n\nQUESTIONS:\nI also do not like how the authors say in the abstract and introduction that the NTK at initialization is suboptimal compared to the learned NTK, without clarifying w.r.t. to which cost. This could suggest that they prove that the learned NTK is better in terms of generalization, which is not proven. Rather the learned NTK is optimal w.r.t. to the MKL optimization with group Lasso, which is a very specific optimization introduced in this paper (and previous paper) as a reformulation of the L2 regularization loss. So in some sense the statement that the learned NTK is optimal in contrast to the initial NTK is a tautology, since it is w.r.t. to a cost that was designed to reflect the effect of learning.\n\nIf we are instead talking about generalization, then we cannot say that the learned kernel is always better than the initial kernel, as there could be some form of `kernel overfitting'. This is the `no free lunch' theorem: there is no statistical model that is strictly better than another, each model is better under certain assumptions on the task. For example we already know that learning with a constant NTK is optimal if the true function $f^*$ is a Gaussian process with covariance equal to the limiting NTK, in such a setting, NTK learning would be detrimental. For these reasons, I ask the authors to be more specific when thay say that the learned kernel is optimal and initial NTK is suboptimal.\n\nLIMITATIONS:\nAs already mentioned, the authors should be clearer about the two following possible issues:\n- That the reformulation may not be that much faster, because the number of required kernels $d$ one has to optimize over could be very large. It is therefore not clear that gated ReLU offer a significant optimization advantage over ReLU networks.\n- That the optimization is not unique outside the training set.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

comments: SUMMARY:
This work presents the insight that the convex formulation of training a gated ReLU network is an instance of Multiple Kernel Learning (MKL) techniques. This contrasts with the NTK limit, which becomes a single kernel learning in the infinite width limit. The main thesis is that the finite-width ReLU network is necessarily superior over the NTK limit, since MKL learning optimizes over all linear combinations over the multiple kernels, but the NTK limit corresponds to a single fixed linear combination of these kernels.

STRENGTHS:
The presentation and the contribution of this paper are very straightforward: (convexified gated ReLU) = (MKL). This main insight, in my view, the key and only non-trivial insight of this work. In my view, all other parts of this work are relatively straightforward consequences of this observation, and the experiments are minimal and compact.

Therefore, in my view, the decision comes down to whether this single insight is a sufficient contribution to warrant a NeurIPS publication. In my view, it definitely is. How the hidden convex structure of NN approach compares with the NTK or the mean-field analysis was an important problem to tackle, and I find the insight of this work to be a very good start. (Although this paper doesn't yet deal with the mean-field analysis.) I recommend this paper be accepted.

WEAKNESSES:
.

QUESTIONS:
.

LIMITATIONS:
.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper establishes the equivalence between shallow neural networks activated by ReLU and multiple kernel learning (MKL) through the convex reformulation of shallow neural networks into gated ReLU networks. By interpreting neural networks in this way, it becomes apparent that the network may not be able to attain the optimal MKL kernel on the training set. To address this issue, the authors propose an iterative reweighted method to solve the corresponding MKL problem, which has the potential to obtain the optimal MKL kernel due to its convex nature. Moreover, the authors conduct an analysis of the in-sample prediction error for the gated ReLU networks.\n\nSTRENGTHS:\nPreviously, a series of recent studies have demonstrated that (shallow) ReLU networks can be reformulated as a convex problem. This paper builds upon the previous work by establishing the equivalence between this network formulation and the mask kernel learning (MKL) model. This finding has the potential to offer valuable insights into comprehending the performance of neural networks during both training and testing stages.\n\nThe paper is well-written, providing an ample and lucid introduction to various topics such as the convex formulation of neural networks, mask kernel learning, NTK, and other related subjects.\n\nWEAKNESSES:\nThe main contribution of this paper is to establish the equivalence between neural networks and multiple kernel learning (MKL) models. However, there are certain assumptions and setups in the paper that may limit the scope of their findings:\n\n1. The convex formulation of neural networks appears to be applicable only to ReLU (or its variants) activation functions, as the cone constraints rely on the non-negativity of ReLU. Consequently, the results might be specific to ReLU and may not generalize to other activation functions.\n2. Since the output layer of the neural network is fully connected, it can be expressed in a feature map format as f(x) = phi(x)^T v, which holds true even for deep neural networks. Thus, merely establishing equivalence to a kernel method may not sufficiently uncover the intricacies of neural networks.\n3. In lines 164-167 of the proof for Lemma~3.1, the authors fail to explain why the simplex constraint can be eliminated through a change of variables. This omission could lead to confusion for readers.\n4. Theorem 4.3 demonstrates that problem (7) can be reformulated in the form of MKL, but the original problem (7) does not include a simplex constraint. However, the authors do not address this discrepancy in the main theorem or the supplementary proof.\n5. The results presented in section 6 do not offer further insights or explain the benefits of formulating the problem as convex optimization or considering it as MKL. This is because the authors solely analyze in-sample prediction, while previous research has shown that shallow ReLU neural networks can generalize well to various underlying distributions (e.g., see https://arxiv.org/abs/1901.08584).\n\nQUESTIONS:\n1. What is the motivation behind using a regularized formulation? Do the main results (Theorem 4.3, Theorem 5.1, Theorem 6.1) hold for the non-regularized version with lambda set to zero?\n2. Is it possible to extend the convex formulation to activation functions other than ReLU? If so, what would be the considerations and challenges involved in adapting the formulation to non-ReLU activations?\n3. Given that the reformulation is convex, why do you suggest using the Iteratively Reweighted Least Squares (IRLS) algorithm instead of a commonly used convex problem solver? Are there specific benefits to using IRLS in this context? Furthermore, can you provide insights into the properties of the solution that IRLS converges to?\n\nLIMITATIONS:\nSee weakness", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-12"}