Summary
This paper studies the training of 2-layer neural networks and proposes a two-timescale limit/regime. In this
limit/regime, the learning rate of the first layer is much smaller than the learning rate of the second layer.
As a result, the training of the network can be viewed as training the first layer and performing linear regression
over the first layer outputs, which is simpler/more structured than training both layers at the same rate. To
demonstrate the usefulness of this strategy, this paper considers a toy model and shows that a 2-layer network can fit
a certain family of 1D piecewise constant functions. The authors also empirically show that SGD can fail to fit this
family of functions outside the two-timescale regime.
Strengths
* The presentation is clear.
* This paper not only contains results in the $\varepsilon \to 0$ limit, but also non-asymptotic results for small but
non-vanishing $\varepsilon$. The derivation in the $\eta \to 0, \varepsilon \to 0$ limit is clean (Sec. 4.2), and
the asymptotic-to-non-asymptotic parts are also easy to follow.
* This paper reminds me of a technique that is gaining popularity in the theory community: training the network for
one (large) step, freezing the first layer, and then performing linear regression in the second layer using the
features learned in the first layer (cf. [AAM22], [DLS22]). The limitation of this technique is that since the
first layer is fixed after the first layer, it lacks the ability to refine the learned features. I feel the
strategy introduced in this paper is a potential remedy to this problem, as here we also have the linear regression
part of the argument, but the first layer can be trained for multiple steps.
[AAM22] Abbe, Emmanuel, 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 Proceedings of Thirty Fifth Conference on Learning Theory, 4782–4887. PMLR, 2022. https://proceedings.mlr.press/v178/abbe22a.html.
[DLS22] Damian, Alex, Jason D. Lee, and Mahdi Soltanolkotabi. “Neural Networks Can Learn Representations with Gradient Descent.” arXiv, June 30, 2022. http://arxiv.org/abs/2206.15144.
Weaknesses
* Almost all things in this paper are 1D and somewhat tailored to this specific piecewise constant function class.
I wonder whether/how this can be generalized to higher dimensions, other network architectures, and more general
function classes.
* It seems that the dynamics are still relatively local. That is, we need some neurons in each target interval at
initialization, which may not be reasonable when the dimension is high.
* The authors should probably add some discussion on the training-for-one-large-step-type technique (see the Strengths
part of the review).
Questions
* How general is this strategy? Is it possible to apply this strategy to problems (even toy ones) with dimension higher
than $1$ and more standard 2-layer networks?
* Can we relax the requirement on the initialization? For example, can we remove the requirement that each interval
contains at least $6$ $u_j$?
Intuitively, when some interval contains no neurons, it is still possible that some neurons from those fitted
intervals can move to that interval and fit the target function on that interval.
I know the current requirement is reasonable in your setting, I am asking this because I feel this type of global
feature learning process is important when the dimension is high.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Limitations
This is a theoretical work and, as far as I can see, has no potential negative societal impacts.