ReLU soothes NTK conditioning and accelerates optimization for wide neural networks

Chaoyue Liu, Like Hui, Xiao Liu

Primary Area: optimization

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Keywords: ReLU, non-linear activation function, condition number, NTK, neural tangent kernel, convergence rate

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

TL;DR: we showcase a new and interesting property of certain non-linear activations, focusing on ReLU: the non-linearity, and the depth of ReLU network, help to decrease the NTK condition number and accelerate optimization for wide neural networks

Abstract:

Non-linear activation functions are well known to improve the expressivity of neural networks, which is the main reason of their wide implementation in neural networks. In this work, we showcase a new and interesting property of certain non-linear activations, focusing on the most popular example of its kind -- Rectified Linear Unit (ReLU). By comparing the cases with and without this non-linear activation, we show that the ReLU has the following effects: (a) *better data separation*, i.e., a larger angle separation for similar data in the feature space of model gradient, and (b) *better NTK conditioning*, i.e., a smaller condition number of neural tangent kernel (NTK). Furthermore, we show that the ReLU network depth (i.e., with more ReLU activation operations) further magnifies these effects. Note that, without the non-linear activation, i.e., in a linear neural network, the data separation and NTK condition number always remain the same as in the case of a linear model, regardless of the network depth.
Our results imply that ReLU activation, as well as the depth of ReLU network, helps improve the worst-case convergence rate of gradient descent, which is closely related to the NTK condition number.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 4073