Weight Conditioning for Smooth Optimization of Neural Networks

TL;DR

This paper introduces a weight conditioning normalization technique that improves the conditioning of neural network weight matrices, leading to smoother loss landscapes and enhanced convergence across diverse architectures.

Contribution

The paper proposes a novel weight normalization method inspired by linear algebra, demonstrating theoretical smoothing of the loss landscape and empirical superiority over existing techniques.

Findings

Improves convergence of stochastic gradient descent.

Outperforms existing weight normalization methods.

Effective across various neural network architectures.

Abstract

In this article, we introduce a novel normalization technique for neural network weight matrices, which we term weight conditioning. This approach aims to narrow the gap between the smallest and largest singular values of the weight matrices, resulting in better-conditioned matrices. The inspiration for this technique partially derives from numerical linear algebra, where well-conditioned matrices are known to facilitate stronger convergence results for iterative solvers. We provide a theoretical foundation demonstrating that our normalization technique smoothens the loss landscape, thereby enhancing convergence of stochastic gradient descent algorithms. Empirically, we validate our normalization across various neural network architectures, including Convolutional Neural Networks (CNNs), Vision Transformers (ViT), Neural Radiance Fields (NeRF), and 3D shape modeling. Our findings indicate that our normalization method is not only competitive but also outperforms existing weight normalization techniques from the literature.