Towards Understanding Generalization via Decomposing Excess Risk Dynamics
Authors: Jiaye Teng
Jianhao Ma
Yang Yuan
Venue: ICLR 2022 Poster (2022)
Keywords: generalization
excess risk
stability
dynamics

Abstract: Generalization is one of the fundamental issues in machine learning. However, traditional techniques like uniform convergence may be unable to explain generalization under overparameterization \citep{nagarajan2019uniform}. As alternative approaches, techniques based on stability analyze the training dynamics and derive algorithm-dependent generalization bounds. Unfortunately, the stability-based bounds are still far from explaining the surprising generalization in deep learning since neural networks usually suffer from unsatisfactory stability. This paper proposes a novel decomposition framework to improve the stability-based bounds via a more fine-grained analysis of the signal and noise, inspired by the observation that neural networks converge relatively slowly when fitting noise (which indicates better stability). Concretely, we decompose the excess risk dynamics and apply the stability-based bound only on the noise component. The decomposition framework performs well in both linear regimes (overparameterized linear regression) and non-linear regimes (diagonal matrix recovery). Experiments on neural networks verify the utility of the decomposition framework.

Published as a conference paper at ICLR 2022
TOWARDS UNDERSTANDING GENERALIZATION VIA DECOMPOSING EXCESS RISK DYNAMICS
Jiaye Teng .Jianhao
Jianhoo M Yang Yuan
Institute for Interdisciplinary Information Sciences, Tsinghua University
2 Department of Industrial and Operational Engineering, University of Michigan, Ann Arbor Shanghai Qi Zhi Institute
t tiy20@mails. ts tsinghua. edu,, jianhaolumiche thhm.ihhh. edu yuanyangetsinguust edu
ABSTRACT
Generalization is one of the fundamental issues in machine learning. However, traditional techniques like uniform convergence may be unable to explain gener- alization under overparameterization (Nagarajan & Kolter 2019). As alternative approaches, techniques based on stability analyze the training dynamics and derive algorithm-dependent generalization bounds. Unfortunately, the stability-based bounds are still far from explaining the surprising generalization in deep learning since neural networks usually suffer from unsatisfactory stability. This paper pro- poses a novel decomposition framework to improve the stability-based bounds via a more fine- greeggrridd analysis of the signal and noise, inspired by the observation that neural networks converge relatively slowly when fitting noise (which indicates better stability). Concretely, we decompose the excess risk dynamics and apply the stability-based bound only on the noise component. The decomposition framework performs well in both linear regimes (overparameterizeed linear regression) and non-linear regimes (diagonal matrix recovery). Experiments neural networks verify the utility of the decomposition framework.
INTRODUCTION
Generalization one of the essential mysteries uncovered modern machine learning (Neyshabur al. 2014; Zhang al. 2016 Kawaguchi et al. (2017). measuring how the trained model performs on unseen data. One of the most popular approaches to generalization is uniform convergence (Mohri et al. 2018), which takes the supremum over parameter space decouple the dependency between the training set and the trained model. However, Nagarajan & Kolter| (2019) pointed out that uniform convergence itself might not be powerful enough to explain generalization, because the uniform bound can still be vacuous under overparameterized linear regimes.
One alternative solution beyond uniform convergence is analyze the generalization dynamics, which measures the generalization gap during the training dynamics. Stability-bare bound is among the most popular techniques in generalization dynamics analysis (Lei & Ying, 2020). which is derived from algorithmic stability (Bousquet & Elisseeff) 2002). Fortunately, one can derive non-vacuous bounds under general convex regimes using stability frameworks (Hardt et al., 2016). However, stability is still far from explaining the remarkable generalization abilities of neural networks, mainly due to two obstructions. Firstly, stability- bound depends heavily on the gradient norm in non-convex regimes (Li et al. 2019), which is typically large at the beginning phase in training neural networks. Secondly, stability-based bound usually does not work well under general non-convex regimes (Hardt al. 2016 Charles & Papailiopollos 2018 Zhou et al. 2018b) but neural networks are usually highly non-convex.
The aforementioned two obstructions mainly stem from the coarse-grained analysis of the signal and noise. As Zhang et al. 2016) argued, neural networks converge fast when fitting signal but converge
"Equal contribution.

---

Published as a conference paper at ICLR 2022
loss
standard training
variance training
standard training
parameter
loss
variance training
1 0
0,
0
parameter
parameter
(b)
Figure (a) The gradient norm of standard training (training over noisy data) is larger than that
of variance training (training over pure noise) at initialization g(0) 0 R 0. We denote the
minimizers of the training loss by é. 6r, respectively (where the y-axis is the training loss). (b) The
whole training loss landscape is highly non-convex while the trajectory of the variance training lies in convex region due to the slow convergence. We defer the details to Appendix B
relatively slowly when fitting noise indicating that the training dynamics over signal and noise significantly different. Consequently, the one hand, the fast convergence of signal-related training contributes to large gradient norm at the beginning phase (see Figure Tap. resulting in poor stability. On the other hand, the training signal forces the trained parameter away from the initialization, making the whole training path highly non-conves (see Figure lb). The above two phenomena inspire to decompose the training dynamics into noise and signal components and only apply the stability-based analysis over the noise component. To demonstrate that such decomposition generally holds in practice, conduct several experiments of neural networks both synthetic and real-world datasets (see Figure 2 for more details).
Based the above discussion, we improve the stability- analysis by proposing decomposition framework on excess risk dynamica where we handle the noise and signal components separately via a bias-variance decomposition. risk (VER) and bias excess risk (BER), where VER measures how the model fits noise and BER measures how the model fits signal. Under the decomposition, we apply the stability-basee to VER and apply uniform convergence to BER inspired by Negrea et al. (020). The decomposition framework accords with the thcoretical and experimental evidence surprisingly well, providing that outperforms stability-basee bounds in both linear (overparameterizeed linear regression) and non-linear (diagonal matrix recovery) regimes.
detail, we decompose the excess risk into variance excess
techniques
We summarize contributions follows:
We propose a new framework aiming at improving the traditional stability-based bounds, which is novel approach generalization dynamics analysis focusing decomposing the excess risk dynamics into variance component and bias component. Starting from the overparameterized linear regression regimes, we show how deploy the decomposition framework in practice, and the proposed framework outperforms the stability -based bounds.
We theoretically analyze the excess risk decomposition beyond linear regimes. As a case study, derive a generalization bound under diagonal matrix recovery regimes. To best knowledge, this is the first work to analyze the generalization performance of diagonal matrix recovery.
We conduct several experiments on both synthetic datasets and real-world datasets (MINIST, CIFAR-10) to validate the utility of the decomposition framework, indicating that the framework provides interesting insights for the generalization community.
In this paper, we refer the signal the clean data without the output noise, and the noise the output noise.
See Section 2 for the formal definitions.
We decompose the excess risk, which is closely related to generalization, purely due to technical reasons. The excess risk dynamics tracks the risk during the training process.
2

---

Published as a conference paper at ICLR 2022
0.6
ER
0.5
BER
10
VER
0.3
0.6
0.2
0.4
ER
BER
0.2
VER
0.0
0.0
0
500
1000
1500
2000
2500
3000
10
20
50
Epoch T
Epoch T
(a)
(b)
Figure 2: Experiments of neural networks on synthetic (linear ground truth, 3-layer NN, SGD) and real-world (MNIST, 3-layer CNN, Adam) datasets. The trend of excess risk dynamics (blue) meets its bias component (BER) at the beginning phase and meets its variance component (VER) afterwards, indicating that ER can indeed decomposed into VER and BER. See Appendix for more details.
1.1 RELATED WORK
Stability-base Generalization. The stability-baad researches can be roughly split into two branches. One branch is about how algorithmic stability leads generalization (Feldman & Vondrak 2018: 2019 Bousquet et al. 2020). Another branch focus on how to calculate the stability parameter for specific problems, e.g., Hardt et al. (2016) prove a generalization bound scales linearly with time in convex regimes. Furthermore, researchers try to apply stability techniques into more general settings, e-g., non-smooth loss (Lei & Ying 2020 Bassily et al. 2020). noisy gradient descent in non-convex regimes (Mou et al. 2018
al. 2019) ) and stochastic gradient descent in non-convex | Charles & Papailiopoulos 2018 Zhang et al., 202T). In this paper, we mainly focus applying the decomposition framework to improve the stability-based bound.
regimes (Zhou al. 2018b]
Uniform Convergence widely used generalization analysis. For bounded losses, the gener- alization gap is tightly bounded by its Rademacher Complexity (Koltchinskil & Panchenko, 2000 Koltchinskil 2001 Koltchinskii et al. (2006). Furthermore, we reach faster rates under realizable assumption (reero al., 2010). A line of work focuses uniform convergence under neural network regimes, which is usually related to the parameter norm (Bartlett et al. 2017 Wei Ma 2019). However, as Nagarajan & Kolter (2019) pointed out, uniform convergence may be unable to explain generalization. Therefore, more techniques are explored to go beyond uniform convergence.
Other Approaches Generalization. There
other approaches to generalization, including PAC-Bayes (Neyshabur et al. 2017a| Dziugaite & Roy 2017 Neyshabur et al.) 20176] Dziugaite &
Roy 2018 Zhou et [0188]] Yang al. 2011). information-based bound (Russo & Zou 2016]]
Xu & Raginsky] 2017 Banerjee & Montúfar, [202]] Haghifam et al. 22200] Steinke & Zakynthinou
2020), and compression-baseeeeeeeeeeeeo bound (Arora al. 2018 Allen-Zhu al.
Arora et al. 2019).
Bias- Variacer Decomposition. Bias-variance decomposition plays an important role in statistical analysis (Lehmann & Casella (2066; [asella & Berger] [2021] Geman et al. [992). Generally, high bias indicates that the model has poor predicting ability on average and high variance indicates that the widely used machine learning analysis, e.g., adversarial training (Yu et al., (2021), double descent (Adlam & Pennington] (020). uncertainty (Hu et al. 2020). Additionally, (Nakkiran et al. 2020) propose decomposition framework from an online perspective relating the different optimization speeds on empirical and population loss. Oymak et al. (2019) applied bias- variance decomposition on the Jacobian of neural networks to explain their different performances clean and noisy data. This paper considers slightly different variance decomposition following the analysis SGD (Dieuleveut et al. 2016 Jinn et al. 2018: Zou et al. 2021), focusing on the decomposition of the noisy output.
model performs unstably. Bias- variance decomposition
3

---

Published as a conference paper at ICLR 2022
Matrix Recovery. Earlier methods for solving matrix recovery problems rely the convex relaxation techniques for minimum norm solutions (Recht et al. 2010 Chandrasekaran et al. [011]. Recently, branch of works focuses the matrix factorization techniques with simple local search methods. It has been shown that there no spurious local minima in the exact regime (Ge et al.) 2016 2017]] Zhang et al., 2019). In the overparameterized regimes, it was first conjectured by Gunasekar al. (2018) and then answered by [i et al. (2018) that gradient descent methods converge to the low-rank solution efficiently. Later, Zhuo et 2021 extends the conclusion to the noisy settings.
PRELIMINARY
In this section, we introduce necessary definitions, assumptions, and formally introduce previous techniques in dealing with generalization.
Data distribution. Let æ E X c R² be the input and y E y C. R be the output, where æ, are generated from joint distribution (æ. P. Define Px; py. and Pul= as the corresponding marginal and conditional distributions, respectively. Given n training samples D - (ar Yi kein) generated from distribution P. we denote its empirical distribution by P", To simplify the notations, we define X Raxp as the design matrix and Y E IR" as the response vector.
Excess Risk. Given the loss function 2(0; y) with parameter 0 and sample (x, y), de- fine the population loss P) E(m.y)> [e(0; x. y)] and its corresponding training loss as CO; P") A 1. E, ,0, 2,, vi). Let A, denote the optimization algorithm which takes the dataset D imput and returns the trained parameter g()) at time t, namely, A((D) @(t). During the analysis, we focus on the excess risk dynamics (,((()) P), which measures how the trained parameter ((t) performs on the population loss:
:c( P) C( P) min L(0;P).
Although the minimizer may not be unique, we define PP)":P) ming L(0;P) where 0" denotes arbitrarily one of the minimizers. Additionally, bounding the generalization gap (()()); ()) P)
L((();; PP) suffices bound the excess risk under Empirical Risk Minimization (ERM) framework by Equation equation] [ Therefore, we mainly discuss the excess risk following Bartlett et al. 2020))
Ec(0();)) P) - C() C(0(); Pm)) C(0" -(0" -
PR)
generalization
0 under ERM
zeO by concentration incqualities
VER and BER. As discussed in Section I we aim to decompose the excess risk into the variance excess risk (VER) and the bias excess risk (BER) defined in Definition I The decomposition focuses on the noisy output regimes and we split the output y into the signal component Elz x) and the noise component E (yy jz). Informally, VER measures how the model performs on pure noise, and BER measures how the model performs on clean data.
Definition (VER and BER) Given (x. y) P, let (x. Eg)(a) Ps denote the signal distri- bution and (x, E Elll=l ~ Py denote the noise distribution. The variance excess risk (VER) PO; and bias excess risk (BER) (0;P) P) are defined as:
( P) Ec(O; E2(0:)) & Ec(O:P))
To better illustrate VER and BER, consider three surrogate training dynamics: Standard Training, Variance Training, and Bias Training, corresponding to ER, VER, and BER, respectively. Standard Training: Training process over the noisy data (X, Y) from the initialization g(0), We
denote the trained parameter at time by ((t)
Variance Training: Training process over the pure noise (X, Y - E|) from the initialization
We denote the trained parameter at time by g())
Bias Training: Training process the clean data (X, [[] from the initialization g We
4

---

Published as a conference paper at ICLR 2022
denote the trained parameter time by 0,
0%0 600 are related to the corre- When the context is clear, although the trained parameters ,
sponding initialization and the algorithms, we omit the dependency. Besides, we denote 0", 8; and 0g the optimal parameters which minimize the corresponding population loss L(0:P) C(0; P.), and L(O;P) Pb), respectively.
Techniques in generalization. We next introduce two techniques in generalization analysis includ- ing stability-base baased techniques in Proposition [ and uniform convergence in Proposition These techniques will be revisited in Section
Proposition 1 (Stability Bound from Feldman & Vondrak (2019) Assume that algorithm A, is e-uniformly-stable at time t, meaning that for any two dataset DD and DD' with only one different data point, we have
sup EA (A(((A(D);; (() =Z,
(a.
Then the following inequality hold f with probability least ô:
( ( + log(1/0)
(elog(n)
=
Proposition 2 (Uniform Convergence from Wainwright (2019)
Uniform convergence decouples
the dependency between the trained parumeter and the training set by taking supremum over parameter space that is independent of training data, namely
[(A(D))P) - L(AA(D); S sup -) C(0;
-
BEB
where B is independent of dataset D and A (D) E B for any time t.
3 WARM-UP: DECOMPOSITION UNDER LINEAR REGIMES
In this section, we consider the overparameterized linear regression with linear assumptions, where the sample size is less than the data dimension, namely, p. We will show how the decomposition framework works and the improvements compared to traditional stability- bbased bounds. Before diving into the details, emphasize the importance of studying overparameterized linear regression. As the arguably simplest models, we expect a framework can at least work under such regimes. Besides, recent works (e.g.. aaacot al. (2018)) show that neural networks converge to overparametrized linear models (with respect to the parameters) as the width goes to infinity under some regularity conditions. Therefore, studying overparameterized linear regression provides enormous intuition to neural networks.
Settings. When the context is clear, we reuse the notations Section 2 Without loss of generality, we assume that E|x] 0. Besides, we assume linear ground truth x 0' where [[c[æc] = 0. Let Ex E(æxx denote the covariance matrix. Given &2 loss {(0; æ. y) (y x: 0)2, we apply gradient descent (GD) with constant step size À and initialization ( = O, namely, ((++1) 0(0) + E, (3: xx! ((9)). We provide the main results of overparameterized linear regression Theorem I and then introduce how to apply the decomposition framework.
Theorem (Linear Regression Regimes) Under the overparameterized linear regression settings, assume that w 0- Vo--- is o², -sssuuusssian. If ||x||| < 1, |e| < V. suPre(m] ||0.0| < B are all
bounded, the following inequality holds with probability at least 8:
max {1, 11,0"TE00'00 VV + BJ}
log (4/6)
10
TXV + B
XT
where the probability is taken over the randomness of training data.
In the rest the paper, we use or S to omit the constant dependency and use Ô to omit the log dependency of n. For example, write 2/n (11/n) 2/n $ 1/n, log(n)/n / ((1/n), and log2(n) (n)/n Õ(1/n).
5

---

Published as a conference paper at ICLR 2022
( ) indicating that By setting T 0(vn) in Theorem the upper bound is approximately Ō Vn
the estimator at time T is consisten Below we show the comparison with some existing bounds, including benign overfitting bound and traditional stability- baseed bound
Comparison with benign overfitting, The bound in Bartlett et al. (2000) mainly focus on the min- norm solution (as T 20, the parameter trained in GD converges to the morm solution). Instead, the bound in Theorem ] does not require the min-norm solution assumption.
Comparison with stability-based bound. The stability- bound (without applying the decompo-
sition framework) is Ô [V + BB
log(2/6)
TV + B313) where B' = sup, |
2n
n
denotes the bound of the trained parameter in standard training. As comparison, the notation B in Theorem I denotes the bound of the trained parameter in variance training. One major difference between B and B' is that: B is independent of ||@"| while B' is closely related to ||@"|l Therefore. with a large time T, the bound in Theorem outperforms the stability- baaled bound when | large compared to V (namely, large signal-to-ne ratio We refer to Appendix F. for more discussion.
[
Proof Sketch. We defer the proof details to Appendix Cuu to space limitations. Firstly, we provide Lemma I to show that the excess risk can be decomposed into VER and BER.
Lemma Under the overparameterized linear regression settings with initialization (()) =0,²
0 0. for any time T, the following decomposition inequality holds:
(c(((T); P) s (0(") P) + 2
We next bound VER and BER separately in Lemma 2 and Lemma 3 When applying the stability-
based bound on VER, we first convert VER into generalization gap (see Equation [ and then bound three terms individually.
Lemma Under the overparameterized linear regression settings, assume that |>|| S 1, |c| S V. sup, [10 < B are all bounded. Consider the gradient descent algorithm with constant stepsize X, with probability at least 8:
E (D)) =0
T\ V + B] log(n) log(2n/8) + max {1.1V + BJ} / 1og(2/6) 2n
(2)
n
In Lemma we apply uniform convergence bound BER. Although we can bound it by directly calculating the excess risk, we still apply uniform convergence, because exact calculating is limited to a specific problem while uniform convergence is much more general. To link the uniform convergence and excess risk, we split the excess risk into two pieces inspired by Negrea et al. (2020). For the first piece, which stands for the generalization gap, we bound it by uniform convergence. For the second piece, which stands for the (surrogate) training loss, we derive a bound decreasing with time T.
Lemma Under the overparameterized linear regression settings, assume that w Vo-.T E
is
of suggaaaaasi,n then with probability at least
8:
( -
f
Σ 0
log
2
+
10"0 ||?
(3)
w
X/2
Although we do not present it explicitly, the results in Lemma are derived under uniform convergence frameworks. Combining Lemmal] [ Lemma 2 and Lemma 3 Teads to Theorem [
*Consistency means the upper bound converges to zero as the number of training samples goes to infinity. 'We remark that the assumptions in Theorem can be relaxed. Firstly, assuming is subGaussian suffices show that w is subGaussian. Secondly, assuming is subGaussian suffices to bound |e| S with a log(n) cost. A large time guarantee that the dominating term in Theorem [ TA [V + BT² We remark that with
large signal-to-noise ratio, have TA [V + B²²
+ B]2

---

Published as a conference paper at ICLR 2022
4 DECOMPOSITION UNDER GENERAL REGIMES
This section mainly discusses how to decompose the excess risk into variance component (VER) and bias component (BER) beyond linear regimes. Due fine-grained analysis the signal and noise, the decomposition framework improves the stability- We emphasize that the methods derived in this section are nevertheless the only way to reach the decomposition goal and we leave more explorations for future work. To simplify the discussion, we assume the minimizer 0 of the excess risk exists and is unique. Before diving into the main theorem, it is necessary to introduce the definition of local sharpness for the excess risk function, measuring the smoothness of the excess risk function around the minimizer 0'
bounds under the general regimes.
Definition 2 (Local Sharpness) For given excess risk function (0; with (P) 0. define & > as the local sharpness factor when the following equation hold M.
P)
0
litm
0+1 X0.
0:1 11-1111100 ||~10
10
Assumption (Uniform Bound) Assume that for excess risk function Pp( with local sharp- ness 8, given Pa. for any Pyle and 0, the uniform bounds < m. m. (Pa) S Ma(Pa) exist:
E0;
m. (,,(P=)
10 1-1 S M..).).
When the context is clear, we omit the dependency of Pa and denote the bounds by ma and M
Roughly speaking, the uniform bounds act as the maximal and minimal eigenvalues of the excess risk function (0; P). For example, under the overparameterized linear regression regimes (Section the uniform bounds are the maximum and minimum eigenvalues of the covariance matrix 2iz-
Assumption (Dynamics Decomposition Condition, DDC) We assume the trained parameters @(T) are (a, C, C*) bbonded at time T, namely:
C
C"
[|0(T) 0 S a(110 - -01 + -10 01 +
+
(4)
VTi
where denotes the size of training samples, T denotes the training time, and 6", 0% and 0 denote the minimizers of the corresponding population loss L(O; P). L(0; P.), and L(O: Po). respectively.
be replaced with any o(1) term which converges to zero when T, go to infinity. Empirically, several experiments demonstrate the generality of DDC and its variety (see Figure 3). From the theoretical perspective, we derive that overparameterized linear regression is (1, O, bo-udddd in DDC, and diagonal matrix recovery regimes satisfy DDC (see Section 4.1) We next show how to transfer DDC to the Decomposition
We emphasize that although
C C' in Assumption 2 VT vT
Theorem in Theorem
Theorem (Decomposition Theorem) Assume that the excess risk has a unique minimizer with the following decomposition
local sharpness factor Under Assumption [ and Assumption P inequality holds:
Ma [4a]
4C
(c(O(T)
( E0 + M.
+
Ma
(5)
M.
Vn
We defer the proofs to Appendix D due to limited space. In Theorem 2 Assumption and As- sumption focus on the population Toss and the training dynamics individually. We remark that our proposed generalization bound is algorithn-dependent due to parameters (a, C, C') in DDC.
By Theorem 2] we successfully decompose the excess risk dynamics into VER and BER. The remaining question how bound VER and BER individually. As we will derive in Section[..1 via a case study on diagonal matrix recovery, we bound VER by stability techniques and BER by uniform convergence.
Here || refers the fg-norm.
7

---

Published as a conference paper at ICLR 2022
2.5
standard
12
standerd
varionce
valience
1.3
upper bound
18
blas
0.6
0.1
upger bound
250
2000
200
400
000
1000
1200
1400
550
.000
1:00
2000
2500
200
Reeation
Eguch
(a) linear regression, n = 300
(b) matrix recovery. n = 200
(c) neural network, 500
standard
standard bias upper bound
standard
2.0
F1
variance
variance
bound
upper bound
1911
14
0.0
230
500
750 1030
Reeratonn
1250 1506
1750 2000
202
400
003
Beration
502
1000
1200
1400
550
7000
Edach
11000
2000
2520
3003
(d) linear regression, 800
(c) matrix recovery, n = 600
(f) neural network, 1500
Figure 3: DDC holds in overparameterized linear regression, general matrix recovery, and neural networks. For linear regression and neural networks, we set a = 1,C = C² = 0 in DDC. For matrix recovery problems, we set 1, C 0, C' 8.5 in DDC. The upper bound (the right hand side in DDC, red line) is over the standard (the left hand side in DDC, orange line), indicating that DDC holds in general. We defer the experimental details to Appendix
4.1 DIAGONAL MATRIX RECOVERY
From the previous discussion, we now conclude the basic procedure when applying the decomposition framework under Theorem Generally, there three questions to be answered in a specific problem: (a) Does the problem satisty DDC? (b) How to bound VER using stability? (c) How to bound BER using uniform convergence? To better illustrate how to solve the three questions individually, provide a case study on diagonal matrix recovery, a special and simplified case of general low-rank matrix recovery problem but still retains the key properties. The diagonal matrix recovery regime and its variants are studied as a prototypical problem in Gissin et al. (2019); HaoChen et al. (2021).
Settings. Let X Diag(oi,
O,.......,.. E Rdxd denote rank-r (low rank) diagonal matrix, where of > ... 2 of > O. We define A Rdxd as the mcasurement matrix which is diagonal and y E 1 as the measurement vector. Assume the ground truth yı A. + E; where Ve, A; denotes their corresponding i-th element, denotes the noise, and of = 0f i S d. Our goal is to recover X from dataset D = {A(1, heEM) with training samples.
In the training process, we use the diagonal matrix U = Diag(un,.... ua} to predict X* via UUT Given the loss function (U) Eー= | Ab) @ CUT 1|² we train the model with Gradient FU bounded with unit variance, and E; is -ussuGGausii with uniform bounds | < V. We emphasize that the training procedure of diagonal matrix recovery is fundamentally different from linear regimes since C, (U) is non-convex with respect to U.
V., (U.). During the analysis, we assume that A; is (1))-uniforll
Theorem (Decomposition Theorem in Diagonal Matrix Recovery) In the diagonal matrix re- covery settings, if we set the initialization Un = aT where o > 0 is the initialization scale, then for
*We use @ to represent the Hadamard product.
We use gradient flow instead of gradient descent mainly due technical simplicity.
8

---

Published as a conference paper at ICLR 2022
any 0 <t<< < with probability at least - 8, we have:
log (d/o)
a²d
Ec(Ur; (18-1 dV2 + da+)
of + alelloss))
(6)
(t+ log(n) log
2dn
1
rv log (r/6)
+ da² + log
8
a2
By setting a - ((dn))+ (dfnn) approximately ① cav2 ||
and = θ (log (dno,) ) /a,) Theorem the upper bound is
indicating that the estimator at time t consistent.
vn
Comparison with stability- based bound. The stability- based bound (without applying the decomposi-
tion framework) Ô
dV <V=++IX' + [|x=+| 1,)+++1) As a comparison, the bound in Theorem
vT.
outperforms the stability-based bound in the sense that deduce the coefficient of the blowing-up term 1+11 from dV2 + ||, to dV2, which is significant improvement with a large signal-to-noise ratio under large time t. Besides, the bound in Theorem captures the bias-variance tradeoff, meaning that the excess risk curve falls first and then rises.
Proof Sketch. We defer the proof details to Appendix P due to space limitations. Similar Section the proof of Theorem 3 consists of the following three lemmas, corresponding to the DDC, BER, and VER terms separately. We first prove that diagonal matrix recovery regimes satisfy DDCin Lemma
Lemma 4 Under the assumptions in Theorem 3 for any S t < oc, with probability at least 1 8, we have
|0UT - 1 S 1 - -x+1 U +0
1 a
rv² log (F/8)
log
Different from the overparameterized linear regression regimes, Lemma suffers from a
factor due the non-linearity. This phenomenon can be verified empirically in Figure 3b and Figure 3e We next apply stability techniques and uniform convergence to tackle VER and BER in Lemma and Lemma respectively.
Lemma For VER, under the assumptions in Theorem 3 with probability at least 6, we have
(2(U:)) 2t log(n) log
2dn
log (2d/o)
+ da² + dV²
(8)
Lemma 6 For BER, under the assumptions in Theorem 3
with probability at least - 8, we have
(UU $ (1x*112+ da*) / log (d/8)
a²d
Σ
(9)
o (a
n
Combining the above lemmas leads to Theorem
Directly applying Theorem 2 under overparameterized linear regression regimes costs loss on condition number. The cost comes from the DDC condition where we transform the function space (excess risk function) to the parameter space (the distance between the trained parameter and the minimizer). We emphasize that the field decomposition framework is indeed much larger than Theorem and there other different approaches beyond Theorem 2 For example, informally, by
Remark:
introducing pre-conditioner matrix A, changing DDC from the form | #* || A0 and setting A = /2/ one can avoid the cost of the condition number. We leave the technical discussion for future work.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
This work has been partially supported by National Key R&D Program of China (2019AAA0105200) and 2030 innovation megaprojects of china (programme on new generation artificial intelligence) The authors would like to thank Ruiqi Gao for his insightful sugges-
Grant No. tions. We also thank Tianle Cai, Haowei He, Kaixuan Huang, and, Jingzhao Zhang for their helpful discussions.
REFERENCES
Ben Adlam and Jeffrey Pennington. Understanding double descent requires a fine- lne-grained bias- variance decomposition. arXiv preprint 1rXiv:201..0332) 2020.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. arXiv preprint arXiv: (811v999 2018.
Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via a compression approach. In International Conference on Machine Learning, pp. 254-263. PMLR, 2018.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pp. 322-332 222.332. PMLR, 2019.
Pradeep Kr. Banerjee and Guido Montúfar. Information complexity and generalization bounds, 2021.
Peter Bartlett, Dylan J Foster, and Matus Telgarsky. Spectrally-normaaizze margin bounds for neural networks. arXiv preprint arXiv: 10K::1706608498, 2017.
Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117749))30063-30077, 2020.
Raef Bassily, Vitaly Feldman, Cristóbal Guzmán, and Kunal Talwar. Stability of stochastic gradient descent nonsmooth convex losses. arXiv preprint arXiv:2006.06914 2020.
Olivier Bousquet and André Elisseeff. Stability and generalization. The Journal of Machine Learning Research, 2:499-524 522,,226222222 2002.
Olivier Bousquet, Yegor Klochkov, and Nikita Zhivotovskiy. Sharper bounds for uniformly stable algorithms. In Conference on Learning Theory, pp. 610-626. PMLR, 2020.
George Casella and Roger L Berger. Statistical inference. Cengage Learning, 2021.
Venkat Chandrasekaran, Sujay Sanghavi, Pablo A Parrilo, and Alan S Willsky. Rank-sparsity incoherence for matrix decomposition. SIAM Journal on Optimization, 21(2):572-566, 2011.
Zachary Charles and Dimitris Papailiopoulos. Stability and generalization of learning algorithms that converge to global optima. International Conference on Machine Learning, pp. 745-754. PMLR, 2018.
Aymeric Dieuleveut, Francis Bach, al. Nonparametric stochastic approximation with large step- sizes. Annals of Statistics, 1344)))363-399) 2016.
Gintare Karolina Dziugaite and Daniel Roy. Entropy-sgd optimizes the prior of pac-bayes bound: Generalization properties sgd and data-dependent priors. International Conference Machine Learning, pp. 1377-1386. PMLR, 2018.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. arXiv preprint arXiv: 1008, 2017.
Vitaly Feldman and Jan Vondrak. Generalization bounds for uniformly stable algorithms. arXiv preprint /xX2:/8//200985/ 2018.
10

---

Published as a conference paper at ICLR 2022
Vitaly Feldman and Jan Vondrak. High probability generalization bounds for uniformly stable algorithms with nearly optimal rate. In Conference on Learning Theory, pp. 1270-1279. PMLR, 2019.
Rong Ge, Jason D Lee, and Tengyu Ma. Matrix completion has spurious local minimum, arXiv 2016.
preprint
Rong Ge, Chi Jin, and Yi Zheng. No spurious local minima in nonconvex low rank problems: A unified geometric analysis. In International Conference on Machine Learning, pp. 1233-1242. PMLR, 2017.
Stuart Geman, Elie Bienenstock,, and René Doursat. Neural networks and the bias/variance dilemma. Neural computation, 4(1):1-58, 1992.
Daniel Gissin, Shai Shalev-Shwartz, and Amit Daniely. The implicit bias of depth: How incremental learning drives generalization. arXiv preprint arXiv:I909.12000 2019.
Suriya Gunasekar, Blake Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nathan Srebro. Implicit regularization matrix factorization. In 2018 Information Theory and Applications Workshop (ITA), pp. 1-10. IEEE, 2018.
Mahdi Haghifam, Jeffrey Negrea, Ashish Khisti, Daniel M Roy, and Gintare Karolina Dziugaite. Sharpened generalization bounds based on conditional mutual information and an application to noisy, iterative algorithms. arXiv preprint arXiv:2004. 2020.
Jeff Z HaoChen, Colin Wei, Jason Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. Conference on Learning Theory, pp. 2315-2357. PMLR, 2021.
Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In International Conference Machine Learning, pp. 1225-1234. PMLR, 2016.
Shi Hu, Nicola Pezzotti, Dimitrios Mavroeidis, and Max Welling. Simple and accurate uncertainty quantification from variance decomposition. arXiv preprint arXiv:2002.05582, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint arXiv: 006600572, 2018.
Prateek Jain, Sham Kakade, Rahul Kidambi, Praneeth Netrapalli, and Aaron Sidford. Parallelizing stochastic gradient descent for least squares regression: mini-batching, averaging, and model misspecification. Journal of Machine Learning Research, 18, 2018.
Kenji Kawaguchi, Leslie Pack Kaelbling, and Yoshua Bengio. Generalization in deep learning. arXiv preprint arXiv:1710.05445 2017.
Vladimir Koltchinskii. Rademacher penalties and structural risk minimization. IEEE Transactions on Information Theory, 47(5):1902-1914 2001.
Vladimir Koltchinskii and Dmitriy Panchenko. Rademacher processes and bounding the risk of function learning. In High dimensional probability 11, pp. 443-457. Springer, 2000.
Vladimir Koltchinskii et al. Local rademacher complexities and oracle inequalities in risk minimiza- tion. Annals of Statistics, 34(6):2593-266 2006.
Erich Lehmann and George Casella. Theory of point estimation. Springer Science Business Media, 2006.
Yunwen Lei and Yiming Ying. Fine-grained analysis of stability and generalization for stochastic gradient descent. International Conference on Machine Learning, pp. 5809-5819. PMLR, 2020.
Jian Li, Xuanyuan Luo, and Mingda Qiao. On generalization error bounds of noisy gradient methods for non-convex learning. arXiv preprint DrXiv:1902.00622) 2019.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang, Algorithmic regularization in pereparraeterizeed matrix sensing and neural networks with quadratic activations. In Conference On Learning Theory, pp. 2-47. PMLR, 2018.
11

---

Published as a conference paper at ICLR 2022
Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learning. MIT press, 2018.
Wenlong Mou, Liwei Wang, Xiyu Zhai, and Kai Zheng. Generalization bounds of sgld for non-convex learning: Two theoretical viewpoints. In Conference on Learning Theory, pp. 605-638. PMLR, 2018.
Vaishnavh Nagarajan and J Zico Kolter. Uniform convergence may be unable to explain generalization in deep learning. arXiv preprint arXiv:1902.04442, 2019.
Preetum Nakkiran, Behnam Neyshabur, and Hanie Sedghi. The deep bootstrap: Good online learners are good offline generalizers. arXiv e-prints, pp. arXiv-2010, 2020.
Jeffrey Negrea, Gintare Karolina Dziugaite, and Daniel Roy. In defense of uniform convergence: Generalization via derandomization with an application to interpolating predictors. In International Conference on Machine Learning, pp. 7263-7272 PMLR, 2020.
Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. In search of the real inductive bias: On the role of implicit regularization deep learning. arXiv preprint orXiv:1412 1412.6614, 2014.
Behnam Neyshabur, Srinadh Bhojanapalli, David McAllester, and Nathan Srebro. Exploring general- ization in deep learning. arXiv preprint arXiv:1766. 08947, 2017a.
Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-bayesian approach spectrally- normalized margin bounds for neural networks. arXiv preprint arXiv:1707.09566 2017b.
Samet Oymak, Zalan Fabian, Mingchen Li, and Mahdi Soltanolkotabi. Generalization guaran- tees for neural networks via harnessing the low-rank structure the jacobian. arXiv preprint 10f.:09900005992, 2019.
Benjamin Recht, Maryam Fazel, and Pablo Parrilo. Guaranteed minimum-rank solutions of linear matrix equations via nuclear norm minimization. SIAM review, 52(3):471-501, 2010.
Daniel Russo and James Zou. Controlling bias in adaptive data analysis using information theory. In Artificial Intelligence and Statistics, pp. 1232-1240. PMLR, 2016.
Nathan Srebro, Karthik Sridharan, and Ambuj Tewari. Optimistic rates for learning with a smooth loss. arXiv preprint 1009v:1009.5880 2010.
Thomas Steinke and Lydia Zakynthinou. Reasoning about generalization via conditional mutual information. In Conference on Learning Theory, pp. 3437-3452. PMLR, 2020.
Martin J Wainwright. Hlih--immnnsooall statistics: non-asymptotic viewpoint, volume 48. Cam- bridge University Press, 2019.
Colin Wei and Tengyu Ma. Improved sample complexities for deep networks and robust classification via an all-layer margin. arXiv preprint (XXiv:11100444 2019.
Aolin Xu and Maxim Raginsky. Informationn-hherrtiicc analysis of generalization capability of learning algorithms. arXiv preprint
2011 2017.
Jun Yang. Shengyang Sun, and Daniel Roy. Fast-rate pac-bayes generalization bounds via shifted rademacher processes. In NeurlPS, pp. 10802-10812, 2019.
Yaodong Yu, Zitong Yang. Edgar Dobriban, Jacob Steinhardt, and Yi Ma. Understanding generaliza- tion in adversarial training via the bias-variance decomposition. arXiv preprint (03.02003309947, 2021.
Chiyuan Zhang, Samy Bengio. Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv: 1.03530, 2016.
Richard Y Zhang, Somayeh Sojoudi, and Javad Lavaei. Sharp restricted isometry bounds for the inexistence of spurious local minima nonconvex matrix recovery. Journal of Machine Learning Research, 20(1 1-34, 2019.
12

---

Published as a conference paper at ICLR 2022
Yikai Zhang, Wenjia Zhang, Sammy Bald, Vamsi Pingali, Chao Chen, and Mayank Goswami. Stability sgd: Tightness analysis and improved bounds. arXiv preprint arXiv:2102.05274, 2021.
Wenda Zhou, Victor Veitch, Morgane Austern, Ryan P Adams, and Peter Orbanz. Non-vacuous generalization bounds the imagenet scale: pac-bayesian compression approach. arXiv preprint arXiv:/804.05862, 2018a.
Yi Zhou, Yingbin Liang, and Huishuai Zhang. Generalization error bounds with probabilistic guarantee for sgd in nonconvex optimization. arXiv preprint arXiv: 182.v0380,.06933, 2018b.
Jiacheng Zhuo, Jeongyeol Kwon, Nhat Ho, and Constantine Caramanis. On the computational and statistical complexity of over-parameterized matrix sensing. arXiv preprint arXiv:2102.02756, 2021.
Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, and Sham M Kakade. Benign overfitting of constant-stepsize sgd for linear regression, arXiv preprint arXiv:2103.1266 2021.
13

---

Published as a conference paper at ICLR 2022
Supplementary Materials
In this section, we first provide additional numerical experiment in Section and figure illustration in Section B We then give all the omitted proofs of lemmas and theorems in Section Section [ and Section We finally provide some further discussions, including why we apply stability on VER, the motivation behind DDC and relaxed version of DDC.
A NUMERICAL EXPERIMENT
In this section we provide additional numerical experiments on neural network to verify that our theory is satisfied for general neural network architectures and optimization algorithms.
3.5
1.0
9.6
33
standard
bies
0.0
standard
bias
1.0
standard
bias
5.5
upper bound
upper bound
veriance
0.0
upper bound
0.6
0.4
0-
54
02
0.0
0.0
0.0
2503
Epoch
Epoch
Epoch
Depth,
(b) Depth, T= 3
Depth,
0.5
3.4
1.0
0.7
standard
3.2
standard
standard
variance
varfance
variance
upper bound
0.8
upper bound
0.6
upper bound
0.6
0.2
0.2
0.2
0.0
0.0
0.0
2002
1200
2000
2500
3000
2002
2503
2002
3000
2000
1505
2000
2003
Epoch
Epoch 7
(poch
(d) Width, 64
(e) Width, 256
(f) Width, 512
0.5
3.0
1.0
standard
standard
veriance
variance
mac
upper
D.C
upper
standard blas
2-2
0.2
0.2
variance
03
upper hound
0.0
520
1003
1500
2200
7500
3000
20
40
(0)
300
100
20
40
60
80 100
Epoch
Epoch 7
Epoch 7
(g) SGD
(h) Adam
ii Rprop
Figure 4: Numerical results in Section A.1
A.1 SYNTHETIC EXPERIMENT
In this section, our goal is to recover sparse linear function f*(x) (0*, a)) (æ E R39) with fully connected ReLU network. For each setting, we run independent trials and report their average performances (the solid line) and the corresponding standard deviations (the error bar).
We first explore the effect of depth and width for the excess risk decomposition property. In this experiment, we fix the width each layer to be 64, and use SGD as the optimizer with Gaussian initialization (N(0, o²) where line, we plot aBERR + VER) to verify the performance of our risk decomposition framework. Here
Effects of depth:
1 10 3) and stepsize
x 10 For the dash
be regarded measure of the success of our theory. Ideally should be small constant that
14

---

Published as a conference paper at ICLR 2022
our bound is non-vacuous. Specific to this experiment, for Figure 4a respectively. Overall, these results are well-predicted by our theory. Moreover, is remarkable to notice that this constant becomes large once increasing the depth. This mainly contributes to the fact that deeper neural network has higher non-linearity.
4b
4c
we set a 1.3, 2.0, 2.3,
Effects of width: In this experiment, we explore the effect of width for 2-layer ReLU network with SGD optimizer (the initialization and the stepsize are the same as above). For Figure 4d we set 1.3, 1.0, 1.0, respectively. As can seen, the wider the neural network, the better our risk decomposition works. We comment that this behavior also comes to the level of non-linearity. It now well-known that when the width tends infinity, the neural network converges linear function (e.g. (Jacot et al. 2018)).
4c 4e
Effeets of optimizer: In this experiment, we try three different optimizers: SGD (the same hyperparameters before), Adam (stepsize 77 0.002, (B1, (22
..9,0.999), c le 08, no weight decay), Rprop (learning rate = 5 x 10 (m, 72) = (0.5, 1.2). stepsizes (1 x 10 50)). For Figure 4h 4i
set 1.3, 1.2, 1.1. It can be seen that our risk decomposition framework is robust to different optimization algorithms, demonstrating its generality and potential to extend to general optimization problems.
A.2 CIFAR-10
In this section we deliver additional experiment in CIFAR-10 dataset. In this experiment, we use ResNet-18 and choose Adam as the optimizer. We regard that optimizing over the original dataset as the bias training (in the sense that all the label is noiseless, referring to the bias part). As for the variance training, we the ground truth to be the uniform logit v = [0.1,-..
.0.]] (here we use one-hot coding and choose the categorical cross entropy loss). As for the labels in the training dataset, we generate it uniformly from (e,,- G10 ) where e; the standard basis in I 10 Under such data generating model, its Bayes risk is provided by -10 x 0.1 x 1og(0.1) x 2.3026 for categorical cross entropy loss. For the standard training, we generate the training label as follows: y = (1-p)y* pej-- Here [0, 1] is the corruption probability, y* is the true label and is uniformly sampled from {1..... 10). In this experiment, we set the corruption probability to be 0.4. The numerical result is provided in Figure Different from the performance on MINIST, here the variance training actually converges to the ground truth. We conjecture that this is due to the structure ResNet, resulting to some kind of implicit regularization towards the training process
B FIGURE ILLUSTRATION
B.1 FIGURE LA
Figure Ta conveys an important message that variance training has a much smaller gradient norm at the beginning compared standard training. This mainly comes from the difference in their landscapes. Since there is no signal term in the variance learning, we can regard the landscape corresponding the variance training as just a translation transformation of the original landscape (Figure Ta the blue line is just translation transformation of the green line). Based on this difference, they pursue different optima. We denote Ô.., θ the empirical loss minimizer, respectively. Due to concentration, we have 0 R 0, and Ô
O*. Note that we use near-zero initialization ( e 0, we immediately have that !! xe Ô.. indicating that the gradient norm is small. On the other hand, since g(0) is far away from the optimum, its initial gradient norm can be very large.
To better illustrate this principle, we use linear regression to provide simple example. For the standard training, the empirical loss is C. (0:P) Σ ((0, Xx) (0', X₂) ex) Hence, we
10We do observe that for certain learning rate regimes, the algorithm will overfit the training dataset for variance training, leading large VER (~ 10).
15

---

Published as a conference paper at ICLR 2022
2.5
standard
2.0
bias
empirical risk
variance
1.5
1.0
0.5
0.0
0
20
40
60
80
100
120
epoch
Figure 5: Three excess risk dynamics for CIFAR-10.
have
2 71
Vo,
Σ
(000), x:) - -0", (x) 61) Xi
1
i=1
(10)
-
2
(00 X e4) X₁ -2"
if we assume X; N (0, laxa). On the other hand, for the variance training, the empirical loss is C, - (i) - Therefore, have
(O. P.) = 2 Σ ,()) 1=1
X.) - 8.) X; Σ
ExX; R O.
(11)
1=1
In conclusion, we have
(o. x 2|
R0
(12)
Obviously, there is a huge initial gradient norm gap between standard and variance training, showing that our bound is significantly better than the direct stability-based bound under such regimes.
B.2 FIGURE
Figure [b] provides an example showing that the trajectory experienced by the variance training is likely stay convex region, though the global landscape can be highly non-smooth and non-convex.
The reason is similar to that of Figure la Since !!") x O.. the gradient norm is small and the training dynamics may stay close to the optimum, which lies smooth and convex region. On the contrary, the standard training needs to cross potentially many bad local minimum or saddle points, suffering from large cumulative gradient norms.
B.3 FIGURE 2
The setting of Figure 2 is similar to that of Section For MINIST dataset, we use three layer convoluted neural network with Adam optimizer.
16

---

Published as a conference paper at ICLR 2022
B.4 FIGURE
Figure provides the numerical verification of Dynamics Decomposition Condition, and our upper bounds for both linear regression and matrix recovery problems
Linear Regression We set the dimension d = 500, the sample sizes are n. 300, 800, respectively. Samples 31 In are i.id. from standard Gaussian distribution A (0, Ldxa), and the output = (z,, 0*) te where E; ~N(0,, o²) and o = 2. We set the initialization ( 0g² = 0 = 0 and the learning rate 0.01.
set the dimension d 20, rank 3, the ground truth X VEVT where Σ = 55 3, 1, ... .0} and V is orthonormal matrix. The numbers of the measurement matrices are 200, 600, respectively. The measurement matrices A.. An are i.idd standard Gaussian matrices. The corresponding measurement 14 (A, + where E; ~ (O, 1). In this numerical experiment, we use the initialization (()) = (6) = (0) = aTaxd where a = 0.01.
Matrix Recovery
We solve this problem via gradient descent with constant stepsize 7 0.1.
For the upper bound, it calculated via upper bound = X - 11 - + || + 8.3 where the last term is corresponding to the additional term in Lemma
Neural Network We the dimension d 30, the ground truth function is f*(x) (0 x) where or is sparse and | = =()). The covariate x is generated from Gaussian distribution. The label for bias training is exactly yi f (x:). For the variance trainng, its label is generated from additional Gaussian distribution &; ~ _f(0, o²) where a 1.5. For the standard training, its label is
y/ = f (xx) t&, For the upper bound, it is calculated via upper bound = fe: - f 1| + fl lp
Here Il is defined to be the L2(P) norm, i.e., l| Ez~p [f())]] For these settings, choose SGD (with stepsize 7 1 x 10-3) as the optimizer where the initialization is randomly generated from a Gaussian distribution.
C PROOF FOR OVERPARAMETERRIZE LINEAR REGRESSION
C.1 PROOF THEOREM
In this section, we prove Theorem stated as follows:
Theorem (Linear Regression Regimes) Under the overparameterized linear regression settings, 0-
assume that v -. bounded, the following inequality holds with probability at least 8:
is o -ubbaussiisa. If [|||| le| S V. suPre(T] ||0!01 < B are all
((() max (1,00, T' VV + BJ²)
log (4/0)
1012
TXV BJ
7
XI
n
where the probability is taken over the randomness of training data.
Proof:
Here, extend the diagonal matrix recovery problem the general matrix recovery problem.
17

---

Published as a conference paper at ICLR 2022
We first calculate the population loss under linear regression regime, namely, for any time
(( P)
Ez,y
y x ((()
(
0" + x 0(t)
E
E
x
00
a (())
2E
Ex T (())
+ +Ee
T
0
000
Σ
(1)
2I, E
cx T TT000/ |x
+ EE²
0.
00)
T
Σn
0
2
x O [e|ar]
+ Be²
(6) 0 ((+)
2 0
((0)
+Ee?
where (a) we use the condition that y T 0 c and (b) we use the assumption that Ejelæ] Therefore, the excess risk can be written
((((1) o O()) Ex (o- 6(0)
Now plugging the Lemma least 1 8:
Lemma
Lemma
into Theorem
2
we derive that with probability at
Ec Ec(0(T);)) S E (0T)
PP) + (0,T)
0 TT 0x
1
log
1011/
vn
A2T + 1]
T
V + B|² log(n) log(2n/0) + max 1, + B]]]]
log(4/0)
n
2n
max 11,00 Σn000 000 V + [V+}]]
log(4/8)
10-12
TX
IV
+ BJ
=
XT
which completes the proof.
C.2 PROOF OF LEMMA II
We first recall Lemma as follows:
Lemma Under the overparameterized linear regression settings with initialization ((0) =00
0 0. for any time T, the following decomposition inequality holds:
( P) S ( P) + (0)
(TT)) PP)
Proof: Firstly, when we use gradient descent, we have that:
A
((+1) 0(() +
x [y: - x ((),
A
0(+1) 0(0)) +
(5 x3 )(0))
n
000) 0(++1
A
Σ
[ 0- xx. 0(0)
x;
n
We next prove the conclusion by induction. Since choose (0) 0 =0(-) 0. the conclusion holds at the initialization
0 ( 0
18

---

Published as a conference paper at ICLR 2022
Now assume that ((() 0. + =gg holds time I, have
@(++1) = ((t) + Σ ly xx! 0
= ( + 00 + Σ
ai E; + x; 0 -x, (0.)) + +0(%))
0 + Σ x.(e; x; 6
+ 0 + Σ x [c, 00 x 0
((++1) + 0 ((t+1)
In summary, we have: Ō(t) 0 k0 holds for any time Besides, we easily calculate that
0- = 0 0g since 0" 0% = 0 and 0: 0.
Therefore, we have (it) Ar 0 0 + 0 0g holds for any By Schwwrz inequality, we derive that for specific time T:
( ; P)
0 (T))
Σx 0- ((T)
(
0: (
0 ((T)
Σ,
(0)
22 0: (T)
+ 0 0.:7
2 :P) + (0 P).
where we denote the norm ||u|3 AAu
C.3 PROOF OF LEMMA 2
Lemma 2 Under the overpurameterized linear regression settings, assume that |>|| S 1, |e| V, sup, |10." < B are all bounded. Consider the gradient descent algorithm with constant stepsize X,
with probability at least 1 1-6 6:
33 (P:) =0
T\ V + log(n) log(2n/6) + max {1, IV + BJ} V
log(2/0)
(2)
2n
Proof:
When the context is clear, we omit in this section and denote the trained parameter by 0(1)
For any time 1, we first split the excess risk into three components following Equation]] [
Ec(0(); (()()) : L( p.)
[()()) P.)) L( P.)) + [C(O": :()):
(III
(1)
(13)
Part (I). We bound the first component in Equation 13 using stability.
Fact A: the loss ((0: x, 3) VV + Lipsschtt
V 0; x. y) = 21|x(E - x T 0) < 2V + B].
Fact B: the stability of the parameter 2 VV + B). Denote the parameter trained on dataset S and the parameter trained on dataset S' 1 and 0 gr respectively. We will show that, when S and S
19

---

Published as a conference paper at ICLR 2022
have only one different sample (e.g,, WLOG S {z1, 22, ..... ".,n)) and 5' 121 ,.,n)).
@ is close to 0g)
0 0
5
( Xx
XTx
-
1
X X
n
11
n
X
0 0
||
|+||
xxTxx x'
x x 0 || + x - X
0
0
A
x TV
+ A
72
1
0
2\
V + B).
n
Summation over time / provides the following result:
| S 0 S [V B).
2t
n
Therefore, the total stability with respect to the f2-loss is
(((05)) ( - 41A
[V + B)
n
By Proposition we have that
C(010; - C(OW) S tA VV + B] log(n) (2n(22/)) +
log(1/8)
Part II. We next show that the second component in Equation 13 is less than zero based on Lemma 7
C(0(1);
1 ||e - x0 1/ - |||
e T 2X
I -- 1 A x
xt +Xt
I 1 1-2x x
X X
1
1 1- x-=XXxx
x
xt E.
n
n
n
The i-th eigenvalue of the matrix
-1- (xxxx]Jxx=> + x1 1.T -- -
[ - T "x]'Xi] can be calculated as follows
2 [1 - (1 (1-Ao;) + [1 - (1- = -11-=00 - 5,
where o; refers to the i-th eigenvalue of the matrix XT X. Therefore, have
(() P.)) -(0 P.) < O.
Part III. We bound the third component in Equation 13 using concentration.
Note that (0: x, 2) ( - T < [V + BJ² so we can use Hoeffding's inequality,
2nu²
P(I[((";P)) L( P.)l u] 2exp
(B + V)
By setting u = [V+ + 2 log(2/6) then with probability at least 1 8, we have
(2/()
IC(0:PP)) P()): S IV + B
2n
20

---

Published as a conference paper at ICLR 2022
Combining Part (I), (II), and (III) together, at time T, have:
TX
V + B|a log(n) log(2m/6) + log(1/5)
log(2/0)
Ec(O Po) S
+ IV B)²
n
2n
TX
log(2/0)
V + B|" log(n) log log(2n/6) max IV + B}²)}
2n
C.4 PROOF OF LEMMA 3
Lemma 3 Under the overparameteritzed linear regression settings, assume that
g- is Vo. 0
o² -ssuuaaasssian, then with probability at least 8:
-000) (0xTT2x0 log
;)) =0
+ 1(27 + 1 0
Vn o
Proof: Since we use the noiseless data (a, E(y(z)) and E(y/2>) = x 0 by Lemma 7 we have
000 - I
x x
x1xo = I I x x 1
x
0
n
n
where we use the fact that 1 - -
-
x
Therefore, BER can be expressed as:
Ec(0:P)) 0-
(1)
Σ 0- 0000)
A
1
A x T x
T
7
X
X
Σ 1
0"
X
A
xj'00 0+
T
T x
Σn
- XT X
-- n
1
XT x
n
n
X T X
1
+ T
X X
x
X
0¹
We split the excess risk into two parts. Intuitively, it is just like the decomposition in Equation I The first part is similar to the generalization gap to use uniform convergence to bound it. Note that here 1 Hence, have require that À supx, | T xo|l
21

---

Published as a conference paper at ICLR 2022
0
XTX
Dx n
- X
0"
n
n
sup 0.T X1
A Xo Xo
1 1
_ x x
1
1
A
Xo Xo
0"
n1
n
n
X Xo Xu
x X
sup Xa
r
Σ 12
x x
I
Xu
0*)
12
n
sup tr
Xo T Xo
2
1 X x
1
A
Xo
Xo
000
n
n
sup X1
I
X X
T X
Ex n
1 x x
Xu
T
Xu
0-0
n
n
sup Xa
[ n
X X T X
X x
I
A No
T Xo
0-0
n
sup X
Σ
X
x
Xo T Xo 0-
n
1
A Xo TTTxY X 7e
sup tr X
r
Σn
x X
1
0")
n
n
sup tr Xa
0
x
x
0-))
r
Σn,
- X
0"
n
where apply Lemma repeatedly since rank (0~T A'A 1 where A stands for arbitrary matrix.
We highlight the difference between Xg and X. The notation X represents the contribution of the training set to the training error. The notation Xo represents the contribution of the training set the training estimator θ. The uniform convergence is taken over the estimator O. Therefore, we take sup operator on Xo-
We next bound T Ex n x with subGaussian norm ow. Obviously, E [w] O,E [w3] 1. Therefore, we have that: w² is sub-Exponential distributions with sub-Exponential norm [2 11 S 0
x 0" Assume that w = 0 2/V//// 0r. Ee00" is SubGaussian
While e/ [0 0] o we have the following tail bound
1
P 10
Eg
- x
0
c
= P
[x]
(xT00² 0) -0 0E 2
n2
im1
1
P
[w] 1] 2 c/ _
n
i=1
[e/ [OT En 011
< eexp
no
where e is a universal constant.
Therefore, by setting e
1xT20
a log[2/0] -EVR
with probability at least 1 - 8, we have
1
T 0] 0 log 220
0..T
Σg
_ x
<
n
Veyn
22

---

Published as a conference paper at ICLR 2022
With abuse to use constant, we have with probability at least 2 1
x 0- c <c(0TE00] o (2g(2/)) Ex
1 x
0 T
vn
For the second part, which is closely related the empirical loss of 00 We will show that: for a general case, the empirical loss converges with rate O(1/t). Denote the eigen- values of _ X as o.. Tn. On.0,... .,,0,,,,,,00. ,0. Then the corresponding eigenvalues of matrix
[I A x x] [
[ A x] are (1 20 Therefore, we have:
0
x
x x n
[
1 X x
n
n
< $11011
m
\o; 2
10
X(2t 1)
The maximum is attained while a; X(2+1))
Combining these terms leads to the conclusion.
C.5 AUXILIARY LEMMAS
Lemma The dynamics of ((t) using GD is:
0(0) - (000--tty] + +xir.
where X represents the pseudo inverse of matrix X.
Proof:
The proof is based on induction. We notice that when = 0. the equation directly follows. Assume when 1 I, the equation holds. Then at time t + 1, we have
( 1) = (+) + A X (Y xx(1)
X x X
[0 -xto-xt]] xty
+ l A
X
X
H1
0 xtr] + -11-XTXX x xty + AxtY
T
1+1
XTX
X
0 x + +xly,
n
where we use the fact that X
xxt XT.
Lemma For rank-1 matrix A, we have that ||A|| |tr [A]]-
D PROOF OF THEOREM
In this section, provide the proof of Theorem 2 We first recall the formal statement of Theorem
Theorem (Decomposition Theorem) Assume that the excess risk has a unique minimizer with the following decomposition
local sharpness factor s. Under Assumption 7 and Assumption 2
inequality holds:
Eg (T())) S [4a]* Ma Ey (T)
(
+ (0 :P] + M
4C 1
M
7C
(5)
VT
Vn
23

---

Published as a conference paper at ICLR 2022
Proof: Under Assumption |
(upper bound), we have:
Ec(0(t;;)
;c())); 1 - 1- 10 0-11"
/ 6-1]
Combining Assumption 2 (DDC) and the lower bound in Assumption 1
we have
||0(0
0-1l
C
C
Sal00
( 01 + +all0 0.
+
Vt
Vn
1'
al 1/ (0.") Po) +
c
- 1
1
aE
(
P.)
Vi + Vn'
C C'
m
my
where we remark that Assumption holds uniformly on distribution pule and Therefore, we have that
(c(00):
P)
Mu 1/s
EE ( +
Mu 1/"
aE - + C'
aEl/s (
C
mn
m.
Vi
Vn
Mu a ) PP)
Mn.
M C
C²
4"
a (E
M//
mu
my
Vi
vn
$
[4a]*
Mu
[Ec(0] (1)) Po) ( P.)]
+ Ms
T47*
+ +Ma
4
my
Vt
Vn
where we use the fact that (a + b + c + d) < (4c)² + (4d) when a, b, c, e,d.p > 0.
(46)),(4c) (4c)P, (4d)² S (4a)² + (4b)² =
E PROOF FOR DIAGONAL MATRIX RECOVERY
In this section, we prove Theorem via proving the proceeding three lemmas separately.
E. PROOF OF THEOREM 3
Theorem (Decomposition Theorem in Diagonal Matrix Recovery) In the diagonal matrix re- covery settings, we set the initialization Uo = al where a is the initialization scale, then for any 0 < t < oo, with probability at least 1 - 8, we have:
Ee(U,;P) S (IX*1 +V + do²) ) log (d/8)
add
o3 + (0000
n
dV (t + log(n) log
2dn
1
rv log (r/6)
+ da² + log?
8
a
Note that although we do optimization based on U' we consider its square X = UU as the parameter since we require the solution is unique during the analysis. Note that we do not change the update rule (where GD is still taken on U). One can directly check that under Diagonal Matrix Recovery regimes, s 2 and M2 = m. 1. Therefore, combining Lemma Lemma and Lemma leads to Theorem 3
E.2 PROOF OF LEMMA
Lemma Under the assumptions Theorem 3 for any S < oc, with probability at least
we have
JUUT - X*11, S <||U:U0-- - -x-11
UTU. (
rv22 log (r/8)
+0
log
(7)
24

---

Published as a conference paper at ICLR 2022
Proof: Since all the measurement matrices A A are diagonal, analyze the dynamic of each singular value separately.
For given singular value a, assume at (),, Es subG(v²). Then define dj Di= af0, and Ø₂ = 1 à Di=1 a,E,, denote o = 01 + 02, and { = n Ei=1"; Denote u (t). (f),), (t) the dynamics of standard training, bias training, and variance training, respectively. Then our goal is show that lo u² (() S lo (f(()] (1) small term holds with high probability.
Before diving into the details, we first give the dynamics of u ((), (1)) o (1). If o > 0, we have
u (1) Eas + Ea2e2ot
01022200 20
(1) =
01 - + Ea2e2out
Furthermore, if 02 0, we have
02a2 2021
(t) 02 Ea² + 20222002
Otherwise, if 02 we have
02 a²
(t = (102) {a?) c² malt - -aa?
If a= 0, we have
a² (t())= {a +
Now we turn to the proof of Lemma If σ = 0, then σ = 02 which implies (²( (f). indicating that the conclusion holds.
Hence, we only consider the case that 2 > 0. By triangle inequality, it suffices to show
[t(()) - (t) O (lo2)).
Without loss of generality, we assume > 0 0. Note that
220
2010
|u²(t) - -2(())]
+ +((02))
2a² 01 + c22
a'oo, (c2at - 20
+ +0((a2))
(o + 201) (01 + (a)e201t)
Via integration-be -artoon-------ttt- formula, we have
20ze 20t ( c -202 ds & 202te 20t
20t 201t
Therefore, suffices to show that
00 2 20t (a + 2at) (01 + (a?62nt) Õ(1) O(1).
We prove that there exists a universal constant C, such that
2022 (a + (ar (01 + (a2e2011) S SC,
which is equivalent to show
<(t) = C (a + faze22t) 20t) (01 + 201) - a 2et 2 0, VD <t<< < xc.
Note that ((0) = C (a + (0?) (ar + Ea?) 2 0, and its gradient
= 2Cc + 000)220001)) + 2Comraa (62ot + 2201)) 2 20t 2020021120 20t
25

---

Published as a conference paper at ICLR 2022
By concentration of Chi-square distribution, have 1- 1og(1/0: S2(1 with probability at
least 1 - 6, providing that 2 log (1) It suffices to show that d'(t) 2 0, which is equivalent to
(x)) Ca Co 00st 0.
attains its minimum o + Co log (2dx2) at the point log (20 /20 Therefore, suffices to set C' log
Overall, we have
[u²(1) - n((1) o [og
_ loal)
i > 0;
[t²(t) - (t)( - (t)
if = 0.
Hence, with proper choice of probability, we have with probability at least 1 - 5, the following holds
ro² log (r/o)
- X S x| x + | +
log
V
a²
To summarize, the case is (a = 1, C C" = 0 (log (+t) / /ro² log (r/55)
-bounded (See
DDC).
E.3 PROOF OF LEMMA 5
Lemma For VER, under the assumptions in Theorem
3 with probability at least 6, we have
P)(U;; x dV 1) log(n) log
2dn
log (2d/5)
+ da² + dV
(8)
n
n
Proof:
Similar to the proof in overparameterrized linear regression regimes, we first split the excess risk in three componenets:
E.) = (x²(1) Ln (x²(1) + L. (x)(1) - Cn (0) + Ln (0) - (0)
(14)
(1)
(III)
Part (I). We bound the first component Equation using stability. We define d2
-i=1 0,Eis
0/2 = , Ei= a;e; + e','' Similarly, we define ç Ei= of and g' ===1
2 + d2
In the first case, we assume both 022 022 > so that their difference can be bounded
1 a²azo's (20xt - 2oge) a E'a2 Eage2osi + a (026' (
(02 Fa² + (a?c202t) (oz l/a2 + ('a2220051)
WLOG, assume 02 oy. Then, similar the proof in Lemma we have
a² a'o20' (220st 2002t)
O (l02 - dglt).
(o2 ta² (22e2ozt) (022 E²a² + ('a2e2a5t)
The second term is dominated by the first term. Hence, we only need to handle the third term. Note that
102l
E
- obg = 102 (g' - E(02 - ob21l =
a - -;²l
+
(anEn - d2
n
we have
a a (0zE' 2(
1ozL OGEl
(02 - 2 + 202) (022 - Ela2 + 2
E5'
02
2 0;"
+
E (a,on -,
26

---

Published as a conference paper at ICLR 2022
Combined the above three parts, finally have
v(++1))
(e((t) - 2t]
In the second case, we assume both 022. 02 < 0. Since both of them would decrease to 0, we have [t] u (t) I a² where O n174 )
In the last case, without loss of generality, we assume d2 d/2 Since u decreases and w (1) increases to ozz we have (x](1) (1) S oe S 102 - 02
Overall, we have
V( (v(++)) n
[t.(() - xt)
Besides, since the loss is O(V)-Lipshhitz:
[[a={(u)]| = 2 (an²² - 2) a S
where 2² is bounded upper bounded by OVV) during the training analysis due small initialization. Then the total stability is
(((u ((4 11
n
Summation over the dimension d, the whole loss (t( has stability
- Ea(U? (t)) s dV (t+
By Proposition we have that
C( P) - C(0());')) x (t++)
log(n) (g(dd//)) + v Vlog(1/ (og(1/5)
Vn
Part (II). We next calculate the second term. Take the explicit formula into the equation, we have
La (x²(1) - L. (0) 1 Σ
(c(u(()) 2a,))((()))
1=1
02 1 Σi 0r8; 0, we have
1 (t) - 20))
(02 - {aa) 202t 200t
(02 - Ea² + (a2e2001) S 0.
If f2 < 0, we have
0²a²
2/02|2a0
(df((()) - (1)) = ((lozl + 2look2 - Ea2) +
(102) + + llaalt - 2
v2
$ Eaa + 2/02 02 2 S
a
n
In conclusion, we have Cn (12(1) - L. (0)
Part (III). For the third term, just need single concentration inequality. Note that for the variance training, the loss function is of the order O (V2). Hence, via the Hoeffding inequality, with probability least 1 6, we have
log(1/8)
Cn (()) V
n
Combining the three terms leads to Lemma 5
27

---

Published as a conference paper at ICLR 2022
E.4 PROOF OF LEMMA 6
Lemma 6 For BER, under the assumptions in Theorem
3
with probability at least - ô, we have
(L, x (1xx1))++a))
+ da+)
/ log (d/o)
Σ
a²d
(9)
n
Proof:
Similar to linear regression, we decompose the excess risk into two parts:
B a² (uā(1)
o) E a²
1 n
af
(n)(t) o)) 2+ Σ at (v)
o)²
- -o)²
We use the uniform convergence to bound the first term. Assume that o > 0, we have
01a22200t
σ - uf(() = o
01 - Ea + fa
001 ςa²o + e (Co
Ea² + E21e
001 Ea²o
9 k + 2200
0² a²o
a² + 28ot
< (1 - a 2 S a.
Now assume 0, we have
a²
(1)())
t + < a
The second term refers to the training loss, if σ > we have
-²))
of
af (1)(1) o)²= =5 (a 1-22 + a2
o² + aao(
If a = 0, we have
(Ea2t + 1)² E a1
a1
of (12(1) -0²²= =
Finally, by assigning probability properly, we have
Ec(U,;PP) $ (1I*112> + dox)) V
log (d/8)
add
1=1 0 + a (0(
F ADDITIONAL DISCUSSION
In this section, we discuss more the details in the main text. We first show case where our bound outperforms previous stability-based bound under linear regimes. We then validate the rationality of applying stability-based bound by showing that the generalization gap (under variance regime) increases with time. We next provide some motivation behind Dynamics Decomposition Condition by showing that it indeed holds the beginning of the training phase. We finally introduce related version of Dynamics Decomposition Condition which does not require the information about 0*, 8 and 0
28

---

Published as a conference paper at ICLR 2022
F.1 COMPARISON PREVIOUS STABILITY-BASEE BOUNDS UNDER LINEAR REGIMES
One of the shortcomings of the stability- based bound is its failure with large time T. Our decompo- sition framework can aid reducing the failure at a large time 7. For example, when considering a the newly-propose decomposition bound
large time T n/3/4 with À
On On + B)
As a comparison, the original stability- based bound
Onn-1/4VV + BT?)).
We next show that B < B' holds with a large hignal-to-nolee ratio.
We write the above formula as
1 B B' 110(.)1
From the formulation of the trained parameter, we have 0 = - - x and A X X - [I x o Therefore, it suffices show that || I x > 11
(t)) [ - I
Note that the first part only related to signal e and the second part related to noise e (by applying concentration, the second part is closely related to the noise level). Therefore, informally, with large signal-to-noise ratio, the last equation ||!/ - [I A X X > 2 2|(0(()) holds. To summary, holds that B B' and our bound outperforms the original stability- -baed bounds.
F.2 THE GENERALIZATION GAP UNDER VARIANCE INCREASES WITH TIME
In this part, we will prove that the (variance regime) generalization gap indeed increases with time, which validates that applying stability-based bound (which also increases with time) is rational.
We provide the following Lemma which validates the above statement.
Lemma Let A. L() P) C(O p") be the generalization gap under variance regime. Assume that Eg = Then under linear regression settings as in Section we have
i
(t) 0.
0t
Proof: Note that the population loss is calculated as follows by Equation C.I
(( P)
(())
Σg
((() + Be?.
And the training loss is calculated as follows:
1 1 E X 11
L(0 pm)
1 1100
2 E
[
1 =1112
[
1
x
x
[0:0)
2 E
x0!
- -
n
Note that by Lemma
we have
T x 78
1" [xte].
(
1 -
29

---

Published as a conference paper at ICLR 2022
Therefore, we have:
(t) C(0 P) - ; P")
E [
+ IEe²
- 1 (000 -(000)
- -xTx X
1
(0:0
+
2 E x
_ ||1
n1
n
[01]
Eg
- X X
1
[000]
+
2
ET x
+ EEe²
1
ET|
n
72
n
ET [xt]]
1 - I A x x xx]
Σx - X x
--
- x X
[xte]
+ x x
1 A x x
[X + Ec² - -
By omitting the terms in A. (t) that is uncorrelated to t, we have (denote as A.(t))
(1) = E [xt]
A x x
Exp
1 x x
x x
x
XtE]
n
n
n
- -2E [xt]
Σg Σx- 1
- x x
I
X x X
7(
[
2 E. x n
x [
E [xt]]
-- A
X X
Ee n2
_ x
I
2
x 1
[xt=]
n
n
2E T [x]] T
[mn]]
-- A Ax x
[xtr]
= E [x]]]
- xx
Ex - x 2Σ,
I -- A x xxxxxx 1 [x(e].
n
where we use the fact that [xt]
XTX x-T==1 = X.
We next focus on the matrix
[
| - n
Ma a [X]]]]
x n
21
[
xT x n
xt.
where we plug in Eg = I by assumption. Denote the SVD of T as U Yu where the i-th eigenvalue is og. Then when o; = 0, the i - th eigenvalue of M is also equal to zero (due to the persudo inverse Xt). When o; > 0, the i th eigenvalue of M, can be derived as
n.
(1 [[1 Ao, σ.) 21 -Ao,)
Its derivation is then
θ (1; Ma)
2(1 \o.)² log(1 - Ao;) [(1 [1--a))(1 Ao,)' -1]
2 0,
(15)
0r
di
since log(1 - >xo,)= S and (1-o,)(1- \o,)² 1 0. Therefore, the derivation of each eigenvalue is larger than zero
We rewrite (t) = EMazT where TT UE is a vector independent of time and EMa is the diagonal matrix with diagonal of (Ma). As a result, we have
A (t) = zT 01
MMa)z T 2 0,
since 2 Esaa 2 0 according to Equation
15
30

---

Published as a conference paper at ICLR 2022
Note that (t) and (t) only differ with term independent of time, we have
0 (.(0))000000 5.(1) Ot A.(t) at
a
F.3 MOTIVATION BEHIND DYNAMICS DECOMPOSITION CONDITION
For a dataset = ((X,, y.) )=))= where yi = fo (X;) + E., and 0. We choose the 2z-loss to solve this problem 2()) 2n Di= (Jo(X.) y1)² Then we can show that the Dynamics Decomposition Condition holds at the beginning phase we initialize 0 = O, and assume fx( = 0, AX supp(P). For sufficient small t 0, the dynamics are well approximated by
@ (
(for (x,) v4) f.gro (xt))t:
000 000
Σ (fo00 - (xo) V. (Xi)t;
(16)
0
0
Σ
- (X.)
Vi (X:)1.
Suppose (00)
0 %00 0 000 0
then based on Taylor expansion, we have the approximation
Σ
V Soro (X,) (0!0) ) = 0)5) + +
(17)
Therefore, the Dynamics Decomposition Condition holds at least at the beginning phase. Moreover, we discuss the main text that for a branch of learning tasks, such as linear regression and matrix recovery, these Dynamies Decomposition Condition uniformly holds for arbitrary 0 I t co.
F.4 RELAXED VERSION OF DDC
One may notice that there exist 0-, O Of in DDC (See Assumption 2). which is the optimal solution the corresponding population loss. They are sometimes unavailable in practice and even have no explicit solution. For more convenient analysis, we propose relaxed version of DDC in Lemma I0
Lemma 10 (Relaxed DDC) Let the loss be l2 loss where Ely. 9) = (y - .) We assume the
predicting model få (x) ) is unique with respect to 1 (when 0: & f. fo. + fo,). Besides, assume fo=o(x) = 0 for all æ. For the ground truth, we assume an additive noise, namely, y = fo- (x) + where is independent of x. Then the following Equation 18 suffices to prove (max(a, 1}. C, boundedness for DDC (See Assumption 2).
C
C" Vn
|10(0 - - S a
(18)
V
Based Lemma [0] one can verify DDC without knowing the optimal solution parameter 0* O,, 05. Besides. since we usually have the iteration equation (the relationship between ((++1) and (()) given an algorithm, it is possible to verify DDC based on the iteration forms. Note that we may not directly verify the condition in practice based on the relaxed DDC if we do not have the explicit split on signal E ar) and noise
E [[yllr)) ar).
31


Reviews:
[{"review_id": "WiuwSyySQH", "reviewer": "Reviewer_HZ9d", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "This paper is written clearly and easy to follow. Under the linear regression setting, the decomposition makes a lot of sense because the dynamics can be split exactly into two parts. Compared with directly applying the stability arguments, the decomposition removes the reliance on the norm of the parameter $\\theta^*$.\n\nThe first concern is, how better is Theorem 1 compared to directly using uniform convergence to characterize the excess risk of linear regression, especially under the high signal-noise-ration scenario. It is argued that Theorem 1 is better than directly applying stability results under the high signal-noise-ration scenario, which makes sense. But I wonder if the proposed decomposition really improves the state-of-the-art risk bounds. Some added discussion may help.\n\nThe next concern is the applicability of this decomposition. For the decomposition to make sense in general regimes, some very strong assumptions are made. For example, Assumption 2 directly assumes the training dynamics can be decomposed into two parts, which I believe can hardly happen in most cases.\n\nI also have concerns with whether this decomposition is meaningful for modern deep learning tasks. In my opinion, this decomposition is meaningful typically under a setting where the 'bias training' is easy and the noise plays a nonnegligible role in the training dynamics. In modern deep learning tasks, usually, the 'bias training' task is the most complicated one, and to characterize this learning dynamic rather than the noise is of the most interest. Therefore, even if the decomposition is approximately correct for deep learning tasks, it could be such a case that how to characterize the 'bias training' dynamic is still unknown, while the 'variance training' dynamic is not important at all.\n\n", "summary_of_the_paper": "This paper proposes a new approach to characterize the excess risk by decomposing the risk into two parts, the risk of learning the exact response and the risk of fitting the noise signal. Under the linear regression setting, the learning dynamic can be decomposed exactly into the proposed two parts. Stability and uniform convergence arguments are applied to two parts respectively. Similar results are shown under the more general setting, though additional assumptions are made to guarantee a reasonable dynamic decomposition. ", "summary_of_the_review": "The decomposition proposed in this paper makes sense in some specific settings like linear regression. However, it is very unclear for me to see how this decomposition can be extended to a more general regime and how it helps understand the learning dynamics of deep learning. There are also no new technical tools developed. Given that I believe the decomposition's applicability and significance are questionable, and the technical contribution seems incremental, I recommend a weak reject.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "T-rZYrTL3sA", "reviewer": "Reviewer_MkVE", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "As a non-specialist in this topic, I found the paper hard to follow. While the idea of decomposing the excess risk and combining the best of stability and uniform techniques to produce a better bound is very interesting and clear, the writing is confusing. Moreover, the quantity of typos and confusing sentences (some of which I list below) don't help.\n\n**Questions**\n\n1. The authors should explicit better what is the setting considered in the non-linear regime. While in the linear regime it is very clear what is the model of data and the loss, there is a conceptual jump between Sec. 3 and Sec. 4 which makes it confusing to understand how general are the results in Sec 4. Does theorem 2 holds for any model and any data distribution as long as assumptions 1 and 2 are valid? If yes, could the authors provide an intuition of how much these assumptions are constraining?\n\n2. For the real data experiments, how is the noise component estimated?\n\n**Comments / Suggestions**\n\n- In the second paragraph in the introduction, both shortcomings of stability-based bounds refereed by the authors are due to non-convexity. The first is due to the norm of the gradient, but it is not clear what the second is due to. The authors should precise better to what aspect of non-convexity the second obstruction comes from.\n\n- On the related work, \"Bias-variance decomposition\": [1] proceeded [2] in the investigation of the bias-variance decomposition for the double descent in random features regression.\n\n- *Besides, recent works (e.g., Jacot et al. (2018)) show that neural networks converge to overparametrized linear models (with respect to the parameters) as the width goes to infinity under some regularity conditions*. The authors could be more precise on what they mean by *regularity* here. As it is, the sentence seems to suggest that almost all infinite-width network is lazy.\n\n- The notation of vectors should be consistent. As it is, vectors with latin letters are denoted with bold characters, while vectors with greek letters (such as the weights $\\theta$) are not bold. \n\n- The authors speak about neural networks in the abstract, and present two plots in Fig. 2 and Fig. 3 related to neural networks. However, these networks are never defined or refereed to in the main manuscript. The authors should either detail the setting they are plotting in the main, or leave these results to the appendix.\n\n**Minor typos**\n\n- Abstract: *As alternative approaches, techniques based on stability analyze the training dynamics and drive algorithm-dependent generalization bounds.*. This sentence is confusing.\n\n- Abstract: *[...] (a) stability-based bound*. \n\n- Introduction, first paragraph: *which takes (the) supremum*. \n\n- Introduction, second paragraph: *Stability-based bound(s)* or *(The) Stability-based bound*. Appears three times in this paragraph.\n\n- Introduction, footnote: *In this paper, we refer the signal to the clean data without the output noise, and the noise to the output noise.* This sentence is confusing. Maybe: \"we refer to the clean data without the output noise as the signal\" would be clearer. \n\n- Introduction, third paragraph: *component(s)*\n\n- Introduction, third paragraph: *on both synthetic and real-world dataset(s)* or *on both (a) synthetic dataset and (a) real-world dataset* (the first option is better).\n\n- Introduction, fourth paragraph: *via (a) bias-variance decomposition*.\n\n- Introduction, contributions: *into (a) variance and bias component*.\n\n- Introduction, contributions: *provides interesting insights into the generalization community.*. into -> for?\n\n- Introduction, related work: *Oymak et al. (2019) applied bias-variance decomposition on (the) Jacobian of neural networks to explain their different performances on clean and noisy data.*\n\n- Section 2, third paragraph: *Let $\\mathcal{A}_{t}$ denote the optimization algorithm which takes (the) dataset $\\mathcal{D}$ as (an) input and return(s) the trained parameter $\\hat{\\theta}^{\u0012t}$ at time $t$*\n\n- Section 2, fourth paragraph: *Although the minimizer (might not) be unique*\n\n- Section 2, fifth paragraph: spilt -> split\n\n- Section 3, page 5: *The notation $B$ in Theorem 1 denote(s)*\n\n- Section 4, page 6: *the only way to reach the* or *the only way of reaching the*\n\n- Section 4, page 6: *minimizer $\\theta^{\\star}$ of (the) excess risk*\n\n[1] Stéphane D'Ascoli, Maria Refinetti, Giulio Biroli, Florent Krzakala, Double Trouble in Double Descent: Bias and Variance(s) in the Lazy Regime, arXiv:2003.01054 [cs.LG]\n\n[2] Ben Adlam, Jeffrey Pennington, Understanding Double Descent Requires a Fine-Grained Bias-Variance Decomposition, arXiv:2011.03321 [stat.ML]\n", "summary_of_the_paper": "This work proposes a new bound for the excess risk based on a decomposition into a bias and variance term. The authors use a stability-based analysis to bound the bias excess risk, and uniform convergence to bound the variance excess risk. They illustrate their framework in two contexts: overparametrised linear regression with SGD and low-rank diagonal matrix recovery with gradient flow. ", "summary_of_the_review": "The work has some interesting new ideas that, to the best of my knowledge, are original. However, the results are presented in a confusing manner, making it hard to understand the scope of applicability of these results to a non-expert. Therefore, I believe this work could benefit from a rewriting. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "LFGSatpoX73", "reviewer": "Reviewer_PEVU", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "The idea of decomposing excess risk dynamics seems to be novel and interesting. However, there are some issues.\n\n- The advantage of the results is not quite clear. For example, the authors indicated that Thm 1 outperforms the existing stability analysis by replacing $B'$ with $B$. However, this improvement is minor as both two quantities are unknown to us and it is not clear which one would be significantly larger. While the authors mentioned that $B$ is independent of $\\|\\theta^* \\|$. This seems not important since the excess risk bound already involves $\\|\\theta^* \\|$.\n\n- The decomposition theorem requires Assumption 1 and Assumption 2, which may be restrictive. While the authors mentioned that they hold for some specific cases, in general cases these conditions are difficult to check in practice and may not hold.\n\n- Assumption 1 on the uniform bound is imposed for $\\theta$ and $P$. However, in the proof of Thm 2, the authors also useed the uniform bound assumption for $(\\theta^{(t)}_b,P_b)$ and $(\\theta^{(t)}_v,P_v)$. The authors should indicate these assumptions clearly in Assumption 1.\n\n- The stability approach has a drawback of not yielding good bounds for nonconvex problems, while the uniform convergence approach has a drawback of yielding dimension-dependent bounds. The authors proposed to use stability approach and uniform convergence approach to tackle the noise and signal separately. The theoretical results therefore may inherit the drawbacks of both stability and uniform convergence: not appealing for nonconvex problems and yielding dimension-dependent bounds.\n\n- The analysis seems not rigorous. In particular, the authors proved in page 19 that $\\ell(\\theta^{(t)},P_v^n)\\leq \\ell(\\theta^*,P_v^n)$. This is not quite intuitive since this result holds for all $t$. In particular, in the first few iterations the performance of gradient descent may not be good and it is unlikely that they have smaller training errors than $\\theta^*$. Furthermore, the notation $\\epsilon$ in page 19 is not clear. It should be a vector according to the norm there. This causes confusion since $\\epsilon$ is a real-valued random variable in Thm 1.\n\nTypos:\n- page 17: $E[\\epsilon x^\\top\\theta^{(t)}]$ should be $-2E[\\epsilon x^\\top\\theta^{(t)}]$\n- page 18: $x_i^\\top\\theta_v+\\theta_b$ should be $x_i^\\top(\\theta_v+\\theta_b)$\n- proof of Lemma 7: some $\\theta$ should be $\\hat{\\theta}$. $\\theta_0$ should be $\\theta^{(0)}$\n- page 23: $M_u^{1/s}/m_u$ should be $(M_u/m_u)^{\\frac{1}{s}}$. There is also a missing factor of $m_u^{1/s}$ ", "summary_of_the_paper": "This paper studies the generalization performance of learning algorithms. The basic idea is to decompose the training dynamics into noise and signal components, and then consider separately the behavior of models trained with noise and signals. The authors then proposed to use the stability approach to tackle the noise components, and the uniform convergence approach to tackle the signal part. The authors consider two specific algorithms: overparameterized linear regression and diagonal matrix recovery. Experimental results are also reported.", "summary_of_the_review": "The idea seems to be novel. However, the advantage of the decomposing excess risk dynamics is not well justified. There are also some issues on the correctness of the deduction.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "D0T_yodhTQY", "reviewer": "Reviewer_Pc9J", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "I believe decomposing the excess risk into the variance excess risk and bias excess risk intuitively makes sense, which aligns well with the empirical observation that neural networks converge fast when fitting signal but converge relatively slowly when fitting noise. Especially, the authors first pointed out two flaws in stability-based bounds, i.e., \n\n- stability-based bound depends heavily on the gradient norm\n- stability-based bound usually does not work well under general non-convex regimes.\n\nThese two flaws make stability-based bound cannot handle well with neural networks, where the gradient norm at initialization is usually large, and the optimization path is highly-nonconvex. By decomposing the excess risk into BER and VER, the authors can hopefully get a better bound. For VER, they apply the stability-based bound on it, as the gradient norm of variance training is usually smaller than that of standard training at initialization and the optimization of VER is closer to convex. Whereas for BER, they adopt the uniform convergence to bound it. \nIn addition, the theorems cover both the linear regimes and the non-linear regimes as well as accompanied with some empirical verifications. Even though, I have the following concerns:\n\n- Both the linear regime and diagonal matrix recovery regime are too different from the neural network case, as the optimization of neural networks is highly non-convex. Although the NTK theory indicates that the neural network at the infinite-width limit is equivalent to kernel methods, it is still far away from explaining the generalization/optimization property of practical neural networks. Therefore, I believe the claims in the paper about neural networks are not well-supported. I would suggest the authors apply their framework to analyze, for example, two-layer relu networks.\n\n- It is unclear to me how does the new decomposition lead to improvements over the original stability bound quantitatively, or the improvement seems quite marginal?  For example, on page 6, you mentioned that \"the bound in Theorem 1 outperforms the stability-based bound when $\\|θ^\\star\\|^2$ is\nlarge compared to V (namely, large signal-to-noise ratio)\". However, there is also a $\\|θ^\\star\\|^2$ in Theorem 1, and both bounds are approximately $\\tilde{\\mathcal{O}}(1/\\sqrt{n})$.  I believe a more concrete comparison between your bound and the original stability bound is necessary (e.g., a bound depends on the signal-to-noise ratio?). Otherwise, I am not convinced that your decomposition will indeed improve the original bound.\n\n\n- The following paper is very relevant to your work, it would be great if you can discuss it in the related work.\n-- The Deep Bootstrap: Good Online Learners are Good Offline Generalizers, Preetum Nakkiran, Behnam Neyshabur, Hanie Sedghi International Conference on Learning Representations (ICLR), 2021.", "summary_of_the_paper": "In this paper, the authors propose to decompose the excess risk into the variance excess risk (VER) and bias excess risk (BER). Based on this decomposition, the authors derive the generalization bound for overparameterized linear regimes and matrix recovery regimes. ", "summary_of_the_review": "Overall, the idea in the paper is well-motivated. However, the results in the paper are mostly about overparameterized linear regression and matrix recovery, but the many of the claims in the paper are about understanding the generalization of neural networks. Besides, I am not convinced that the derived bounds in the paper really improve over the original bounds. Therefore, I am currently leaning towards a weak rejection. If the authors can address my concerns, I will consider increasing my rating. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The main contribution is a way of analyzing the generalization error of neural nets by breaking it down into bias and variance components, and using separate principles to analyze each of the two components. The submission first proves rigorous generalization bounds for overparameterized linear regression (motivated in a general sense by the NTK); there are settings where this improves upon existing bounds. It extends the case to a matrix recovery model, showing that it's not limited to the linear regime. Finally, experimental results show that the risk decomposition holds empirically for neural nets.

The numerical scores would place this paper slightly below the cutoff. The reviewers feel that the paper is well written and have not identified anything that looks like a critical flaw. They have a variety of concerns, mostly centered around whether the results apply to practical situations. Specifically, they're worried about (1) the theory not applying directly to neural nets, (2) the high-noise setting being less relevant for modern deep learning, and (3) whether there's a realistic situation where it improves over past bounds. Regarding (1), the theory covers not only the linear regime, but also the nonlinear matrix recovery regime; combined with the empirical results, this seems pretty solid by the standards of a DL theory paper. Regarding (2), even though the most common benchmarks indeed have low label noise, the high-noise regime still seems worth understanding (after all, we'd like our nets to work in domains like medicine). I haven't dug deeply enough to properly evaluate (3), but the author response seems believable to me.

Overall, the paper strikes me as creative and well-executed. Regardless of whether the theory is easily extendable to neural nets, this seems like an interesting paper that can be built on in future work. I recommend acceptance.
Decision: Accept (Poster)