Compute Better Spent: Replacing Dense Layers with Structured Matrices

TL;DR

This work targets the dense linear layers that dominate compute in foundation models and investigates structured matrices as universal, compute-efficient alternatives. By extending Maximal Update Parameterization (μP) to structured layers, the authors derive structure-aware initialization and learning-rate scales, enabling stable and scalable training across diverse architectures. They introduce Block Tensor-Train (BTT), showing it can outperform dense layers at the same compute on CIFAR and ImageNet, and improve GPT-2 efficiency; they also demonstrate stability improvements via weight normalization for large transformers. The results identify the importance of compute-per-dimension trade-offs and reveal that certain structures can exhibit superior scaling laws, offering practical directions for more efficient large-scale modeling.

Abstract

Dense linear layers are the dominant computational bottleneck in foundation models. Identifying more efficient alternatives to dense matrices has enormous potential for building more compute-efficient models, as exemplified by the success of convolutional networks in the image domain. In this work, we systematically explore structured matrices as replacements for dense matrices. We show that different structures often require drastically different initialization scales and learning rates, which are crucial to performance, especially as models scale. Using insights from the Maximal Update Parameterization, we determine the optimal scaling for initialization and learning rates of these unconventional layers. Finally, we measure the scaling laws of different structures to compare how quickly their performance improves with compute. We propose a novel matrix family containing Monarch matrices, the Block Tensor-Train (BTT), which we show performs better than dense matrices for the same compute on multiple tasks. On CIFAR-10/100 with augmentation, BTT achieves exponentially lower training loss than dense when training MLPs and ViTs. BTT matches dense ViT-S/32 performance on ImageNet-1k with 3.8 times less compute and is more efficient than dense for training small GPT-2 language models.

Figures (12)

• Figure 1: Controlling for compute, replacing dense layers with structured matrices enables wider models and can lead to better performance. (a) A neural network with structured matrices can be made much wider, but its learning rate needs to be scaled differently as a function of width since not all connections are present (\ref{['sec:sucessful']}). The width $d$ of a dense layer scales as $C^{1/2}$ where $C$ is the compute per forward pass, while the width $\tilde{d}$ of a block diagonal layer is exponentially larger, scaling as $C^{2/3}.$ The optimal learning rate $\eta$ of the dense layer and $\tilde{\eta}$ of the block diagonal layer scales differently as $d^{-1}$ and $\tilde{d}^{-1/2}$ respectively. (b) Structured matrices can improve the training error scaling laws of MLPs on CIFAR-10 with data augmentation (\ref{['sec:experiments']}). (c) Scaling the learning rate in a structure-aware fashion ($\bullet$) is crucial for performance (\ref{['sec:sucessful']}), without which the benefit of structured layers does not emerge ($\blacktriangledown$).
• Figure 2: Structure-aware learning rate scaling results in stable feature learning and stable optimal learning rate as we vary the structure and model size. (a) The RMS of the changes $\Delta h$ of the last layer features is stable as the models are scaled up in width, but is smaller or vanishes if we naively use the learning rate for the dense model. (b) The optimal learning rate is stable as we vary the structure and width, provided we use structure-aware learning rates. Here we use Monarch with 16 blocks.
• Figure 3: Structure-aware learning rates improve performance even after tuning the learning with grid search. Test error of ViT ($d=1024$) on CIFAR-10 where the feed-forward layers are replaced using BTT.
• Figure 4: Using structured matrices changes the scaling laws of MLPs and ViTs trained on CIFAR-100. We find 1) BTT achieves the best scaling, and 2) structures with FLOPs equal to parameters scale better than those with parameter sharing (Kronecker and TT)
• Figure 5: Less compute per dimension is more compute-efficient on CIFAR-10. (a) BTT with a lower rank achieves lower train error per FLOP. (b) Monarch with more blocks achieves lower train error per FLOP. A lighter color indicates less compute per dimension.