SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhenmei Shi*, Junyi Wei*, Yingyu Liang
zhmeishi@cs.wisc.edu, jwei53@wisc.edu, yliang@cs.wisc.edu
Neural networks have achieved remarkable empirical performance, while the current theoretical analysis is not adequate for understanding their success, e.g., the Neural Tangent Kernel approach fails to capture their key feature learning ability, while recent analyses on feature learning are typically problem-specific. This work proposes a unified analysis framework for two-layer networks trained by gradient descent. The framework is centered around the principle of feature learning from gradients, and its effectiveness is demonstrated by applications in several prototypical problems such as mixtures of Gaussians and parity functions. The framework also sheds light on interesting network learning phenomena such as feature learning beyond kernels and the lottery ticket hypothesis.
Neural network (NN) learning has achieved remarkable empirical success and has been a main driving force for the recent progress in machine learning and artificial intelligence. On the other hand, theoretical understandings significantly lag behind. Traditional analysis approaches are not adequate due to the overparameterization of practical networks and the non-convex optimization in the training via gradient descent. One line of work (e.g. [9, 31, 38, 60, 71, 123] and many others) shows under proper conditions, heavily overparameterized networks are approximately linear models over data-independent features, i.e., a linear function on the Neural Tangent Kernel (NTK). While making weak assumptions about the data and thus applicable to various settings, this approach requires the network learning to be approximately using fixed data-independent features (i.e., the kernel regime, or fixed feature methods). It thus fails to capture the feature learning ability of networks (i.e., to learn a feature mapping for the inputs which allow accurate prediction), which is widely believed to be the key factor to their empirical success in many applications (e.g., [54, 77, 117, 119]). To study feature learning in networks, a recent line of work (e.g. [5, 6, 14, 33, 52, 72, 76, 116] and others) shows examples where networks provably enjoy advantages over fixed feature methods (including NTK), under different settings and assumptions. While providing more insights, these studies typically focus on specific problems, and their analyses exploit the specific properties of the problems and appear to be unrelated to each other. Is there a common principle for feature learning in networks via gradient descent? Is there a unified analysis framework that can clarify the principle and also lead to provable error guarantees for prototypical problem settings?
In this work, we take a step toward this goal by proposing a gradient feature learning framework for analyzing two-layer network learning by gradient descent. (1) The framework makes essentially no assumption about the data distribution and can be applied to various problems. Furthermore, it is centered around features from gradients, clearly illustrating how gradient descent leads to feature learning in networks and subsequently accurate predictions. (2) It leads to error guarantees competitive with the optimal in a family of networks that use the features induced by gradients on the
data distribution. Then for a specific problem with structured data distributions, if the optimal in the induced family is small, the framework gives a small error guarantee.
We then apply the framework to several prototypical problems: mixtures of Gaussians, parity functions, linear data, and multiple-index models. These have been used for studying network learning (in particular, for the feature learning ability), but with different and seemingly unrelated analyses. In contrast, straightforward applications of our framework give small error guarantees, where the main effort is to compute the optimal in the induced family. Furthermore, in some cases, such as parities, we can handle more general data distributions than in the existing work.
Finally, we also demonstrate that the framework sheds light on several interesting network learning phenomena or implications such as feature learning beyond the kernel regime, lottery ticket hypothesis (LTH), simplicity bias, learning over different data distributions, and new perspectives about roadmaps forward. Due to space limitations, we present implications about features beyond the kernel regime and LTH in the main body but defer the other implications in Appendix C with a brief here. (1) For simplicity bias, it is generally believed that the optimization has some implicit regularization effect that restricts learning dynamics to a low capacity subset of the whole hypothesis class, so can lead to good generalization [53, 90]. Our framework provides an explanation that the learning first learns simpler functions and then more sophisticated ones. (2) For learning over different data distributions, we provide data-dependent non-vacuous guarantees, as our framework can be viewed as using the optimal gradient-induced NN to measure or quantify the "complexity" of the problem. For easier problems, this quantity is smaller, and our framework can give a better error bound to derive guarantees. (3) For new perspectives about roadmaps forward, our framework suggests the strong representation power of NN is actually the key to successful learning, while traditional ones suggest strong representation power leads to vacuous generalization bounds [19, 33]. Thus, we suggest a different analysis road. Traditional analysis typically first reasons about the optimal based on the whole function class then analyzes how NN learns proper features and reaches the optimal. In contrast, our framework defines feature family first, and then reasons about the optimal based on it.
Neural Networks Learning Analysis. Recently there has been an increasing interest in the analysis of network learning. One line of work connects the sufficiently over-parameterized neural network to linear methods around its initialization like NTK (e.g. [9, 11, 20, 21, 31, 38, 49, 60, 62, 69, 71, 78, 82, 91, 93, 95, 114, 121, 122] and more), so that the neural network training is a convex problem. The key idea is that it suffices to consider the first-order Tyler expansion of the neural network around the origin when the initialization is large enough. However, NTK lies in the lazy training (kernel) regime that excludes feature learning [29, 50, 68, 113]. Many studies (e.g. [2, 5, 6, 8, 12, 14, 22, 26, 33, 37, 51, 52, 57, 58, 70, 72, 73, 76, 99, 112, 115, 116] and more) show that neural networks take advantage over NTK empirically and theoretically. Another line of work is the mean-field (MF) analysis of neural networks (e.g. [27, 28, 36, 79, 80, 100, 106] and more). The insight is to see the training dynamics of a sufficiently large-width neural network as a PDE. It uses a smaller initialization than the NTK so that the parameters may move away from the initialization. However, the MF does not provide explicit convergence rates and requires an unrealistically large width of the neural network. One more line of work is neural networks max-margin analysis (e.g. [30, 47, 48, 56, 61, 63, 74, 75, 83, 85, 86, 107, 109] and more). They need a strong assumption that the convergence starts from weights having perfect training accuracy, while feature learning happens in the early stage of training. To explain the success of neural networks beyond the limitation mentioned above, some work introduces the low intrinsic dimension of data distributions [17, 18, 23, 24, 25, 44, 67, 104, 108, 124]. Another recent line of work is that a trained network can exactly recover the ground truth or optimal solution or teacher network [3, 4, 10, 39, 84, 87, 94, 96, 120], but they have strong assumptions on data distribution or model structure, e.g., Gaussian marginals. [1, 40, 55, 110, 111] show that training dynamics of neural networks have multiple phases, e.g., feature learning at the beginning, and then dynamics in convex optimization which requires proxy convexity [43] or PL condition [65] or special data structure.
Feature Learning Based on Gradient Analysis. A recent line of work is studying how features emerge from the gradient. [7, 46] consider linear separable data and show that the first few gradient steps can learn good features, and the later steps learn a good network on neurons with these features. [33, 45, 105] have similar conclusions on non-linear data (e.g., parity functions), while in their problems one feature is sufficient for accurate prediction (i.e., single-index data model).
[32] considers multiple-index with low-degree polynomials as labeling functions and shows that a one-step gradient update can learn multiple features that lead to accurate prediction. [13, 81] studies one gradient step feature improvements at different learning rates. [97] proposes Recursive Feature Machines to show the mechanism of recursively feature learning but without giving a final loss guarantee. These studies consider specific problems and exploit properties of the data to analyze the gradient delicately, while our work provides a general framework applicable to different problems.
3 Gradient Feature Learning Framework
Problem Setup. We denote  $[n] := \{1, 2, \dots, n\}$  and  $\tilde{O}(\cdot), \tilde{\Theta}(\cdot), \tilde{\Omega}(\cdot)$  to omit the log term inside. Let  $\mathcal{X} \subseteq \mathbb{R}^d$  denote the input space,  $\mathcal{Y} \subseteq \mathbb{R}$  the label space. Let  $\mathcal{D}$  be an arbitrary data distribution over  $\mathcal{X} \times \mathcal{Y}$ . Denote the class of two-layer networks with  $m$  neurons as:
$$
\mathcal {F} _ {d, m} := \left\{f _ {\left(\mathbf {a}, \mathbf {W}, \mathbf {b}\right)} \mid f _ {\left(\mathbf {a}, \mathbf {W}, \mathbf {b}\right)} (\mathbf {x}) := \mathbf {a} ^ {\top} \left[ \sigma \left(\mathbf {W} ^ {\top} \mathbf {x} - \mathbf {b}\right) \right] = \sum_ {i \in [ m ]} \mathbf {a} _ {i} \left[ \sigma \left(\left\langle \mathbf {w} _ {i}, \mathbf {x} \right\rangle - \mathbf {b} _ {i}\right) \right] \right\}, \tag {1}
$$
where  $\sigma(z) = \max(z,0)$  is the ReLU activation function,  $\mathbf{a} \in \mathbb{R}^m$  is the second layer weight,  $\mathbf{W} \in \mathbb{R}^{d \times m}$  is the first layer weight,  $\mathbf{w}_i$  is the  $i$ -th column of  $\mathbf{W}$  (i.e., the weight for the  $i$ -th neuron), and  $\mathbf{b} \in \mathbb{R}^m$  is the bias for the neurons. For technical simplicity, we only train  $\mathbf{a}$ ,  $\mathbf{W}$  but not  $\mathbf{b}$ . Let superscript  $(t)$  denote the time step, e.g.,  $f_{(\mathbf{a}^{(t)},\mathbf{W}^{(t)},\mathbf{b})}$  denote the network at time step  $t$ . Denote  $\Xi := (\mathbf{a},\mathbf{W},\mathbf{b}), \Xi^{(t)} := (\mathbf{a}^{(t)},\mathbf{W}^{(t)},\mathbf{b})$ . The goal of neural network learning is to minimize the expected risk, i.e.,  $\mathcal{L}_{\mathcal{D}}(f) := \mathbb{E}_{(\mathbf{x},y) \sim \mathcal{D}} \mathcal{L}_{(\mathbf{x},y)}(f)$ , where  $\mathcal{L}_{(\mathbf{x},y)}(f) = \ell(yf(\mathbf{x}))$  is the loss on an example  $(\mathbf{x},y)$  for some loss function  $\ell(\cdot)$ , e.g., the hinge loss  $\ell(z) = \max\{0,1-z\}$ , and the logistic loss  $\ell(z) = \log[1 + \exp(-z)]$ . We also consider  $\ell_2$  regularization. The regularized loss with regularization coefficient  $\lambda$  is  $\mathcal{L}_{\mathcal{D}}^{\lambda}(f) := \mathcal{L}_{\mathcal{D}}(f) + \frac{\lambda}{2} (\|\mathbf{W}\|_F^2 + \| \mathbf{a}\|_2^2)$ . Given a training set with  $n$  i.i.d. samples  $\mathcal{Z} = \{(x^{(l)},y^{(l)})\}_{l \in [n]}$  from  $\mathcal{D}$ , the empirical risk and its regularized version are:
$$
\widetilde {\mathcal {L}} _ {\mathcal {Z}} (f): = \frac {1}{n} \sum_ {l \in [ n ]} \mathcal {L} _ {\left(\mathbf {x} ^ {(l)}, y ^ {(l)}\right)} (f), \quad \widetilde {\mathcal {L}} _ {\mathcal {Z}} ^ {\lambda} (f): = \widetilde {\mathcal {L}} _ {\mathcal {Z}} (f) + \frac {\lambda}{2} \left(\| \mathbf {W} \| _ {F} ^ {2} + \| \mathbf {a} \| _ {2} ^ {2}\right). \tag {2}
$$
Then the training process is summarized in Algorithm 1.
Algorithm 1 Network Training via Gradient Descent
Initialize  $(\mathbf{a}^{(0)},\mathbf{W}^{(0)},\mathbf{b})$
Sample  $\mathcal{Z}^{(t - 1)}\sim \mathcal{D}^n$
$$
\mathbf {a} ^ {(t)} = \mathbf {a} ^ {(t - 1)} - \eta^ {(t)} \nabla_ {\mathbf {a}} \widetilde {\mathcal {L}} _ {\mathcal {Z} ^ {(t - 1)}} ^ {\lambda^ {(t)}} \left(f _ {\Xi^ {(t - 1)}}\right), \quad \mathbf {W} ^ {(t)} = \mathbf {W} ^ {(t - 1)} - \eta^ {(t)} \nabla_ {\mathbf {W}} \widetilde {\mathcal {L}} _ {\mathcal {Z} ^ {(t - 1)}} ^ {\lambda^ {(t)}} \left(f _ {\Xi^ {(t - 1)}}\right)
$$
In the whole paper, we need some natural assumptions about the data and the loss.
Assumption 3.1. We assume  $\mathbb{E}[\| \mathbf{x}\| _2]\leq B_{x1}$ ,  $\mathbb{E}[\| \mathbf{x}\| _2^2 ]\leq B_{x2}$ ,  $\| \mathbf{x}\| _2\leq B_x$  and for any label  $y$ , we have  $|y|\leq 1$ . We assume the loss function  $\ell (\cdot)$  is a 1-Lipschitz convex decreasing function, normalized  $\ell (0) = 1$ ,  $|\ell^{\prime}(0)| = \Theta (1)$ , and  $\ell (\infty) = 0$ .
Remark 3.2. The above are natural assumptions. Most input distributions have the bounded norms required, and the typical binary classification  $\mathcal{Y} = \{\pm 1\}$  satisfies the requirement. Also, the most popular loss functions satisfy the assumption, e.g., the hinge loss and logistic loss.
3.1 Warm Up: A Simple Setting with Frozen First Layer
To illustrate some high-level intuition, we first consider a simple setting where the first layer is frozen after one gradient update, i.e., no updates to  $\mathbf{W}$  for  $t\geq 2$  in Algorithm 1.
The first idea of our framework is to provide guarantees compared to the optimal in a family of networks. Here let us consider networks with specific weights for the first layer:
Definition 3.3. For some fixed  $\mathbf{W} \in \mathbb{R}^{d \times m}$ ,  $\mathbf{b} \in \mathbb{R}^d$ , and a parameter  $B_{a2}$ , consider the following family of networks  $\mathcal{F}_{\mathbf{W},\mathbf{b},B_{a2}}$ , and the optimal approximation network loss in this family:
$$
\mathcal {F} _ {\mathbf {W}, \mathbf {b}, B _ {a 2}} := \left\{f _ {(\mathbf {a}, \mathbf {W}, \mathbf {b})} \in \mathcal {F} _ {d, m} \mid \| \mathbf {a} \| _ {2} \leq B _ {a 2} \right\}, \quad \operatorname {O P T} _ {\mathbf {W}, \mathbf {b}, B _ {a 2}} := \min  _ {f \in \mathcal {F} _ {\mathbf {W}, \mathbf {b}, B _ {a 2}}} \mathcal {L} _ {\mathcal {D}} (f). \tag {3}
$$
The second idea is to compare to networks using features from gradient descent. As an illustrative example, we now provide guarantees compared to networks with first layer weights  $\mathbf{W}^{(1)}$  (i.e., the weights after the first gradient step):
Theorem 3.4 (Simple Setting). Assume  $\widetilde{\mathcal{L}}_{\mathcal{Z}}\left(f_{(\mathbf{a},\mathbf{W}^{(1)},\mathbf{b})}\right)$  is  $L$  -smooth to a. Let  $\eta^{(t)} = \frac{1}{L},\lambda^{(t)} = 0,$  for all  $t\in \{2,3,\dots ,T\}$  . Training by Algorithm 1 with no updates for the first layer after the first gradient step, w.h.p., there exists  $t\in [T]$  such that
$$
\mathcal {L} _ {\mathcal {D}} \big (f _ {(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b})} \big) \leq \mathrm {O P T} _ {\mathbf {W} ^ {(1)}, \mathbf {b}, B _ {a 2}} + O \Big (\frac {L (\| \mathbf {a} ^ {(1)} \| _ {2} ^ {2} + B _ {a 2} ^ {2})}{T} + \sqrt {\frac {B _ {a 2} ^ {2} (\| \mathbf {W} ^ {(1)} \| _ {F} ^ {2} B _ {x} ^ {2} + \| \mathbf {b} \| _ {2} ^ {2})}{n}} \Big).
$$
Intuitively, the theorem shows that if the weight  $\mathbf{W}^{(1)}$  after a one-step gradient gives a good set of neurons in the sense that there exists a classifier on top of these neurons with low loss, then the network will learn to approximate this good classifier and achieve low loss. The proof is based on standard convex optimization and the Rademacher complexity (details in Appendix D.1).
Such an approach, while simple, has been used to obtain interesting results on network learning in existing work, which shows that  $\mathbf{W}^{(1)}$  can indeed give good neurons due to the structure of the special problems considered (e.g., parities on uniform inputs [15], or polynomials on a subspace [32]). However, it is unclear whether such intuition can still yield useful guarantees for other problems. So, for our purpose of building a general framework covering more prototypical problems, the challenge is what features from gradient descent should be considered so that the family of networks for comparison can achieve a low loss on other problems. The other challenge is that we would like to consider the typical case where the first layer weights are not frozen. In the following, we will introduce the core concept of Gradient Features to address the first challenge, and stipulate proper geometric properties of Gradient Features for the second challenge.
3.2 Core Concepts in the Gradient Feature Learning Framework
Now, we will introduce the core concept in our framework, Gradient Features, and use it to build the family of networks to derive guarantees. As mentioned, we consider the setting where the first layer is not frozen. After the network learns good features, to ensure the updates in later gradient steps of the first layer are still benign for feature learning, we need some geometric conditions about the gradient features, which are measured by parameters in the definition of Gradient Features. The conditions are general enough, so that, as shown in Section 4, many prototypical problems satisfy them and the induced family of networks enjoys low loss, leading to useful guarantees. We begin by considering what features can be learned via gradients. Note that the gradient w.r.t.  $\mathbf{w}_i$  is
$$
\begin{array}{l} \frac {\partial \mathcal {L} _ {\mathcal {D}} (f)}{\partial \mathbf {w} _ {i}} = \mathbf {a} _ {i} \mathbb {E} _ {(\mathbf {x}, y)} \left[ \ell^ {\prime} (y f (\mathbf {x})) y \left[ \sigma^ {\prime} \left(\langle \mathbf {w} _ {i}, \mathbf {x} \rangle - \mathbf {b} _ {i}\right) \right] \mathbf {x} \right] \\ = \mathbf {a} _ {i} \mathbb {E} _ {(\mathbf {x}, y)} \left[ \ell^ {\prime} (y f (\mathbf {x})) y \mathbf {x} \mathbb {I} [ \langle \mathbf {w} _ {i}, \mathbf {x} \rangle > \mathbf {b} _ {i} ] \right]. \\ \end{array}
$$
Inspired by this, we define the following notion:
Definition 3.5 (Simplified Gradient Vector). For any  $\mathbf{w} \in \mathbb{R}^d$ ,  $b \in \mathbb{R}$ , a Simplified Gradient Vector is
$$
G (\mathbf {w}, b) := \mathbb {E} _ {(\mathbf {x}, y) \sim \mathcal {D}} [ y \mathbf {x} \mathbb {I} [ \mathbf {w} ^ {\top} \mathbf {x} > b ] ]. \tag {4}
$$
Remark 3.6. Note that the definition of  $G(\mathbf{w}, b)$  ignores the term  $\ell'(yf(\mathbf{x}))$  in the gradient, where  $f$  is the model function. In the early stage of training (or the first gradie i.e.,  $\ell'(yf(\mathbf{x})) \approx \ell'(0)$ ) due to the symmetric initialization
Definition 3.7 (Gradient Feature). For a unit vector  $D \in \mathbb{R}^d$  with  $\| D \|_2 = 1$ , and a  $\gamma \in (0,1)$ , a direction neighborhood (cone)  $\mathcal{C}_{D,\gamma}$  is defined as:
$$
\mathcal {C} _ {D, \gamma} := \left\{\mathbf {w} \mid | \langle \mathbf {w}, D \rangle | / \| \mathbf {w} \| _ {2} > (1 - \gamma) \right\}. \tag {5}
$$
Gradient Feature being cones under Mixture of Gaussians data
Figure 1: An illustration of Gradient Feature, i.e., Definition 3.7 with random initialization (Gaussian), under Mixture of three Gaussian clusters in 3-dimension data space with blue/green/orange color. The Gradient Feature stays in three cones, where each center of the cone aligns with the corresponding Gaussian cluster center.
Let  $\mathbf{w} \in \mathbb{R}^d$ ,  $b \in \mathbb{R}$  be random variables drawn from some distribution  $\mathcal{W}, \mathcal{B}$ . A Gradient Feature set with parameters  $p, \gamma, B_G$  is defined as:
$$
S _ {p, \gamma , B _ {G}} (\mathcal {W}, \mathcal {B}) := \left\{(D, s) \mid \Pr_ {\mathbf {w}, b} \left[ G (\mathbf {w}, b) \in \mathcal {C} _ {D, \gamma}, \| G (\mathbf {w}, b) \| _ {2} \geq B _ {G}, s = b / | b | \right] \geq p \right\}. \tag {6}
$$
Remark 3.8. When clear from context, write it as  $S_{p,\gamma,B_G}$ . Gradient features (see Figure 1 for illustration) are simply normalized vectors  $D$  that are given (approximately) by the simplified gradient vectors. (Similarly, the normalized scalar  $s$  is given by the bias  $b$ .) To be a useful gradient feature, we require the direction to be "hit" by sufficiently large simplified gradient vectors with sufficient large probability, so as to be distinguished from noise and remain useful throughout the gradient steps. Later we will use the gradient features when  $\mathcal{W},\mathcal{B}$  are the initialization distributions.
To make use of the gradient features, we consider the following family of networks using these features and with bounded norms, and will provide guarantees compared to the best in this family:
Definition 3.9 (Gradient Feature Induced Networks). The Gradient Feature Induced Networks are:
$$
\mathcal {F} _ {d, m, B _ {F}, S} := \left\{f _ {(\mathbf {a}, \mathbf {W}, \mathbf {b})} \in \mathcal {F} _ {d, m} \mid \forall i \in [ m ], | \mathbf {a} _ {i} | \leq B _ {a 1}, \| \mathbf {a} \| _ {2} \leq B _ {a 2}, \left(\mathbf {w} _ {i}, \mathbf {b} _ {i} / | \mathbf {b} _ {i} |\right) \in S, | \mathbf {b} _ {i} | \leq B _ {b} \right\},
$$
where  $S$  is some Gradient Feature set and  $B_{F} \coloneqq (B_{a1}, B_{a2}, B_{b})$  are some parameters.
Remark 3.10. In above definition, the weight and bias of a neuron are simply the scalings of some item in the feature set  $S$  (for simplicity the scaling of  $\mathbf{w}_i$  is absorbed into the scaling of  $\mathbf{a}_i$  and  $\mathbf{b}_i$ ).
Definition 3.11 (Optimal Approximation via Gradient Features). The optimal approximation network and loss using Gradient Feature Induced Networks  $\mathcal{F}_{d,r,B_F,S}$  are defined as:
$$
f ^ {*} := \underset {f \in \mathcal {F} _ {d, r, B _ {F}, S}} {\operatorname {a r g m i n}} \mathcal {L} _ {\mathcal {D}} (f), \quad \mathrm {O P T} _ {d, r, B _ {F}, S} := \underset {f \in \mathcal {F} _ {d, r, B _ {F}, S}} {\operatorname {m i n}} \mathcal {L} _ {\mathcal {D}} (f). \tag {7}
$$
3.3 Provable Guarantee via Gradient Feature Learning
To obtain the guarantees, we first specify the symmetric initialization. It is convenient for the analysis and is typical in existing analysis (e.g., [7, 32, 33, 105]), though some other initialization can also work. Formally, we train a two-layer network with  $4m$  neurons,  $f_{(\mathbf{a},\mathbf{W},\mathbf{b})}\in \mathcal{F}_{d,4m}$ . We initialize  $\mathbf{a}_i^{(0)},\mathbf{w}_i^{(0)}$  from Gaussians and  $\mathbf{b}_i$  from a constant for  $i\in \{1,\dots ,m\}$ , and initialize the parameters for  $i\in \{m + 1,\ldots ,4m\}$  accordingly to get a zero output initial network. Specifically:
$$
\text {f o r} i \in \{1, \dots , m \}: \quad \mathbf {a} _ {i} ^ {(0)} \sim \mathcal {N} (0, \sigma_ {a} ^ {2}), \mathbf {w} _ {i} ^ {(0)} \sim \mathcal {N} (0, \sigma_ {w} ^ {2} I), \mathbf {b} _ {i} = \tilde {b},
$$
$$
\text {f o r} i \in \{m + 1, \dots , 2 m \}: \quad \mathbf {a} _ {i} ^ {(0)} = - \mathbf {a} _ {i - m} ^ {(0)}, \mathbf {w} _ {i} ^ {(0)} = - \mathbf {w} _ {i - m} ^ {(0)}, \mathbf {b} _ {i} = - \mathbf {b} _ {i - m}, \tag {8}
$$
$$
\text {f o r} i \in \{2 m + 1, \dots , 4 m \}: \quad \mathbf {a} _ {i} ^ {(0)} = - \mathbf {a} _ {i - 2 m} ^ {(0)}, \mathbf {w} _ {i} ^ {(0)} = \mathbf {w} _ {i - 2 m} ^ {(0)}, \mathbf {b} _ {i} = \mathbf {b} _ {i - 2 m},
$$
where  $\sigma_{a}^{2},\sigma_{w}^{2},\tilde{b} >0$  are hyper-parameters. After initialization,  $\mathbf{a}$ ,  $\mathbf{W}$  are updated as in Algorithm 1. We are now ready to present our main result in the framework.
Theorem 3.12 (Main Result). Assume Assumption 3.1. For any  $\epsilon, \delta \in (0,1)$ , if  $m \leq e^d$  and
$$
m = \Omega \left(\frac {1}{p \epsilon^ {4}} \left(r B _ {a 1} B _ {x 1} \sqrt {\frac {B _ {b}}{B _ {G}}}\right) ^ {4} + \frac {1}{\sqrt {\delta}} + \frac {1}{p} \left(\log \left(\frac {r}{\delta}\right)\right) ^ {2}\right),
$$