Low-rank bias, weight decay, and model merging in neural networks

TL;DR

This paper investigates the low-rank structure and properties of neural network weights at stationary points under L2 regularization, revealing alignment, norm preservation, and a multitask learning phenomenon through weight merging.

Contribution

It provides a theoretical analysis of low-rank bias and weight properties at stationary points, and demonstrates a novel multitask learning method by merging trained networks.

Findings

Low-rank bias and parameter alignment at stationary points.

Norm preservation across layers in neural networks.

Merging trained networks enables multitask learning with comparable performance.

Abstract

We explore the low-rank structure of the weight matrices in neural networks at the stationary points (limiting solutions of optimization algorithms) with $L2$ regularization (also known as weight decay). We show several properties of such deep neural networks, induced by $L2$ regularization. In particular, for a stationary point we show alignment of the parameters and the gradient, norm preservation across layers, and low-rank bias: properties previously known in the context of solution of gradient descent/flow type algorithms. Experiments show that the assumptions made in the analysis only mildly affect the observations. In addition, we investigate a multitask learning phenomenon enabled by $L2$ regularization and low-rank bias. In particular, we show that if two networks are trained, such that the inputs in the training set of one network are approximately orthogonal to the inputs in the training set of the other network, the new network obtained by simply summing the weights of the two networks will perform as well on both training sets as the respective individual networks. We demonstrate this for shallow ReLU neural networks trained by gradient descent, as well as deep linear networks trained by gradient flow.