SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lottery Ticket Hypothesis (LTH). Another interesting phenomenon is the LTH [41]: randomly-initialized networks contain subnetworks that when trained in isolation reach test accuracy comparable to the original network in a similar number of iterations. Later studies (e.g., [42]) show that LTH is more stable when subnetworks are found in the network after a few gradient steps.
Our framework provides an explanation for two-layer networks: the lottery ticket subnetwork contains exactly those neurons whose gradient feature approximates the weights of the "ground-truth" network  $f^{*}$ ; they may not exist at initialization but can be found after the first gradient step. More precisely, Lemma 3.14 shows that after the first gradient step, there is a sparse second-layer weight  $\tilde{\mathbf{a}}$  with  $\| \tilde{\mathbf{a}} \|_0 = O(r\sqrt{mp})$ , such that using this weight on the hidden neurons gives a network with a small loss. Let  $U$  be the support of  $\tilde{\mathbf{a}}$ . Equivalently, there is a small-loss subnetwork  $f_{\Xi}^{U}$  with only neurons in  $U$  and with second-layer weight  $\tilde{\mathbf{a}}_U$  on these neurons. Following the same proof of Theorem 3.12:
Proposition 5.1. In the same setting of Theorem 3.12 but only considering the subnetwork supported on  $U$  after the first gradient step, with the same requirements on  $m$  and  $T$ , with proper hyperparameter values, we have the same guarantee: with probability  $\geq 1 - \delta$ , there is  $t \in [T]$  with
$$
\Pr \left[ \operatorname {s i g n} \left(f _ {\Xi^ {(t)}} ^ {U}\right) (\mathbf {x}) \neq y \right] \leq \operatorname {O P T} _ {d, r, B _ {F}, S _ {p, \gamma , B _ {G}}} + r B _ {a 1} B _ {x 1} \sqrt {2 \gamma + O \left(\frac {\sqrt {B _ {x 2} \log n}}{B _ {G} \sqrt {n}}\right)} + \epsilon .
$$
This essentially formally proves LTH for two-layer networks, showing (a) the existence of the winning lottery subnetwork and (b) that gradient descent on the subnetwork can learn to similar loss in similar runtime as on the whole network. In particular, (b) is novel and not analyzed in existing work.
We provide our work's broader impacts and limitations (e.g., statement of recovering existing results and some failure cases beyond our framework) in Appendix A and Appendix B respectively.
Conclusion. We propose a general framework for analyzing two-layer neural network learning by gradient descent and show that it can lead to provable guarantees for several prototypical problem settings for analyzing network learning. In particular, our framework goes beyond fixed feature methods, e.g., NTK. It sheds light on several interesting phenomena in NN learning, e.g., the lottery ticket hypothesis and simplicity bias. Future directions include: (1) How to extend the framework to deeper networks? (2) While the current framework focuses on the gradient features in the early gradient steps, whether feature learning also happens in later steps and if so how to formalize that?
The work is partially supported by Air Force Grant FA9550-18-1-0166, the National Science Foundation (NSF) Grants 2008559-IIS, 2023239- DMS, and CCF-2046710.
[1] Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory. PMLR, 2022.
[2] Emmanuel Abbe, Samy Bengio, Elisabetta Cornacchia, Jon Kleinberg, Aryo Lotfi, Maithra Raghu, and Chiyuan Zhang. Learning to reason with neural networks: Generalization, unseen data and boolean measures. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[3] Shunta Akiyama and Taiji Suzuki. On learnability via gradient method for two-layer relu neural networks in teacher-student setting. In International Conference on Machine Learning, pages 152-162. PMLR, 2021.
[4] Shunta Akiyama and Taiji Suzuki. Excess risk of two-layer reLU neural networks in teacher-student settings and its superiority to kernel methods. In The Eleventh International Conference on Learning Representations, 2023.
[5] Zeyuan Allen-Zhu and Yuanzhi Li. What can resnet learn efficiently, going beyond kernels? In Advances in Neural Information Processing Systems, 2019.
[6] Zeyuan Allen-Zhu and Yuanzhi Li. Backward feature correction: How deep learning performs deep learning. arXiv preprint arXiv:2001.04413, 2020.
[7] Zeyuan Allen-Zhu and Yuanzhi Li. Feature purification: How adversarial training performs robust deep learning. In 2021 IEEE 62nd Annual Symposium on Foundations of Computer Science (FOCS), pages 977-988. IEEE, 2022.
[8] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. In Advances in neural information processing systems, 2019.
[9] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning, 2019.
[10] Sanjeev Arora, Nadav Cohen, Noah Golowich, and Wei Hu. A convergence analysis of gradient descent for deep linear neural networks. In International Conference on Learning Representations, 2018.
[11] 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, pages 322-332. PMLR, 2019.
[12] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Ruslan Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. arXiv preprint arXiv:1904.11955, 2019.
[13] Jimmy Ba, Murat A Erdogdu, Taiji Suzuki, Zhichao Wang, Denny Wu, and Greg Yang. High-dimensional asymptotics of feature learning: How one gradient step improves the representation. arXiv preprint arXiv:2205.01445, 2022.
[14] Yu Bai and Jason D Lee. Beyond linearization: On quadratic and higher-order approximation of wide neural networks. In International Conference on Learning Representations, 2019.
[15] Boaz Barak, Benjamin L Edelman, Surbhi Goel, Sham Kakade, Eran Malach, and Cyril Zhang. Hidden progress in deep learning: Sgd learns parities near the computational limit. arXiv preprint arXiv:2207.08799, 2022.
1 Introduction 1
2 Related Work 2
3 Gradient Feature Learning Framework 3
3.1 Warm Up: A Simple Setting with Frozen First Layer 3
3.2 Core Concepts in the Gradient Feature Learning Framework 4
3.3 Provable Guarantee via Gradient Feature Learning 5
4 Applications in Special Cases 7
4.1 Mixtures of Gaussians 7
4.1.1 Mixtures of Gaussians: Beyond the Kernel Regime 8
4.2 Parity Functions 8
4.2.1 Parity Functions: Beyond the Kernel Regime 10
5 Further Implications and Conclusion 10
A Broader Impacts 19
B Limitations 19
C More Further Implications 20
D Gradient Feature Learning Framework 21
D.1 Simplified Gradient Feature Learning Framework 21
D.2 Gradient Feature Learning Framework under Expected Risk 22
D.2.1 Feature Learning 23
D.2.2 Good NetworkExists 24
D.2.3 Learning an Accurate Classifier 26
D.3 More Discussion abut Setting 30
D.4 Gradient Feature Learning Framework under Empirical Risk with Sample Complexity 31
E Applications in Special Cases 38
E.1 Linear Data 38
E.2 Mixture of Gaussians 40
E.2.1 Problem Setup 40
E.2.2 Mixture of Gaussians: Feature Learning 41
E.2.3 Mixture of Gaussians: Final Guarantee 46
E.3 Mixture of Gaussians - XOR 48
E.3.1 Mixture of Gaussians - XOR: Feature Learning 48
E.3.2 Mixture of Gaussians - XOR: Final Guarantee 52
E.4 Parity Functions 52
E.4.1 Problem Setup 52
E.4.2 Parity Functions: Feature Learning 54
E.4.3 Parity Functions: Final Guarantee 58
E.5 Uniform Parity Functions 59
E.5.1 Uniform Parity Functions: Final Guarantee 60
E.6 Uniform Parity Functions: Alternative Analysis 61
E.6.1 Modified General Feature Learning Framework for Uniform Parity Functions 61
E.6.2 Feature Learning of Uniform Parity Functions 66
E.7 Multiple Index Model with Low Degree Polynomial 68
E.7.1 Problem Setup 68
E.7.2 Multiple Index Model: Final Guarantee 68
F Auxiliary Lemmas 70
Appendix A discusses the potential societal impact of our work. Appendix B describes the limitations of our work. In Appendix C, we present our framework implications about simplicity bias. The complete proof of our main results is given in Appendix D. We present the case study of linear data in Appendix E.1, mixtures of Gaussians in Appendix E.2 and Appendix E.3, parity functions in Appendix E.4, Appendix E.5 and Appendix E.6, and multiple-index models in Appendix E.7. We put the auxiliary lemmas in Appendix F.
Our paper is purely theoretical in nature, and thus we do not anticipate an immediate negative ethical impact. We provide a unified theoretical framework that can be applied to different theoretical problems. We propose the two key ideas of gradient feature and gradient feature-induced neural networks not only to show their ability to unify several current works but also to open a new direction of thinking with respect to the learning process. These notations have the potential to be extended to multi-layer gradient features and multi-step learning, and this work is only our first step.
On the other hand, this work may lead to a better understanding and inspire the development of improved network learning methods, which may have a positive impact on the theoretical machine-learning community. It may also be beneficial to engineering-inclined machine-learning researchers.
Recover Existing Results. The framework may or may not recover the width or sample complexity bounds in existing work.
1. The framework can give matching bounds as the existing work in some cases, like parities over uniform inputs (Appendix E.5).
2. In some other cases, it gives polynomial error bounds not the same as those in the existing work (e.g., for parities over structured inputs). This is because our work is analyzing general cases, and thus may not give better than or the same bounds as those in special cases, since special cases have more properties that can be exploited to get potentially better bounds. On the other hand, our bounds can already show the advantage over kernel methods (e.g., Proposition 4.9).
We would like to emphasize that our contribution is providing an analysis framework that can (1) formalize the unifying principles of learning features from gradients in network training, and (2) give
polynomial error bounds for prototypical problems. Our focus is not to recover the guarantees in existing work.
Failure Cases. There are some failure cases that gradient feature learning framework cannot cover:
1. In [101], they constructed a function that is easy to approximate using a 3-layer network but not approximable by any 2-layer network. Since the function is not approximable by any 2-layer network, it cannot be approximated by the gradient-induced networks as well, so OPT will be large. As a result, the final error will be large.
2. In uniform parity data distribution, considering an odd number of features rather than even, i.e.,  $k$  is an odd number in Assumption E.28, we can show that our gradient feature set is empty even when  $p$  in Equation (6) is exponentially small, thus the OPT is a positive constant since the gradient induced network can only be constants. Meanwhile, the neural network won't be able to learn this data distribution because its gradient is always 0 through the training, and the final error equals OPT.
The first case corresponds to the approximation hardness of 2-layer networks, while the second case gives a learning hardness example. The above two cases show that if there is an approximation or learning hardness, our gradient feature learning framework may be vacuous because the optimal model in the gradient feature class has a large risk, then the ground-truth mapping from inputs to labels is not learnable by gradient descent. These analyses are consistent with previous works [15, 101].
C More Further Implications
Our general framework also sheds some light on several interesting phenomena in neural network (NN) learning observed in practice. Feature learning beyond the kernel regime has been discussed in Section 4.1.1 and Section 4.2.1. The lottery ticket hypothesis (LTH) has been discussed in Section 5. Below we discuss other implications.
Implicit Regularization/Simplicity Bias. It is now well known that practical NN are overparameterized and traditional uniform convergence bounds cannot adequately explain their generalization performance [59, 88, 118]. It is generally believed that the optimization has some implicit regularization effect that restricts learning dynamics to a subset of the whole hypothesis class, which is not of high capacity so can lead to good generalization [53, 90]. Furthermore, learning dynamics tend to first learn simple functions and then learn more and more sophisticated ones (referred to as simplicity bias) [89, 102]. However, it remains elusive to formalize such simplicity bias.
Our framework provides a candidate explanation: the learning dynamics first learn to approximate the best network in a smaller family of gradient feature induced networks  $\mathcal{F}_{d,r,B_F,S}$  and then learn to approximate the best in a larger family. Consider the number of neurons  $r$  for illustration. Let  $r_1 \ll r_2$ , and let  $T_1$  and  $T_2$  be their corresponding runtime bounds for  $T$  in the main Theorem 3.12. Clearly,  $T_1 \ll T_2$ . Then, at time  $T_1$ , the theorem guarantees the learning dynamics learn to approximate the best in the family  $\mathcal{F}_{d,r_1,B_F,S}$  with  $r_1$  neurons, but not for the larger family  $\mathcal{F}_{d,r_2,B_F,S}$ . Later, at time  $T_2$ , the learning dynamics learn to approximate the best in the larger family  $\mathcal{F}_{d,r_2,B_F,S}$ . That is, the learning first learns simpler functions and then more sophisticated ones where the simplicity bias is measured by the size of the family of gradient feature-induced networks. The implicit regularization is then restricting to networks approximating smaller families of gradient feature-induced networks. Furthermore, we can also conclude that for an SGD-optimized NN, its actual representation power is from the subset of NN based on gradient features, instead of the whole set of NN. This view helps explain the simplicity bias/implicit regularization phenomenon of NN learning in practice.
Learning over Different Data Distributions. Our framework articulates the following key principles (pointed out for specific problems in existing work but not articulated more generally):
- Role of gradient: the gradient leads to the emergence of good features, which is useful for the learning of upper layers in later stages.
- From features to solutions: learned features in early steps will not be distorted, if not improved, in later stages. The training dynamic for upper layers will eventually learn a good combination of hidden neurons based on gradient features, giving a good solution.
Then, more interesting insights are obtained from the generality of the framework. To build a general framework, the meaningful error guarantees should be data-dependent, since NN learning on general data distributions is hard and data-independent guarantees will be vacuous [34, 35]. Comparing the optimal in a family of "ground-truth" functions (inspired by agnostic learning in learning theory) is a useful method to obtain the data-dependent bound. We further construct the "ground-truth" functions using properties of the training dynamics, i.e., gradient features. This greatly facilitates the analysis of the training dynamics and is the key to obtaining the final guarantees. On the other hand, the framework can also be viewed as using the optimal by 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. So this provides a united way to derive guarantees for specific problems.
New Perspectives about Roadmaps Forward. We argue a new perspective about the connection between the strong representation power and the successful learning of NN. Traditionally, the strong representation power of NN is the key reason for hardness results of NN learning: NN has strong representation power and can encode hard learning questions, so they are hard to learn. See the proof in SQ bound from [33] or NP-hardness from [19]. The strong representation power also causes trouble for the statistical aspect: it leads to vacuous generalization bounds when traditional uniform convergence tools are used.
Our framework suggests a perspective in sharp contrast: the strong representation power of NN with gradient features is actually the key to successful learning. More concretely, the optimal error of the gradient feature-induced NN being small (i.e., strong representation power for a given data distribution) can lead to a small guarantee, which is the key to successful learning. The above new perspective suggests a different analysis road than traditional ones. Traditional analysis typically first reasons about the optimal based on the whole function class, i.e. the ground truth, then analyze 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.
Our framework provides the foundation for future work on analyzing gradient-based NN learning, which may inspire future directions including but not limited to (1) defining a new feature family for 2-layer NN rather than gradient feature, (2) considering deep NN and introducing new gradient features (e.g., gradient feature notion for upper layers), (3) defining different gradient feature family at different training stages (e.g., gradient feature notion for later stages). In particular, the challenges in the later-stage analysis are: (a) the weights in the later stage will not be as normal as the initialization, and we need new tools to analyze their properties; (b) to show that the later-stage features eventually lead to a good solution, we may need new analysis tools for the non-convex optimization due to the changes in the first layer weights.
D Gradient Feature Learning Framework
We first prove a Simplified Gradient Feature Learning Framework in Appendix D.1, which only considers one-step gradient feature learning. Then, we prove our Gradient Feature Learning Framework, e.g., no freezing of the first layer. In Appendix D.2, we consider population loss to simplify the proof. Then, we provide more discussion about our problem setup and our core concept in Appendix D.3. Finally, we prove our Gradient Feature Learning Framework under empirical loss considering sample complexity in Appendix D.4.
D.1 Simplified Gradient Feature Learning Framework
Algorithm 2 Training by Algorithm 1 with no updates for the first layer after the first gradient step Initialize  $f_{(\mathbf{a}^{(0)},\mathbf{W}^{(0)},\mathbf{b})}\in \mathcal{F}_{d,m}$  ; Sample  $\mathcal{Z}\sim \mathcal{D}^n$  Get  $(\mathbf{a}^{(1)},\mathbf{W}^{(1)},\mathbf{b})$  by one gradient step update and fix  $\mathbf{W}^{(1)}$  , b for  $t = 2$  to  $T$  do  $\mathbf{a}^{(t)} = \mathbf{a}^{(t - 1)} - \eta^{(t)}\nabla_{\mathbf{a}}\widetilde{\mathcal{L}}_{\mathcal{Z}}(f_{\Xi^{(t - 1)}})$  end for
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  $\mathbf{a}$ . Let  $\eta^{(t)} = \frac{1}{L}$ ,  $\lambda^{(t)} = 0$  for all  $t \in \{2,3,\ldots,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}}(f_{(\mathbf{a}^{(t)}, \mathbf{W}^{(1)}, \mathbf{b})}) \leq \mathrm{OPT}_{\mathbf{W}^{(1)}, \mathbf{b}, B_{a2}} + O\left(\frac{L(\|\mathbf{a}^{(1)}\|_2^2 + B_{a2}^2)}{T} + \sqrt{\frac{B_{a2}^2(\|\mathbf{W}^{(1)}\|_F^2B_x^2 + \|\mathbf{b}\|_2^2)}{n}}\right)$ .
Proof of Theorem 3.4. Recall that
$$
\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 \mathrm {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 {13}
$$
We denote  $f^{*} = \operatorname{argmin}_{f\in \mathcal{F}_{\mathbf{W},\mathbf{b},B_{a2}}}\mathcal{L}_{\mathcal{D}}(f)$  and  $\tilde{f}^{*} = \operatorname{argmin}_{f\in \mathcal{F}_{\mathbf{W},\mathbf{b},B_{a2}}}\widetilde{\mathcal{L}}_{\mathcal{Z}}(f)$ . We use  $\mathbf{a}^*$  and  $\tilde{\mathbf{a}}^*$  to denote their second layer weights respectively. Then, we have
$$
\begin{array}{l} \mathcal {L} _ {\mathcal {D}} \left(f _ {\left(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) = \mathcal {L} _ {\mathcal {D}} \left(f _ {\left(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) - \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) (14) \\ + \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) - \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\tilde {\mathbf {a}} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) (15) \\ + \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\tilde {\mathbf {a}} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) - \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) (16) \\ + \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) - \mathcal {L} _ {\mathcal {D}} \left(f _ {\left(\mathbf {a} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) (17) \\ + \mathcal {L} _ {\mathcal {D}} \left(f _ {\left(\mathbf {a} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) (18) \\ \leq \left| \mathcal {L} _ {\mathcal {D}} \left(f _ {\left(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) - \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) \right| (19) \\ + \left| \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) - \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\tilde {\mathbf {a}} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) \right| (20) \\ + 0 (21) \\ + \left| \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) - \mathcal {L} _ {\mathcal {D}} \left(f _ {\left(\mathbf {a} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) \right| (22) \\ + \mathrm {O P T} _ {\mathbf {W} ^ {(1)}, \mathbf {b}, B _ {a 2}}. (23) \\ \end{array}
$$
Fixing  $\mathbf{W}^{(1)}$ ,  $\mathbf{b}$  and optimizing a only is a convex optimization problem. Note that  $\eta \leq \frac{1}{L}$ , where  $\widetilde{\mathcal{L}}_{\mathcal{Z}}$  is  $L$ -smooth to  $\mathbf{a}$ . Thus with gradient descent, we have
$$
\frac {1}{T} \sum_ {t = 1} ^ {T} \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} ^ {(t)}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) - \widetilde {\mathcal {L}} _ {\mathcal {Z}} \left(f _ {\left(\mathbf {a} ^ {*}, \mathbf {W} ^ {(1)}, \mathbf {b}\right)}\right) \leq \frac {\left\| \mathbf {a} ^ {(1)} - \mathbf {a} ^ {*} \right\| _ {2} ^ {2}}{2 T \eta}. \tag {24}
$$
Then our theorem gets proved by Lemma F.9 and generalization bounds based on Rademacher complexity.
D.2 Gradient Feature Learning Framework under Expected Risk
We consider the following training process under population loss to simplify the proof. We prove our Gradient Feature Learning Framework under empirical loss considering sample complexity in Appendix D.4.
Algorithm 3 Network Training via Gradient Descent
Initialize  $(\mathbf{a}^{(0)},\mathbf{W}^{(0)},\mathbf{b})$  as in Equation (8)