Weight Conditioning for Smooth Optimization of Neural Networks
This addresses optimization challenges in neural network training for researchers and practitioners, offering a novel method that improves convergence, though it is incremental relative to prior normalization techniques.
The paper tackles the problem of ill-conditioned weight matrices in neural networks by introducing weight conditioning, a normalization technique that narrows the gap between singular values to smooth the loss landscape, resulting in enhanced convergence and outperforming existing methods across architectures like CNNs, ViTs, and NeRFs.
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.