Optimal criterion for feature learning of two-layer linear neural network in high dimensional interpolation regime
Authors: Keita Suzuki
Taiji Suzuki
Venue: ICLR 2024 poster (2024)
Keywords: two-layer linear neural network
feature learning
optimal criterion
multi-output linear reggression
high-dimensional setting
Primary area: learning theory

Abstract: Deep neural networks with feature learning have shown surprising generalization performance in high dimensional settings, but it has not been fully understood how and when they enjoy the benefit of feature learning. In this paper, we theoretically analyze the statistical properties of the benefits from feature learning in a two-layer linear neural network with multiple outputs in a high-dimensional setting. For that purpose, we propose a new criterion that allows feature learning of a two-layer linear neural network in a high-dimensional setting. Interestingly, we can show that models with smaller values of the criterion generalize even in situations where normal ridge regression fails to generalize. This is because the proposed criterion contains a proper regularization for the feature mapping and acts as an upper bound on the predictive risk. As an important characterization of the criterion, the two-layer linear neural network that minimizes this criterion can achieve the optimal Bayes risk that is determined by the distribution of the true signals across the multiple outputs. To the best of our knowledge, this is the first study to specifically identify the conditions under which a model obtained by proper feature learning can outperform normal ridge regression in a high-dimensional multiple-output linear regression problem.

Published as a conference paper at ICLR 2024
OPTIMAL CRITERION FOR FEATURE LEARNING OF TWO-LAYER LINEAR NEURAL NETWORK IN HIGH DI- MENSIONAL INTERPOLATION REGIME
Keita Suzuki
Taiji Suzuki
University of Tokyo and Preferred Netowrks, Inc. keitasuzzuii@preeerrees
University of Tokyo and RIKEN AIP tai taijimmist.. wwt.ss....u.tooo.o u-tokyo ac. ip
ABSTRACT
Deep neural networks with feature learning have shown surprising generalization performance in high dimensional settings, but it has not been fully understood how and when they enjoy the benefit of feature learning. In this paper, we theoretically analyze the statistical properties of the benefit from feature learning in a two-layer linear neural network with multiple outputs in a high- high-dimensional setting. For that purpose, we propose a new criterion that allows feature lerning of two-layer lin- ear neural network in a high- nigh-diimessionl setting. Interestingly, we can show that models with smaller values of the criterion generalize even in situations where normal ridge regression fails to generalize. This is because the proposed crite- rion contains a proper regularization for the feature mapping and acts as an upper bound on the predictive risk. As important characterization of the criterion, the two-layer linear neural network that minimizes this criterion can achieve the opti- mal Bayes risk that is determined by the distribution of the true signals across the multiple outputs. To the best of our knowledge, this is the first study to specifically identify the conditions under which model obtained by proper feature learning can outperform normal ridge regression a high-dimensional multiple-output lin- ear regression problem.
INTRODUCTION
Thanks to the benefit of feature learning, neural networks have demonstrated superior generaliza- tion performance across numerous high-dimensional problems, including image classification and natural language processing. This empirical observation is surprising as it conflicts with conven- tional statistical theory, which claims model shouldn't 't be too complex to predict unseen randomly labeled data well. Therefore it has drawn great attention from the community of learning theory (Belkin et al., 2018). To understand the generalization of high-dimensional statistical models, many researchers have studied the generalization mechanism of simpler models such as linear models, which are amenable to thcoretical analysis in over- werrpaaameetrzzdd settings. Thanks to intensive re- search, we' ve gained an understanding of benign overfitting phenomenon in many simple models. In this phenomenon, models generalize well to unseen test data even when they interpolate train- ing data with significant label noise, and
this in linear regression (Bartlett et al.) 2020
Hastie et al. [022 Muthukumar et al. 2020 Negrea et al. [2200] Xu and Hsu 2019), ridge regres- sion (Tiggler and Bartlet 2020;; Dobriban and Wager 2018 Mei and Montanari 2022). logistic
regression (Montanari el al. 2019; Chatterji and Long 2021 and kernel-based estimators (Liang
and Rakhlin 2020; Belkin et al. 2019 Li al. [2022). However, practical use. neural networks Trequently surpass these simpler models in - gigh--imensional problem settings. Hence, understand- ing the benefits of feature learning and dominance over simple models without feature learning remains crucial in order to elucidate the generalization of deep learning. For this reason, many re- searchers have studied feature learning of two-layer neural networks, which is the simplest case of neural networks. However, much of this research is confined to the setting of classification prob- lems (Frei et al. 2022 Cao et al. 2022), with few studies dedicated to regression problems. This is because there is a difficulty in feature learning for over- parameterizzee regression problems that increasing the number of parameters excessively for feature mapping can lead to overfitting. Accord- ingly, it's crucial apply regularization for parameters including feature mapping for appropriate

---

Published as a conference paper at ICLR 2024
feature learning. As a regularization for such parameters, considerable number of researchers focusing on implicit regularization. This form of regularization is automatically induced by opti- mization methods and deep learning architectures. Several studies have examined the influence of implicit regularization on the generalization of two-layer neural networks in high-dimensional re- gression problems (Chatterji al. 2022 Azulay al.) (02T). However, these analyses on implicit regularization aren't i enough to fully understand the benefit of feature learning. This is because such implicit regularization does not always leverage the entire potential of feature learning, that is, isn't necessarily an optimal regularization in terms of feature learning. In particular, the existing work on implicit regularization as listed above missed important properties of feature learning such as the Bayes risk optimality that will be shown in our study. Therefore, the goal of this work is to answer the following question:
Can we design optimal regularization of feature mapping to fully exploit the benefit of feature learning and demonstrate its improvement over simpler models like ridge regression?
1.1 CONTRIBUTIONS
For the goal above, we consider a two-layer linear neural network which the simplest setting of feature learning in a multi-output linear regression problem setting. We use the ridge estimator as the second-layer parameter and design a proper regularization on the first-layer parameter which enables feature learning. Our contribution can be summarized as follows:
We propose a criterion that gives an upper bound of the predictive risk and show that the predictive risk for an estimator with smaller values of the criterion converges to zero. This criterion has a proper regularization term on the first-layer parameter.
An important characterization of the proposed criterion is that the minimizer of the criterion achieves near-optimal Bayes risk under some conditions, which represents the adequacy of the criterion. This is a merit of considering a multi-output setting.
We characterize when the normal ridge estimator does not generalize and show the two-layer neural network selected by the proposed criterion generalizes even such situations using examples. Indeed, the predictive risk of such an estimator can converge to zero even when that of the naive ridge regression does not. This result demonstrates the benefit of feature learning and its dominance over simpler models in high-dimensional regressinn problems.
1.2 RELATED WORK
Benign overfitting of linear models. One line of research sought theoretical understanding of the benign overfitting phenomenon in linear models. HHstic al. (2022) gave an asymptotic eval- uation of the bias and variance of linear models and random feature regression. Additionally, they demonstrated that the predictive risks associated with certain models exhibit a double descent curve as the ratio of parameter size to sample size becomes large. [Bartlttt al. (2020) derived a tight up- per bound for the predictive risk of the minimum-nore interpolator in an over-parametio ser-parametrized setting. They further elucidated the role of the input covariance matrix's 's spectrum in enabling the model to generalize well. Furthermore, Tsigler and Bartlett (020) extended this result to ridge regression. Belkin et al. (2019) showed that the Nadaraya- Watson estimator with singular kernel that perfectly fits training data enjoys the benign overfitting phenomenon, because thanks to singular kernlls, this model fits the training data with spike-like shape. [u and Hsu (20T9) gave evaluation of the predictive risk depending on the number of features in principal component regression and showed much in high dimensional settings. Li et al.
that the number of features shouldn't be reduced (202) showed that in the framework of (Bartlett et al.] (020), noises added to features play a role of regularization and demonstrated that the double descent phenomenon occurs in this model. How- ever, these methods are not optimal and can sometimes fail to generalize in our settings as this study demonstrates.
Benign overfitting of two-layer neural network. Many recent works have characterized the gen- cralization performance of two-layer neural networks. CCatterji et al. (2022) studied the two-layer linear neural network setting we consider here and analyzed the generalization performance of two-layer linear neural network trained by gradient flow. FFei et al. (2022) studied benign overfitting of a two-layer neural network with nonlinear activation function in classification problems. They showed that if the ground truth is linear, a two-layer neural network trained by gradient descent can achieve min-max optimal test emrrr even if its training error is vanishingly small. Cao et al. (2022)
2

---

Published as a conference paper at ICLR 2024
also studied benign overfitting in classification problem settings. They characterized the benign overfitting phenomenon in terms of hignal-tonooin ratio in a convolutional neural network whose activation function is the power of ReLU, Additionally, they demonstrated that if the hignal-tonooise ratio satisfies a certain condition, a two-layer CNN trained by gradient descent achieves arbitrary small training and test loss, and otherwise, the CNN only achieve constant level test loss. Other works (Ba et al. 2022] Damian et al.) 2022) showed that the two-layer neural networks with one gradient step on the first-layer parameter can outperform fixed kernel regression. Especially, [Ba et al. (2022) highlighted the importance of a large learning rate in the early phase of training and demonstrated that if the learning rate small, the two-layer neural network can't defeat the best lin- ear model. However, these analyses handled only one-step gradient learning, they were motivated by optimizability, and it's unknown what kind of statistical optimality the actual optimal solution satisfies in situations where << d. Indeed, both were considering settings where n. 2 d.
2 PROBLEM SETTINGS AND NOTATIONS
Multivariate linear regression. In this paper, focus the multi-output lin- ear regression problem. Suppose that we observe n training examples Dn = ...... ((())) ..... (Inx (yh) ..... My ) where 2 E Rd and (v." ....,,;'" (m) R", { (y1 (1)
(I)
generated by
( Vi = Pnja; + 6 (1) j=
with the unknown signals B.g (j = 1,...,,, m) that we want to learn, where (a:) are i.i.d. sample from some distribution Px with mean 0 and covariance matrix Eg. and (e())²
i=1 are noise inde- pendent of 24. This setting is called multivariate linear regression, which has been studied for a long time (Breiman and Friedman 1997 Bai et al. 1992). Indeed, we encounter a multivariate output setting in several situations. Especially, in standard deep learning applications, we typically deal with multivariate output like multi-label classification problems. Our problem setting could be considered as the simplest problem setting to consider such problem.
Assumption. Throughout the paper, we put the following sub-gaussian assumption on the input random variable. We define || to be sub-gaussian norm defined by || inf {t > 0: Texp (z²/t²) s 2} for any random variable Z taking its value in R. and |Z]|gg, = sup, 4/ (1ss, /) |||l p2 for any random variable vector Z taking its value in Rd. Let design matrix X : ===(11 (xI,..., ,x) E Raxd, and then we assume the following condition on the distribution of X. Assumption 1. For some positive constants o O(1) and O(1), the rows of XΣE are sub- gaussian vectors with the sub-gaussian norm at most ag and components of (0) independent and have sub- sub-gaussian norms bounded by a.
Two-layer linear neural network. In this paper, we consider estimating the target linear functions 3 by a two-layer linear neural network which is given by a form of fe(x; W, = 3, Wæ where W e Rddd and B, E Id (i = I,..., m). When the first-layer's 's parameter W is given, we estimate the second layer's parameter B3 as the ridge regression estimator:
m. Note that the first layer's parameter W works linear coordinate transform. Accordingly, we estimate the target function by fa;;; W, (W)) (() kind of W yields better performance.
(W) W (XWT T + nAI,)
(0)
(2.1)
Twx. Wa. Then, our question becomes what
Notations. We write a S b if there exists an absolute constant e such that a S cb, and a Soz if there exists an absolute constant c, that only depends on constant o such that a S c, For positive semidefinite matrix M, we define || M4 E x to be the matrix which is composed of the first columns of M, M E Ruxd-k to be the matrix which is composed of the rest of d - & columns of M. Accordingly, for any matrix M E Rdddd we define the eigenvalues of M in decreasing order by μ4 (M) 2 ... 2 (MM),, Mi. E R to be the submatrix of M which is composed of the upper left k x k components
||
Vx Mæ. For any matrix M E gnxdd we define
3

---

Published as a conference paper at ICLR 2024
of M, and M- E Rd-kxd=k to be the submatrix of M which is composed of the lower right components M. For matrix A, B, B is positive semidefinite, we write
A & B.
MAIN RESULTS
3.1 SELECTING FIRST-LAYER PARAMETER W
Given the ridge estimator for the second layer's parameter, interest is what kind of W yields generalization. A natural strategy choosing corresponding to vanishingly small predictive risk which is unknown priori. For this purpose, aim construct estimator of the predictive risk as a function of W. More specifically, we introduce a criterion R(W) as an upper bound of the predictive risk defined
R(W) = min - B EE
n =1100 XW TW e1| + +111
d² II (Sure (Zwo + +Xla)
where Eue == TWXT XWT and "² is a hyper-parameter. Note that the last term plays role of regularization on feature mapping W as stated in the introduction. For Eur := WE WT, we obtain an upper bound of the predictive risk by bias-variance decomposition follows:
EF (2 (Z - WW.TTBBBB (w)) s B B+V,
==1
where,
(1--WX' WXT (XWTwx" + +ndnn) -'x W W-T3ms
B
W (X +nda)) c00
V
=
EEw
B and V correspond to bias and variance terms respectively. R(W) is essentially important because indeed yields an upper bound of the bias and variance as stated in the following theorem.
Theorem 1. There exist (large) constants C1, C2 and 03 which only depend on a such that I satisfies t E (1, (n/cI,n/ n/4c]]), and o²2 satisfies to² o then it holds that with probability
least 2-e-t/e/// t/a, Se-ean can (log 2)) max[[]8::11 forr all W which satisfies R(W) o² S ng
max e2))
B B+V So, max(R((W) 02,
Theorem asserts that if we take /² moderately large, R(W) plays a role of a criterion to select first-layer representation W because R(W)- 0² can be regarded an estimator for the predictive risk. Therefore, a small function value of R(W) directly leads to small predictive risk, and if we obtain W which makes R(W)- - x² close to zero, the two-layer neural network with such W gen-
eralizes in the sense of predictive risk. The regularization term o "fT (Smm (Sice + Ala)
called "degrees of freedom", which we can see in Widely Applicable Information Criterion (WAIC) Watanabe 2010) or Mallows' Cp Malloww 1973) of ridge regression defined as Í((A)
|1 - X nAla) - x| 20 Σ (É. ALa)
X
This criterion is used to select
optimal X in d < settings. Our proposed function can be regarded as kind of extension in the sense that we can select optimal basis W in d 2 n settings. Note that, the theoretical properties of Mallows' Cp and WAIC are valid only in d & n settings, and so theoretical guarantee of the de- grees of freedom as a penalty term in nigh-dimensional settings is a novelty of our work. Typically, in high dimensional setting where the eigenvalues of Σg have slow decay, the ridge regression with positive À is not necessarily optimal (Wu and Xu) and this kind of classic degrees freedom may not work because the of eigenvalues diverges (Bartlett al. 2020
Tsigler and Bartlett
4

---

Published as a conference paper at ICLR 2024
2020).
Indeed, diverging eigenvalue-sum is essentially important for obtaining benign overfitting. However, in a two-layer neural network setting, the coordinate transformation by W can change the problem into like a kernel regime. The W chosen by R(W) yields an optimal kernel basis in terms of predictive risk, enabling the model to generalize even in situations where it could not generalize without the transformation each other. Such a condition will be distilled to Assumption 3 under which we can show the ridge regression with the optimal W minimizing R(W) gives the Bayes optimal risk.
we will see later. This can happen if the Baj are well correlated to
3.2 EVALUATION OF THE MINIMUM RISK ESTIMATE R(W) - o²
Based on Theorem D we may have an upper bound of the predictive risk by evaluating how small R(W) - x² can be. Here, we give an explicit form of the upper bound by using the eigenvalues of
Σg and the covariance matrix of (B,r)]=1 We let Σ3 := m 2i=1 S. and define of 2 2 oj and U,,..., Ud Rd as eigenvalues of Mig and corresponding eigenvectors respectively. Let
(.1...., v2). Then, we can obtain upper bound of the infimum of the objective R(W) 0² as stated in the theorem below.
Theorem 2. Suppose there exists & S n such that pk+1 (Σ
- S ps (ΣB For
W = diag (Wi...., wa) V such that wi² = nd o there exists a constant e which only depends on o₂ such that, with probability at least - 4on
R(W) - -² So. Σ min {
002 -e (232.213)
:- UNN.
(3.1)
Therefore, as long as the right-hand side of (3.I) is close to O, W in Theorem 2 can make the predictive risk itself close O. This means that the two-layer neural network selected by R(W) can generalize. Qualitatively, this result is due to "good" regularization via feature learning using R(W) on the second-layer parameter. Normal l2 regularization is often not enough because this regularization can be too strong even in a direction where signals have large contribution, which leads to large bias. In an ideal regularization, directions which the signals have large contribution should be weakly regularized, thereby resulting in small bias. Conversely, directions where signals have small contribution ought to be strongly regularized to facilitate small variance. For this ideal regularization, information about the distribution of signals is essential which is usually not accessible. Thanks to degrees freedom as regularization on W, R(W) enables us to make use of the information on the distribution of signals, and this is one of the biggest merits of selecting W by R(W).
3.3 BAYES RISK OPTIMALITY
So far, we have derived upper bound of the predictive risk. In this section, we show that the derived upper bound can achieve a lower bound under some conditions using the optimal Bayes risk. As stated in the introduction, this is the first attempt to show that optimized two-layer neural network can achieve Bayes risk optimality. We suppose Σg is positive definite. To construct a lower bound of the predictive risk of the two-layer linear neural network we consider here, we introduce the
Ey (x=NN(X, [13. B1B3_ ]] where
optimal Bayes risk Rop (X,o)=
inf Bg,~A((0L,,) Acestimator
X= (zi...., ) is given training input data, Y X3, + with c; ~ NO, d²). and the infimum is taken over all estimators depending on (X, Y). The Bayes risk coincides with the predictive risk of the multi-output linear regression averaged over the multiple-output since the Bayes risk is an average of risks over the choice of B.. Therefore, it is natural to expect that attaining the Bayes optimal risk directly leads to minimizing the predictive risk of the multiple-output regression. The Bayes estimator 8Bayes is defined as the minimizer that attains the infimum.
Proposition 1. For positive definite matrices Σg, Eg, we suppose x ~ N(0,Ex). B. (0, Eg), &~ (N(,000)) 0²) and training data (I,, y,) E Rd x R. is generated by y = B x + co (EI....,Ea) and y(B,,e) (y1,...,Ba)) the Bayes estimator Baayes
For X = (ZI...., can be explicitly given by
(Ba,e) = (xTx+o'2p)) X ((
5

---

Published as a conference paper at ICLR 2024
This regularization term gg enables ideal regularization discussed above.
Further-
more, for Wn vnd Σ 1/2 we can write this optimal Bayes estimator Bayes
WT (WaXTXX)) + nid) WBXTy y. Therefore. we can represent this Bayes estimator with a two-layer linear neural network and this WB coincides with W in Theorem 2 exceep for constant factor, As we will see later, minimum norm interpolator or more generally, normal ridge regression, does not achieve Bayes risk optimality. Asymptotically, this result coincides with the result of (Wu and Xub when Di and Dg have the same eigenvectors. To construct a lower bound of this optimal Bayes risk, for 8² Var [c]]
we make the following assumptions:
Assumption 2. We assume that the variance of label noise is not too large, i.e., it holds 26 < 1. Assumption We assume that the contribution to the ouiput of the components correspond- ing to small eigenvalues of Ez is not too large, that is, there exists k S n, such that E 5²
E
This assumption ensures that Σ L has a rapidly decreasing eigenvalues, which indicates that alignment between Eg and 2g has anisotropic property and then we may have better performance by finding the informative direction in the first layer. Under these assumptions, we can construct lower bound of the predictive risk in the regime of (Tsigler and Bartlet 2020) as follows:
Theorem Write eigenvectors of ΣE as uB.1....: UB,d E Fd and UB
(UB.I...,,BB uB.d) E Rdxd Suppose that components of the rows of XW UB are independent. 3 for l in Assumption 3 Athere exists some absolute constant e such that
Then under Assumption 2
for any which satisfies < < min {
a and k + 20² + vklo² < n/2, it holds that with 26
probability at least -//0
m
Bx -W Br (w))
T Mi (E (
min W
2 Ropt (x,0)
min
m
Recalling that the upper bound is UNN = Ei min { Me (2) En Theorem 3 claims that the upper bound of the bias and variance of the optimized two-layer neural network (3.1 achieves the lower bound of the optimal Bayes risk as long as a²2 5 (ō²). This result characterizes the effectiveness of feature learning using the proposed criterion R(W). ). In other words, minimizing R(W) can properly find the informative direction, which can be realized thanks to the multiple- output setting. This result indicates that the degrees of freedom is the optimal regularization of feature mapping W in terms of feature learning in the multivariate linear regression problem. Even though R(W) isn't employed in real-world applications, we can establish that in practical terms, neural networks utilize the full potential feature learning by identifying the implicit regularization that lowers the value of this criterion.
4 COMPARISON WITH NORMAL RIDGE REGRESSION
In this section, we diseuss the generalization performance of the two-layer linear neural network whose first-layer representation W is chosen based on R(W), and compare this to a normal ridge estimator without feature learning. Here, we compare the predictive risk of the selected two-layer linear neural network with that of the normal ridge regression (i.c., W I). For that purpose, we derive an upper bound of the predictive risk of the normal ridge regression. This is the first attempt to clarify when two-layer linear neural network surpasses normal ridge regression in the setting of high-dimensional, multivariate linear regression problems. To construct an upper bound, first, we define eigenvalues of Σz and their corresponding eigenvectors as Å1 2 2 dd and W1...., Ud E Rd. Accordingly, we define U == (u1..... ua) E Raddd and à? == u E.gu;- In the same way as Tsigler and Bartlet| 2020), for any & n, we define Pk := ndE41 (nd + E.sh Ar)
and (XU, A) == (XU) - (XU) nAln. Then, we can evaluate the predictive risk of normal ridge regression as stated below.
Corollary 1. Suppose it is known that for some < 1 <1-4 4-/2 n/c2, with probability at least 1 6, the condition number of Ag. (XU, X) is at most L, Then there exists constant Ca, which only depends

---

Published as a conference paper at ICLR 2024
on oz. such that for any X E R. k Sn/c,, and t (1,m/cz), if (XU, A) is positive defnite and
Pk O (+), then it holds that with probability least 1 - δ 20e t//5e
(3 -iaal 27
EF
x BNorm + VNorm
m
BNorm
Σ
A
(4.1)
Eize {1, mdt2isa
X2
L max
im 1m
VNom
kto²
to²
(4.2)
L
n
n
PE
i=k+1
where BNorm corresponds to the bias term and VNomm corresponds to the variance term.
This result is extension of the result of Tsigler and Bartlett (020) for multi-output linear regres- sion. k is called effective dimensionality (Bartlett et al.) 2020 TTigler and Bartlett 2020), which stands for the effective complexity of the ridge regression estimator in high-dimensional settings. Notice that À in Corollary | can be negative. In fact. some negative À can show higher general- ization performance than any À >> under some conditions (Kobak al., 2020 Wu and Xu). In particular, when À 0 this ridge estimator coincides with minimum norm interpolator (Bartlett al., 2020, Liang and Rakhlin 2020). In a similar way as (Tsigler and Bartlett (2020). we can construct lower bounds of the bias ZNorM and the variance VNorm (See Lemma O and Lemma in Appendix C). and we can also show that the upper bounds and the lower bounds match under some conditions (See Appendix 3.. Furthermore, selected model and that of normal ridge regression coincide when Σg is isotropic (See Appendix G This result is natural because when Eg is isotropic, the ideal regularization is equivalent to normal l2 regularization, that is, W selected by the criterion R(W) also becomes close to being isotropic, which indicates that the selected two-layer neural network comes close to normal ridge estimator. In the following part of this section, we discuss the condition that the selected two-layer linear neural network outperforms normal ridge regression comparing UNN and BNorm + VNorm- Note that since BNorm and V VNorm achieve their lower bounds, this comparison is reasonable. For simplicity, we assume that Σn and Σg have the same eigenvectors.
demonstrate that the upper bound UNN of the
4.1 WHEN A DECREASE SLOWLY
In this section, we show that there appears separation between feature learning by R(W) and the naive ridge regression when the eigenvalues X; decrease slowly. The situation where the decay of A; slow appears the benign overfitting regime (Bartlett al., 2020; Xu and Hsu 2019). The following argument shows selecting W by R(W) is effective in such a setting, Indeed, for the case
of normal ridge regression, observing Corollary if X is large in comparison with p4 P
XÇ
in the first term of the bias Bxorm can be seen as a weight that controls the impact of X to the bias. However, if the decay of X; is gradual, the weight of x² can be close to 1, thereby enhancing the
effect of the bias. On the other hand, in the context of the two-layer linear neural network which is g22 m
chosen based R(W), the upper bound of Theorem is written UNN Σí= min
This bound implies that the effect of bias λ is no larger than a²² and this prevents the bias term from deteriorating the predictive risk of the model. The example below reflects the difference in the generalization performance of both models discussed above.
Example 1. We suppose vn E and eigenvalues of Ex, decay slowly, that is, we consider the situation as follows:
i n i2n 0. log
0² 1 0² 1,, 2n.
141-n
Then for any k < in Corollary
1 it holds that
log n
BNorm 1, UNN
Example [ claims that the two-layer linear neural network selected by R(W) can generalize even when the decay of Ex is slow and normal ridge regression can't generalize. This result reflects the
7

---

Published as a conference paper at ICLR 2024
effect of the good regularization discussed above, using W obtained by feature learning via R(W). This allows for weak regularization in directions where signals make significant contributions, and strong regularization where the signals' contributions are minor.
4.2 WHEN z AND B., ARE MISALIGNED
In this section, we compare the predictive performance of the selected two-layer neural network and normal ridge regression when x and 94r are misaligned. First, we introduce the notion of alignment and misalignment. We say that and B.i aligned when the directions of large variance for and that of Ba, closely match. In other words, an aligned situation is that the eigenvectors of Σg corresponding to large eigenvalues and those of Σg are close directions. On the other hand, the direction in which x has large variance and that of 3.i are different, that is, if eigenvectors of Eg corresponding to large eigenvalues and those of Σg are orthogonal, we refer to x and f.:; as being misaligned. We may define the alignment in more rigorous way by considering the situation where Di and Eg have the same eigenvectors. In this setting, for Å 2 ... 2 Xa
01 2 2 od then we say z and B.s are aligned, and on the contrary. if of S od. then we say I' and B.. are misaligned. Here, we consider a misaligned situation, that is, of < ... S d². For the case of normal ridge regression, for any < n1, we can write the bias term as BNorm Σ \ x
Ei=++1 doo Since of S I oz R4t1...,, 02 can be large and this leads to large bias. On the other hand, the two-layer neural network selected by R(W) can decrease both bias and variance thanks to good regularization via learning with R(W) as the example below demonstrates.
X
Example 2. We suppose a and Bug are misaligned, that is, we consider the situation as follows:
A = of ie i-d
o² = 1, 0² =
Then for any < n in Corollary 7 holds that
log
Bxorm Z 1,
UNN x
n
This result indicates that the estimator selected by R(W) outperforms the normal ridge estimator in misaligned situations.
4.3
WHEN TAIL EIGENVALUES OF E, ARE L.RGE
In this section, we show that the selected two-layer linear neural network outperforms normal ridge regression when tail eigenvalues Mg are large. In normal ridge regression, for any o(n), the variance term is written as L = + mpt E/=+11 X? We can see that if tail eigenvalues of
Σg are large, that is. if E k X3 is large, then the second term of VNorm can also be large. This indi- cates that normal ridge regression doesn't 't generalize. On the other hand, because of feature learning via R(W), the selected two-layer linear neural network can generalize even when the normal ridge regression can't 't generalize as the following example shows.
Example 3. Suppose that tail eigenvalues of Exg are large, that is, we consider the situation as follows:
isn
A
i>n of
= 12 0²= 1, 11, d= =
1
Then for any n in Corollary 7 it holds that
BNorm + VNomm 21, UNn S n
Qualitatively, this result also reflects the good regularization obtained by R(W). W selected by R(W) enables strong regularization in directions where signals don't contribute significantly, lead- ing decrease in variance without a drastic increase in bias.
4.4 WHEN (+) IS LARGE
In this section, demonstrate that the selected two-layer linear neural network surpasses normal ridge regression when the norm of ( is notably large. In normal ridge regression, the bias can
be written BNorm Et= oooo 22 x
Ei ooo and the weights A2 11 in the first term X
contribute reducing the impact the bias term. However, if (1) is large, that is, oo is large, then
8

---

Published as a conference paper at ICLR 2024
the bias term doesn't 't vanish even if the weight is small. This means that an isotropic regularization is not enough to effectively decrease the effect of bias in such situations. On the contrary, the selected model with good regularization via feature learning by R(W) is capable of bringing both the bias and variance close to zero. The following example demonstrates the difference between the selected two-layer neural network and normal ridge regression.
Example 4. We suppose n E N and norm of y is large, that is, we consider the situation as follows:
x= of {
n
i S n
, ,0 = 1, 0== =1
i-1
i> n
Then for any k < n in Corollary 1 it holds that
B + VNorm 2 1, UNN Sn
5 NUMERICAL EXPERIMENTS
In this section, we conduct numerical experiments to justify our theoretical results. We consider the setting that (d, n,,,m, m) = (1, 000, 500, 1, 000), Σg = diag (1, 2-,21,.. 11,.2,,.,,,,,.. 1,,..,,,,00000 1,000-1), x ~ NN0 Ex), Σg = diag (1.2-a...... 1, 1,000) (a 0,0 10), (a², = (1,2) and (!) Lidd o(( Since the criterion R(W) is non-convex, we optimized this function with gradient Langevin dynamics, which is a noisy gradient descent method aiming to escape from sharp local minima. For a random variable matrix & which satisfies (G))ss Liidd. M(0 1), learning rate rj and inverse temperature parameter B. the update of the gradient Langevin dynamics can be written as
W(++)= w 7 (W(") + V 2n St. In this experiment, we set (1, () (0.1, 100, 000) and
iterated the update 500 times and selected W corresponding to the smallest R(W). Figure 1 shows the comparison of the predictive risks with a base-10 logarithm between the obtained two-layer linear neural network, normal ridge regression, and Bayes-optimal estimator. The regularization parameters for ridge regression were chosen from 1.0, 0.5, 0.1, 0.01, 0.001, 0.0001 to minimize the predictive risk calculated using additional data. When is small and the eigenvalues of Eg de- cay slowly, the proposed method has equivalent performance However, as becomes larger and the eigenvalues of Σ, decay more quickly, it can be seen that the proposed method begins to outperform ridge regression. This reflects the fact that the more isotropic Eg is. the more isotropic the optimal parameter of the first layer becomes, hence the benefit of find- ing informative directions by feature learning is diminished in such cases. [ shows the minimum values of the proposed criterion R(W) -W> o² with base-10 logarithm obtained by gradient Langevin dynamics and predictive risks with a base-10 10 logarithm of the two-layer linear neural network with corresponding W. We can see that the behavior of the function values of R(W) 0² captures that of predictive risk. This result shows that the proposed criterion R(W)- - o² can play a role of an estimator for an upper bound of the predictive risk.
or falls below the ridge regression.
Laster vahu
log of predictive
-45
-05
+1.0
-10
-15
-30
-20
-25
decay speed
decay speed
Figure 1: Coompariso of predictive risks.
Figure 2: R(W) - g² and predictive risk.
6 CONCLUSION AND FUTURE WORK
In this paper, we have considered the feature learning of multi-output linear regression problem using a two-layer linear neural network. We have proposed a criterion with penalty term on feature mapping W like Mallows' Cp and have shown that estimator selected by the criterion leads to making its predictive risk close to zero. Then, we have shown the optimized model by the proposed criterion can achieve nearly optimal Bayes risk. Furthermore, have characterized when normal
9

---

Published as a conference paper at ICLR 2024
ridge regression doesn't generalize and shown that the selected two-layer linear neural network gen- eralizes even in such situations. We state the future directions of this research. On the optimization aspect, have guarantee that whether we can obtain the optimal solution by standard optimization techniques such as stochastic gradient descent because the objective non-convex. Another future direction is replacing the acti- vation function. Although we considered identity in this paper, we usually use nonlinear activation functions such as ReLU in real applications.
minimize the proposed criterion well. It is not obvious
ACKNOWL.EDGEMENT
was partially supported by JST CREST (JPMJCR2015). TS was partially supported by JSPS KAKENHI (20H00576) and JST CREST (PMICR2115);
REFERENCES
S. Azulay, E. Moroshko, M. S. Nacson, B. E. Woodworth, N. Srebro, A. Globerson, and D. Soudry. On the implicit bias of initialization shape: Beyond infinitesimal mirror descent. In International Conference on Machine Learning, pages 168-477. PMLR, 2021.
J. Ba, M. A. Erdogdu, T. Suzuki, Z. Wang, D. Wu, and G. Yang. High-dimensional asymptotics of feature learning: How gradient step improves the representation. In Advances in Neural Information Processing Systems, 2022.
Z. Bai, C. R. Rao, and Y. Wu. M-estimation of multivariate linear regression parameters under convex discrepancy function. Statistica Sinica, 2(1):237-254, 1992.
P. L. Bartlett, P. M. Long, G. Lugosi, and A. Tsigler. Benign overfitting in linear regression. Pro- ceedings of the National Academy of Sciences, 117(48):30063-3007 2020.
M. Belkin, S. Ma, and S. Mandal. To understand deep learning we need to understand kernel learning. In International Conference Machine Learning, pages 541-549. PMLR, 2018.
M. Belkin, A. Rakhlin, and A. B. Tsybakov. Does data interpolation contradict statistical optimality? In International Conference on Artificial Intelligence and Statistics, pages 1611-1619, PMLR, 2019.
L. Breiman and J. H. Friedman. Predicting multivariate responses in multiple linear regression. Journal of the Royal Statistical Society: Series (Statistical Methodology). 59(1):3-54, 1997.
Y. Cao, Z. Chen, M. Belkin, and Q Gu. Benign overfitting in two-layer convolutional neural net- works. arXiv preprint arXiv:2202.06526, 2022.
N. S. Chatterji and M. Long. Finite-sample analysis of interpolating linear classifiers in the overparameterized regime. Journal of Machine Learning Research, 22:129-1, 2021.
N. S. Chatterji, P. M. Long, and P. L. Bartlett The interplay between implicit bias and benign overfitting in two-layer linear networks. Journal of Machine Learning Research, 23(263):1-48, 2022.
A. Damian, J. Lee, and M. Soltanolkotabi. Neural networks can learn representations with gradient descent. In Conference Learning Theory, pages 5413-5452, PMLR, 2022.
E. Dobriban and S. Wager. High-dimensional asymptotics of prediction: Ridge regression and classification. The Annals of Statistics, 46(1):247-279, 2018.
S. Frei, N. S. Chatterji, and Bartlett. Benign overfitting without linearity: Neural network classi- fiers trained by gradient descent for noisy linear data. In Conference on Learning Theory, pages 2668-2703, PMLR. 2022.
Hastic, A. Montanari, S. Rosset, and R. Tibshirani. Surprises in high-dimensional ridgeless least squares interpolation. The Annals of Statistics, 50(2):949 986, 2022.
10

---

Published as a conference paper at ICLR 2024
D. Kobak, Lomond, and B. Sanchez. The optimal ridge penalty for ceal-world tigh-dimensional data can be zero or negative due to the implicit ridge regularization. The Journal of Machine Learning Research, 21(1):6863-687) 2020.
Z. Li, W. J. Su, and D. Sejdinovic. Benign overfitting and noisy features. Journal of the American Statistical Association, pages 1-13, 2022.
T. Liang and A. Rakhlin. Just interpolate: Kernel "ridgeless" regression can generalize. The Annals of Statistics, (8(3):1329-133)) 2020.
C. Mallows. Some Comments on Cp. Technometrics, 15(4):661-675, 1973.
S. Mei and A. Montanari. The generalization error of random features regression: Precise asymp- totics and the double descent curve. Communications on Pure and Applied Mathematics, 75(4): 667-766, 2022.
A. Montanari, F. Ruan, Y. Sohn, and J. Yan. The generalization error of max-mann max-margin lin- ear classifiers: High-dimensional asymptoties the overparametrized regime. arXiv preprint arXiv: 1911v1191,.1154, 2019.
V. Muthukumar, K. Vodrahalli, V. Subramanian, and A. Sahai. Harmless interpolation of noisy data in regression. IEEE Journal on Selected Areas in Information Theory, 1(1):67-33, 2020.
J. Negrea, G. K. Dziugaite, and D. Roy. In defense of uniform convergence: Generalization via derandomization with application interpolating predictors. In International Conference on Machine Learning, pages 7263-7272. PMLR, 2020.
A. Tsigler and P. L. Bartlett. Benign overfitting in ridge regression. arXiv:2009. 14286, 2020.
arXiv preprint
Watanabe. Asymptotic equivalence of Bayes cross validation and widely applicable information criterion in singular learning theory. Journal of Machine Learning Research, 11:3571-3594, 2010.
D. Wu and J. Xu. On the optimal weighted ln regularization in overparameterized linear regression. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, pages 10112-10123, Curran Associates, Inc.
J. Xu and D. J. Hsu. On the number of variables to use in principal component regression. Advances in Neural Information Processing Systems, 32, 2019.


Reviews:
[{"review_id": "qzBXnNuK8J", "reviewer": "Reviewer_pQvZ", "summary": "This paper studies two layer linear neural networks with multiple outputs. The authors propose a penalty term (regularization) and add it to their least squares optimization problem. This penalty term is effective because it enables \"feature learning\" They analyze the estimator in the high-dimensional regime and show that in some scenarios, it can outperform ridge regression.", "questions": "Please see above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Feature learning is a very important problem and is attracting a lot of attention from the deep learning theory community. The problem is very well motivated.\n\n- The penalty term introduced is related to the classic Millow's $C_p$ and WAIC. The properties of these criteria were already studied in the classic statistical setting where dimension is fixed and the number of samples is large. The authors analyze these in the high-dimensional regime. This is interesting on its own.\n\n- The paper is solid, the proofs seem to be correct, and the paper is well written in general.", "weaknesses": "* The authors argue that the proposed method can beat ridge regression. However, how does it compare to ridge regression with optimally tuned regularization? For example, in figure 1, \"Normal Ridge\" corresponds to ridge regression with $\\lambda = 1/n$ which is not the optimally tuned. More importantly, the authors consider a case where the $\\beta$s are not isotropic. In that case, one might try to look at the optimization problem $\\hat\\beta_{\\Omega} = \\min_{\\beta}||y - X \\beta||_2^2 + ||\\Omega \\beta||_2^2$ where $\\Omega$ is a $d\\times d$ matrix. Then, try to tune $\\Omega$ optimally (similar to Wu and Xu 2020). One can also use the min-norm interpolation version of this that was studied in (Sun et al., 2022), section 2.2.\n\n\nYue Sun, Adhyyan Narang, Halil Ibrahim Gulluk, Samet Oymak, Maryam Fazel. Towards Sample-efficient Overparameterized Meta-learning, NeurIPS 2022.\n\n For example, setting $\\Omega \\asymp \\Sigma_\\beta^{-1}$ can exactly give the Bayes optimal in Proposition 1. It is also doing the feature learning because it regularizes the dimensions with small signal power more that the directions with strong signal. What is the benefit of the method proposed by the authors to this? How do they compare? What are the benefits?\n\n\n* Right before section 2, it is mentioned that (Ba et al., 2022) studies the problem in the setting where $n\\geq d$. Is this true? They consider the regime where $d$ and $n$ are in the same order but $d$ can be large or smaller than $n$. How does the results in this paper compare to the results of (Ba et al., 2022) with $O(\\sqrt{n})$ gradient step size?\n\n* How does the proposed method compare to sketched linear regression? See e.g. (Chen et al., 2023).\n\nXin Chen, Yicheng Zeng, Siyue Yang, Qiang Sun. Sketched ridgeless linear regression: The role of downsampling, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G7IFhbuuTI", "reviewer": "Reviewer_hzqz", "summary": "In this paper, the primary objective of the authors is to understand how and why a linear two-layer neural network outperforms classical ridge regression in a multivariate output regression problem, thereby illustrating how, why and when proper feature feature learning can benefit the learning. To that end, the authors propose an optimal regularization for training a two-layer linear neural network. This loss incorporates a regularization term inspired by the Widely Applicable Information Criterion (WAIC). The authors show that this loss represents an upper bound for the predictive regression error. Moreover, the authors show that the upper bound can reach a lower bound using Bayes optimal risk. The authors compare ridge regression with the proposed optimal regularization, highlighting a set of cases and examples where optimal regularization yields better results, showing when good regularization leads to greater learning.", "questions": "1- The sentence: \"This is because there is a difficulty in feature learning ...\" in the introduction seems not to be understandable, at least to me.\n\n2- I think the main question in the introduction is a bit confusing: \"Can we design an optimal regularisation...\". I think this question is interesting, but maybe it's not related to the learning process in real NN. At least, even if implicit regularisation doesn't explain all the phenomena, it is implicitly given by the learning process and the dynamics of learning. Also, this question seems a bit challenging and too broad as I would expect this regularisation term to depend on the data modelling under consideration.\n\n3- I think some explanations would have been welcome for assumption 1. Why do you consider putting the sub-Gaussianity assumption on the rows of $X\\Sigma_x^{-1/2}$ and not on the rows of $X$?\n\n4- The explanation at the end of page 4 and the beginning of page 5 is confusing to me. A lot of information is given, sometimes without any link, and some terms are not explained. In particular, several topics are mentioned in the same paragraph which are not related to each other (decay of eigenvalues for benign overfitting, high dimensional regime considered for $d, n$, ...). Furthermore, the kernel regime is not defined.\n\n5- In numerical experiments, it would have been fairer to tune the regularization parameter for ridge regression and take the best one. What if the poor performance of ridge is only due to a bad choice of the regularization parameter?  It would have been interesting to add the important baseline that would train a linear neural network as usual on the synthetic data and also some real data, at least to show the robustness of the approach with respect to the data distribution and the model considered.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1- The purpose of the paper is quite important. Explaining theoretically why a neural network is able to outperform classical linear models through efficient feature learning is an interesting problem even in the simple framework of a linear two-layer network.\n\n2- The proof that the proposed regularization allows us to obtain the lower bound given by the optimal Bayes is also interesting since it justifies the optimality of the approach showing that the approach allows us to obtain the best performance in this class of functions of linear two-layer networks. In addition, the comparison with ridge regression and the illustration of specific cases where feature learning can dominate a non feature learning approach is also interesting.\n\n3- Even though few experiments have been carried out, they tend to suggest that regularization does in fact enable good performance to be achieved compared with conventional ridge regression.", "weaknesses": "1- The methodology and design of regularization in neural networks seem arbitrary. They suggest that the regularization used in previous cases is applied to neural networks, and by chance, it provides an upper bound of the true risk and constitutes a lower bound. The author postulates that this type of regularization appears implicitly in neural networks. The question is why this regularization is exploited by neural networks and not another that would also allow reaching the Bayes risk. For example, in recursive feature machines where authors apply an optimal feature learning step and a ridge regression step, one could expect results that reach the lower bound. One could say that it is possible to look for regularization that already exists in the literature or create new ones that will also show that it is an upper bound of the predictive loss and lead to a lower bound of the Bayes risk. But how can we ensure that the regularization proposed by the authors is unique and actually the one implicitly used by linear two-layer Neural Network? Otherwise, it is possible to write a lot of studies with different regularizations. The only thing we seem to be sure of is that ridge regression is not optimal.\n\n2- Along the same lines, one could ask whether this regularization is not optimal in the restricted setting considered by the authors (hypothesis on the data, generative model considered for the data, zero mean data, etc.). In a neural network study, one would expect to study the learning dynamics that naturally bring out the learning model and the implicit regularization that could also depend on the type of data and not be generic. And probably not the reverse: propose an arbitrary regularization and showing by its properties (lower bound of Bayes risk, upper bound of predictive risk) that it is the actual regularization achieved by linear neural network.\n\n3- Even though the theoretical contributions are important, it would have been interesting to illustrate the theoretical conclusions a bit more. In particular, it would have been interesting to compare the W learned by a linear NN network (without the regularization proposed (by a gradient descent method, for example)) and compare it to that obtained by regularization. This could represent an empirical argument that the actual learning process follows that of the proposed regularization. It would also be interesting to test the robustness of the conclusions on real data to see how dependent they are on the model considered and not obsolete for real data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ix82Ws8SUG", "reviewer": "Reviewer_Uw19", "summary": "The authors consider a two-layer linear network for multivariate regression. They introduce an upper bound on the standard MMSE loss, which they treat as a surrogate minimization objective for minimizing the predictive loss. They provide a thorough analysis of this bound and show that if the true regression coefficients are misaligned with the data, then a two-layer network trained using the upper bound as a loss function, outperforms ridge regression and closely matches they Bayes optimal estimator. They also provide experiments corroborating their theoretical results.", "questions": "- The theoretical results give sufficient conditions for the superior performance of the network over ridge regression. Do you believe these conditions are also necessary and did you try any numerical experiments in that direction?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Introduction of an upper bound obtained by adding a regularization term (degrees of freedom), which has previously been used in the setting where the number of samples is much larger than the dimension.\n- Theoretical analysis of the proposed upper bound and the neural network trained on it, which provides insight as to how/why it works\n- Characterization of a regime in which the network outperforms ridge regression and closely matches Bayes-optimal performance\n- Clear presentation of the results", "weaknesses": "- Since the bounds in the theorems only hold up to a constant, they only give quantitative results if the misalignment is strong enough for the corresponding bound to be of vanishing order", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UOfTgF1bzT", "reviewer": "Reviewer_8UMM", "summary": "The paper considers solving multi-output linear regression with two-layer linear networks. First, the authors construct $R(W)$ which serves as an (asymptotic) upper bound of the predictive risk. Next, they further bound $R(W)$ by a quantity $U_\\text{NN}$ which depends on the eigenvalues of $\\Sigma_\\beta^\\frac{1}{2} \\Sigma_x \\Sigma_\\beta^\\frac{1}{2}$. Then, it is argued that the Bayes risk can be lower bounded by the same expression. Under certain conditions, it is shown that the two-layer linear neworks obtained by minimizing $R(W)$ can outperform standard ridge regression. Results are supported by an experiment on synthetic data.", "questions": "- What happens in the experiment if we use other stochastic gradient descent algorithms (SGD, Adam, etc.) instead of the Langevin dynamics? Other algorithms are also capable of escaping local minima.\n\n- What is the justification for using $\\lambda = \\frac{1}{n}$ in the experiment? What happens if we optimally tune the ridge parameter?\n\n- As mentioned above, it is hard to judge how restrictive the assumptions are and how often the suggested method improves over ridge regression. Therefore, I would like to see some experiments with real-world data.", "limitations": "", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Results are sound and fully justified.\n\n- Implication of the main results are well explained.\n\n- Comprehensive case study is provided.", "weaknesses": "- There are many additional assumptions stated inside the theorems: \"for all $W$ which satisfies $R(W) - \\sigma^2 \\leq$...\"(Theorem 1), \"there exists $k\\leq n$ such that ..\" (Theorem 2), \"the rows of $X W_B^\\top U_B$ are independent\" (Theorem 3), \"the condition number of $A_k$..\" (Corollary 1). To me, the assumptions looked very artificial, and it was hard to see why it is reasonable to assume those. It would strengthen the paper if the authors could state all necessary assumptions separately before stating the theorems. Without this, it is hard to judge how non-trivial the results are.\n\n- Although the authors have provided Example 1--4 to show the advantage of using $R(W)$ over ridge regression, the settings seem very artificial and restrictive. I wouldn't be too surprised even if there exists a pathological setting where standard ridge regression does not generalize well and there is a quick fix for it.\n\n- This happens again in the experiment; the authors choose a very specific parameters and argue that the proposed method improves over ridge regression. I would be interested to see further justification for considering such setting.\n\n\n**Minor points**\n\n- Theorems are hard to parse. It would be better if the authors can simplify the statement using $o(\\cdot), O(\\cdot)$ notation. For example, in Theorem 1, many unnecessary constants are displayed, making it hard to see why it holds with a vanishing probability.\n\n- The authors should used a different notation for $V$ in Theorem 2 as $V$ is already used for the variance term.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)