Provable Guarantees for Nonlinear Feature Learning in Three-Layer Neural Networks

Summary

This paper studies the feature learning capability of a 3-layer neural network, where the bottom layer is random and fixed, the middle layer is trained for only one step from zero, and the upper layer is trained in the rest of the gradient descent steps. The paper characterizes the richer feature learning capability of this 3-layer network. That is, the one-step training of the middle-layer weights essentially maps the original random features to another space, which enlarges the type of functions that can be learned. Two special cases are then considered: (i) the single index model (similar to [18]); and (ii) more significantly, the quadratic feature model (more complex than [18]). The results demonstrate a significant reduction in the sample complexity compared to kernel methods. Further, the latter class of functions in (ii) cannot be approximated by 2-layer networks, and thus a depth-separation between 2-layer and 3-layer networks is established.

Strengths

1. The results are novel and solid. The precisely characterization of the feature learning capability of this particular 3-layer network (Theorem 1), and the ability to learn functions of quadratic features, are both significant contributions to the field. 2. The paper is also very well-written and presents the main intuitions quite clearly.

Weaknesses

Nothing particular.

Questions

1. Can the authors comment on other possible features that this approach can learn effectively (beyond quadratic features)? 2. The one-step training of the middel-layer weights seems to be a crucial element for learning a specific type of features. At the same time, one-step training could be limiting. Can the authors comment on possible generalizations for multi-step training on the middle-layer? Post rebuttal phase: The reviewer wishes to thank the authors for their response and preliminary thoughts on these questions.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The discussion of limitations is adequate.

Summary

The paper theoretically studies the feature learning in three layer neural networks. For the analysis, it considers layer-wise GD; more precisely, the first layer is not trained, the second layer is trained for one step, and then there is the training for last layer. Particularly, they show that three layer neural networks achieve a better bound for learning functions of the form $g(x^TAx)$ comparing to the known bounds for two-layer neural networks. They also show an optimization-based separation showing that three-layer NNs can learn functions that two-layer NNs cannot learn with polynomial width and time (assuming a non-increasing learning rate).

Strengths

- The paper considers the optimization of three-layer NNs with a layer-wise training and provides a fairly general approach for studying their feature learning capabilities (in the paper, the features have also been characterized for the examples considered). - For functions of the type $g(x^TAx)$, a sample complexity upper bound for three-layer NNs is proved which is more efficient than the known upper bounds for two-layer NNs. - There is an optimization-based separation result provided between two-layer and three-layer NNs. - Paper is generally very well written.

Weaknesses

- Generally the limitations of the work have been discussed in the paper. These limitations are usually quite common in the deep learning theory literature. For example, the training is layer-wise, and the first layer is not trained at all and second layer is only trained for one iteration. In this sense, the work is similar to theoretical analysis done on two-layer NNs in many recent works. - The simulations have been implemented for the settings as the theorems. However, it would have been interesting to also run experiments with the more common setting (e.g., training all parameters together) for both two-layer and three-layer NNs and compare the results with the results of the model in the theoretical settings (e.g., with layer-wise training). This would potentially show if the assumptions made in theory are too restrictive.

Questions

- Can the analysis be done in the SGD setting as well? Here are also a few typos: - line 246: $\mathbb{E}[(g^{*})'] = 0$ instead of $\mathbb{E}[g'] = 0$ - line 291: inconsistency regarding $P_2 f^*$ (e.g., with equation 16) (I think also the same problem appeared in the appendix.) - line 298: is it the intended equation?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and 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.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

There is no negative societal impact. The limitations of theory have been discussed in the work and are (unfortunately) quite common in the current literature of deep learning theory.

Summary

In this work the authors show that there is a three layer neural network setup with better provable learning guarantees than the current best bound for two layer setups. The setup involves a randomly initialized layer with frozen weights, which feeds into a two layer network where first the hidden layer weights are trained, and then the output weights are trained. In the case of functions of quadratic features, this three layer (but two learnable layers) network is able to use feature learning to re-weight the initial random feature kernel to improve training efficiency.

Strengths

The model used in the paper is a simple and sensible extension of the two layer network. The learning algorithm is also quite reasonable, both in its connection to previous work as well as to the practical end of deep learning. The arguments seem correct (though I did not fully validate the detailed proofs in the appendices), and the intuitive explanation relating feature learning and eigenspace weighting helps shed light on the importance (and perhaps mechanisms of) feature learning in deep networks. The bounds in the 3 layer case are a marked improvement over the kernel learning and 2 hidden layer case bounds.

Weaknesses

There are two concerns with this paper. The first is a question of the tightness of the bounds. While the bounds show great improvement over the quoted bounds in the two hidden layer case, it is not clear how tight these bounds are in the various scenarios discussed. Relatedly, there is a big question as to whether or not the bounds are useful for understanding the success of deep learning systems even on simple problems (e.g. FCN on MNIST). In particular, the proof sketch seems to suggest that just a single step of GD in the middle layer is enough to induce massive improvements in the sample complexity; however even in simple settings, it seems that it is helpful to both learn in multiple layers, as well as over multiple timepoints. I note that this is a general weakness of similar sample complexity analyses and not of this particular paper, and the authors do mention this in the discussion.

Questions

What evidence is there that the difference between the 2 and 3 layer bounds will persist even as theoretical techniques/the bounds themselves are improved? How do the 3 layer feature learning bounds differ from NTK learning of a similar 3 layer network? (This discussion would be useful to add to the main text to emphasize the benefits of feature learning over depth alone.) Can the method be used to show anything about the usefulness of running gradient descent over a longer period of time? How does this trade off with the sample complexity? Is there any intuition about good choices of the function q in the bounds, in a more general setting?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, 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.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

This paper analyzed the features learned by a three-layer network trained with layer-wise gradient descent as existing analyses are largely restricted to two-layer networks. It presented a general purpose theorem that upper bounds the sample complexity and width needed to achieve low test error when the target has a certain hierarchical structure.

Strengths

- This paper has shown that a three-layer network can learn nonlinear features, although existing work only has shown that two-layer neural networks learn only linear functions of input. - It’s good that the paper has an example section (Sec. 4.)

Weaknesses

- The “Hierarchical function” that this paper deals with is not early enough defined or introduced. It will make the problem that the paper solved more clear if that point is made clear in an earlier part of the paper. - This work still does not answer the same questions for convolutional networks. - This paper certainly showed that three-layer networks have provably richer feature learning capabilities than two-layer networks. Nonetheless, it is not guaranteed that the findings and proof can still be applied to any general “depth” as it has shown for the change from the two-layer to three-layer (although it might be high-likely.) - In other words, this issue is about generalizability on depth scalability. - This paper lacks discussion of its own limitations.

Questions

Do the authors think this can be applied to convolutional networks and networks with more depth than three layers?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

No limitations are discussed in the paper. It would have been nicer if the authors discussed the generalizability of the work to a deeper network.

Response to authors

I thank the authors for their detailed responses. Between their comments here and their comments to the rest of the reviewers, my main concerns have been addressed and I will update my review score.