SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rajat Vadiraj Dwaraknath
Stanford University rajatvd@stanford.edu
Tolga Ergen  
LG AI Research  
tergen@lgresearch.ai
Mert Pilanci
Stanford University
pilanci@stanford.edu
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 estimizable 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.
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 \times k$  as  $\mathbf{I}_k$ . To denote the set  $\{1, 2, \dots, n\}$ , we use  $[n]$ . We also use  $\| \cdot \|_p$  and  $\| \cdot \|_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  $(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
$$
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  $\pmb{\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\| \sum_ {j = 1} ^ {m} \left(\mathbf {X} \mathbf {w} _ {j} ^ {(1)}\right) _ {+} 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 {2}
$$
where  $\lambda > 0$  is the regularization coefficient. We also use  $f(\mathbf{X},\pmb{\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.
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
$$
\min  _ {\mathbf {w} _ {i}, \mathbf {w} _ {i} ^ {\prime}} \left\| \sum_ {i = 1} ^ {p} \mathbf {D} _ {i} \mathbf {X} \left(\mathbf {w} _ {i} - \mathbf {w} _ {i} ^ {\prime}\right) - \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 .} \quad \begin{array}{l l} (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}, \forall i, \end{array} \tag {3}
$$
where  $\mathbf{D}_i\in \mathcal{D}_{\mathbf{X}}$  are  $n\times n$  binary masking diagonal matrices given by  $p\coloneqq |\mathcal{D}_{\mathbf{X}}|$  and
$$
\mathcal {D} _ {\mathbf {X}} := \left\{\operatorname {d i a g} \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:  $(\mathbf{Xw}_i)_+ = \mathrm{diag}(\mathbb{1}\{\mathbf{Xw}_i\geq 0\})\mathbf{Xw}_i = \mathbf{D}_i\mathbf{Xw}_i$  provided that  $(2\mathbf{D}_i - \mathbf{I}_n)\mathbf{Xw}_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 making 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}} (\mathbf {x}, \boldsymbol {\theta}) := \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} \operatorname {d i a g} \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\{\operatorname {d i a g} \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  $|\mathcal{G}| = |\mathcal{D}_{\mathbf{X}}| = 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} \operatorname {d i a g} \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} \| \mathbf {w} _ {i} \| _ {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.
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 \to \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(\mathbf{x}, \boldsymbol{\theta})$  by linearizing it with respect to the parameters  $\boldsymbol{\theta}$  around its initialization  $\boldsymbol{\theta}_0$ :
$$
\hat {f} (\mathbf {x}, \boldsymbol {\theta}) \approx f (\mathbf {x}, \boldsymbol {\theta} _ {0}) + \nabla_ {\boldsymbol {\theta}} f (\mathbf {x}, \boldsymbol {\theta} _ {0}) ^ {T} (\boldsymbol {\theta} - \boldsymbol {\theta} _ {0}).
$$
The authors of [1] show that if  $f$  is a neural network (with appropriately scaled output), and parameters  $\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  $\phi(\mathbf{x}) = \nabla_{\theta} f(\mathbf{x}, \theta_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  $\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(\mathbf{x}, \mathbf{x}') \coloneqq \mathbb{E}\left[\nabla_{\hat{\theta}} f(\mathbf{x}, \hat{\theta}_0)^T \nabla_{\hat{\theta}} f(\mathbf{x}', \hat{\theta}_0)\right]$ , where  $\hat{\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} _ {\mathcal {G}} (\mathbf {x}, \boldsymbol {\theta}) := \frac {1}{\sqrt {2 m}} \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 {9}
$$
In the next lemma, we provide the infinite width NTK of the scaled gated ReLU network in (9).
Lemma 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}(\mathbf{0},\mathbf{I_d})$  and randomly initialized parameters as  $\mathbf{w}_j^{(1)}\sim \mathcal{N}(\mathbf{0},\mathbf{I_d})$  and  $w_{j}^{(2)}\sim \mathcal{N}(0,1)$  is
$$
H \left(\mathbf {x}, \mathbf {x} ^ {\prime}\right) := \frac {1}{2 \pi} \left(\pi - \arccos  \left(\frac {\mathbf {x} ^ {T} \mathbf {x} ^ {\prime}}{\| \mathbf {x} \| _ {2} \| \mathbf {x} ^ {\prime} \| _ {2}}\right)\right) \mathbf {x} ^ {T} \mathbf {x} ^ {\prime}. \tag {10}
$$
Additionally, we introduce a reparameterization of the standard ReLU network (1) with  $\pmb{\theta} \coloneqq \left\{(\mathbf{w}_j^{(+)}, \mathbf{w}_j^{(-)})\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 {2 m}} \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  $\phi_i: \mathbb{R}^d \to \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  $\Phi_i := [\phi_i(\mathbf{x}_1)^T; \ldots; \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 := \Phi_i \Phi_i^T$ . A convex combination of these kernels can be written as  $\mathbf{K}(\boldsymbol{\eta}) := \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 \mathbf{0}\}$  is a set of weights in the unit simplex. By noticing that the feature map corresponding to  $\mathbf{K}(\boldsymbol{\eta})$
<sup>1</sup>All the proofs and derivations are presented in the supplementary material.
is obtained by taking a weighted concatenation of  $\phi_{i}$  with weights  $\sqrt{\eta_i}$ , we can write the MKL optimization problem in terms of the feature matrices as
$$
\min  _ {\boldsymbol {\eta} \in \Delta_ {p}, \mathbf {v} _ {i} \in \mathbb {R} ^ {d _ {i}}} \left\| \sum_ {i = 1} ^ {p} \sqrt {\eta_ {i}} \boldsymbol {\Phi} _ {i} \mathbf {v} _ {i} - \mathbf {y} \right\| _ {2} ^ {2} + \hat {\lambda} \sum_ {i = 1} ^ {p} \| \mathbf {v} _ {i} \| _ {2} ^ {2}, \tag {12}
$$
where  $\hat{\lambda} > 0$  is a regularization coefficient. For a set of fixed weights  $\pmb{\eta}$ , the optimal objective value of the kernel regression problem over  $\mathbf{v}$  is proportional to  $\mathbf{y}^T\left(\sum_{i=1}^{p} \eta_i \mathbf{K}_i + \hat{\lambda} \mathbf{I}_n\right)^{-1} \mathbf{y}$  up to constant factors [16]. Thus, the MKL problem can be equivalently written as the following problem
$$
\min  _ {\boldsymbol {\eta} \in \Delta_ {p}} \mathbf {y} ^ {T} \left(\mathbf {K} (\boldsymbol {\eta}) + \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).
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
$$
\min  _ {\mathbf {w} _ {i} \in \mathbb {R} ^ {d _ {i}}} \quad \left\| \sum_ {i = 1} ^ {p} \boldsymbol {\Phi} _ {i} \mathbf {w} _ {i} - \mathbf {y} \right\| _ {2} ^ {2} + \hat {\lambda} \left(\sum_ {i = 1} ^ {p} \| \mathbf {w} _ {i} \| _ {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} \| \mathbf {w} _ {i} \| _ {2}\right) ^ {2} = \min _ {\eta \in \Delta_ {p}} \sum_ {i = 1} ^ {p} \frac {\| \mathbf {w} _ {i} \| _ {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  $\pmb{\eta}$  as follows
$$
\min _ {\boldsymbol {\eta} \in \Delta_ {p}} \quad \min _ {\mathbf {w} _ {i} \in \mathbb {R} ^ {d _ {i}}} \quad \left\| \sum_ {i = 1} ^ {p} \boldsymbol {\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} \Phi_ {i} \mathbf {w} _ {i} - \mathbf {y} \right\| _ {2} ^ {2} + \lambda \sum_ {i = 1} ^ {p} \| \mathbf {w} _ {i} \| _ {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).
3.2 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-35] to solve the group lasso problem (15) using the following variational formulation of the group  $\ell_1$ -norm [34]
$$
\sum_ {i = 1} ^ {p} \| \mathbf {w} _ {i} \| _ {2} = \min  _ {\boldsymbol {\eta} \in \mathbb {R} _ {+} ^ {p}} \frac {1}{2} \sum_ {i = 1} ^ {p} \left(\frac {\| \mathbf {w} _ {i} \| _ {2} ^ {2}}{\eta_ {i}} + \eta_ {i}\right).
$$
Based on this, we rewrite the group lasso problem (15) as the following minimization problem
$$
\min _ {\boldsymbol {\eta} \in \mathbb {R} _ {+} ^ {p}} \min _ {\mathbf {w} _ {i} \in \mathbb {R} ^ {d _ {i}}} \left\| \sum_ {i = 1} ^ {p} \boldsymbol {\Phi} _ {i} \mathbf {w} _ {i} - \mathbf {y} \right\| _ {2} ^ {2} + \frac {\lambda}{2} \sum_ {i = 1} ^ {p} \left(\frac {\| \mathbf {w} _ {i} \| _ {2} ^ {2}}{\eta_ {i}} + \eta_ {i}\right).
$$
Since the objective is jointly convex in  $(\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} \boldsymbol {\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  $\pmb{\eta}$  is also directly solved by setting  $\eta_{i} = \| \mathbf{w}_{i}\|_{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  $\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-34].
Algorithm 1 Iteratively Reweighted Least Squares (IRLS) for gated ReLU and ReLU networks