Initialization and Regularization of Factorized Neural Layers

2021/05/03 by Mikhail Khodak, Neil Tenenholtz, Khodak, Mikhail +5 · 4 citations
Computer Science · Mathematics · Physics and Astronomy · #Advanced Neural Network Applications #Artificial Intelligence (cs.AI) #Computation and Language (cs.CL) #Computer Vision and Pattern Recognition (cs.CV) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Model Reduction and Neural Networks #Tensor decomposition and applications

paper · pdf · doi:10.48550/arxiv.2105.01029

openalex publication_date 2021/05/03 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Factorized layers--operations parameterized by products of two or more matrices--occur in a variety of deep learning contexts, including compressed model training, certain types of knowledge distillation, and multi-head self-attention architectures. We study how to initialize and regularize deep nets containing such layers, examining two simple, understudied schemes, spectral initialization and Frobenius decay, for improving their performance. The guiding insight is to design optimization routines for these networks that are as close as possible to that of their well-tuned, non-decomposed counterparts; we back this intuition with an analysis of how the initialization and regularization schemes impact training with gradient descent, drawing on modern attempts to understand the interplay of weight-decay and batch-normalization. Empirically, we highlight the benefits of spectral initialization and Frobenius decay across a variety of settings. In model compression, we show that they enable low-rank methods to significantly outperform both unstructured sparsity and tensor methods on the task of training low-memory residual networks; analogs of the schemes also improve the performance of tensor decomposition techniques. For knowledge distillation, Frobenius decay enables a simple, overcomplete baseline that yields a compact model from over-parameterized training without requiring retraining with or pruning a teacher network. Finally, we show how both schemes applied to multi-head attention lead to improved performance on both translation and unsupervised pre-training.

Citations

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Deeper Insights into Weight Sharing in Neural Architecture Search
• Learning Multiple Layers of Features from Tiny Images

Cited by

• BOOST: BOttleneck-Optimized Scalable Training Framework for Low-Rank Large Language Models
• Q3R: Quadratic Reweighted Rank Regularizer for Effective Low-Rank Training
• Accelerating Attention with Basis Decomposition
• LOST: Low-rank and Sparse Pre-training for Large Language Models

Related

• Training Thinner and Deeper Neural Networks: Jumpstart Regularization
• Traditional and Heavy-Tailed Self Regularization in Neural Network Models
• Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning
• Compression-aware Training of Deep Networks
• Weight decay induces low-rank attention layers