Nonlinear Initialization Methods for Low-Rank Neural Networks

TL;DR

This paper introduces a new low-rank initialization method for deep neural networks that focuses on function approximation rather than parameter approximation, providing theoretical insights and practical algorithms validated on ImageNet.

Contribution

It presents a novel low-rank initialization framework, proves the computational tractability of low-rank ReLU approximation, and demonstrates improved training performance on large-scale models.

Findings

Significant gap between parameter and function approximation for low-rank methods.

First provably efficient algorithm for ReLU low-rank approximation.

Validated approach on ResNet and EfficientNet models on ImageNet.

Abstract

We propose a novel low-rank initialization framework for training low-rank deep neural networks -- networks where the weight parameters are re-parameterized by products of two low-rank matrices. The most successful prior existing approach, spectral initialization, draws a sample from the initialization distribution for the full-rank setting and then optimally approximates the full-rank initialization parameters in the Frobenius norm with a pair of low-rank initialization matrices via singular value decomposition. Our method is inspired by the insight that approximating the function corresponding to each layer is more important than approximating the parameter values. We provably demonstrate that there is a significant gap between these two approaches for ReLU networks, particularly as the desired rank of the approximating weights decreases, or as the dimension of the inputs to the layer increases (the latter point holds when the network width is super-linear in dimension). Along the way, we provide the first provably efficient algorithm for solving the ReLU low-rank approximation problem for fixed parameter rank $r$ -- previously, it was unknown that the problem was computationally tractable to solve even for rank $1$. We also provide a practical algorithm to solve this problem which is no more expensive than the existing spectral initialization approach, and validate our theory by training ResNet and EfficientNet models (He et al., 2016; Tan & Le, 2019) on ImageNet (Russakovsky et al., 2015).