Research map / Architectures and efficiency / RNNs and normalization
Compute Better Spent: Replacing Dense Layers with Structured Matrices
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Systematically studies structured matrices (low-rank, Kronecker, Monarch, TT, and a new Block Tensor-Train family) as drop-in replacements for dense layers, using MUP-derived structure-aware learning rate and initialization scaling to compare their compute scaling laws on MLPs, ViTs, and GPT-2.
Problem
Dense linear layers dominate compute in foundation models, but structured alternatives often underperform dense layers because they need different initialization and learning rates, and their compute-scaling behavior is unclear.
Approach
Derive structure-aware initialization and Adam learning-rate multipliers for arbitrary structured layers by treating them as deep linear networks of dense components under muP, and propose Block Tensor-Train (BTT) matrices, which remove parameter sharing from TT so parameters equal FLOPs and scaling laws improve over dense.
Key results
- BTT achieves exponentially lower training loss than dense when training MLPs and ViTs on CIFAR-10/100 with augmentation, with BTT having the largest scaling exponent among structures.
- BTT matches dense ViT-S/32 performance on ImageNet-1k with up to 3.8x fewer FLOPs.
- Structure-aware learning rate scaling yields a stable optimal learning rate across structures and widths, whereas naive dense-derived scaling does not.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace dense layers with BTT | MLPs and ViTs · CIFAR-10/100 with augmentation | training loss | Helped |
| Replace dense layers with BTT | ViT-S/32 · ImageNet-1k | compute at matched performance: 3.8 times less compute | Helped |
| Replace dense layers with BTT, including the language modeling head | GPT-2 · OpenWebText | validation loss as a function of compute | Helped |
| Replace dense layers with BTT, excluding language-modeling-head compute | GPT-2 · OpenWebText | validation loss | No clear effect |
| Use structure-aware learning-rate scaling instead of dense-derived naive scaling (ablation) | structured layers · CIFAR-10 | training error | Helped |
| Use naive dense-derived learning-rate scaling instead of structure-aware scaling (ablation) | structured layers · CIFAR-10 | feature-update RMS | Hurt |
| Add weight normalization to BTT cores (ablation) | GPT-2 with BTT · OpenWebText | validation loss | Helped |
| Add weight normalization to BTT cores (ablation) | GPT-2 with BTT · OpenWebText | activation stability | Helped |
| Remove weight normalization (ablation) | GPT-2 with BTT · OpenWebText | activation stability | Hurt |
| Increase Monarch blocks from 4 to 16 (ablation) | ViT · ImageNet-1k | compute efficiency | Helped |
| Increase BTT rank (ablation) | BTT with 2 cores · CIFAR-10 | training error per FLOP | Hurt |
| Increase BTT core count beyond 2 (ablation) | BTT · CIFAR-10 | performance controlling for FLOPs | No clear effect |
Surprising findings
- Kronecker and TT, which share parameters across block dimensions, underperform dense matrices despite being cheaper, while structures with parameters equal to FLOPs scale better.
- Without weight normalization, BTT-based transformers show unbounded growth of activations that eventually leads to NaN; lowering the learning rate or raising weight decay only slows this at the cost of performance.
Limitations
- Wall-clock runtime overhead for small matrices due to current implementation; FLOPs used as proxy.
- Experiments limited to CIFAR, ImageNet and small GPT-2 models; larger-scale validation left to future work.
- Heuristic delta_i = 1/k for structure-aware learning rate constant is not theoretically determined.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 ImageNet-1k OpenWebText
Benchmarks: CIFAR-10 CIFAR-100 ImageNet-1k OpenWebText
Metrics: training error test error validation loss scaling exponent alpha_C FLOPs
Baselines: Dense Low-rank Kronecker Monarch Tensor-Train GPT-2 Small
Builds on: Maximal Update Parameterization Monarch matrices Tensor-Train decomposition Kronecker product structure CoLA
Closest papers by meaning
- Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices (NeurIPS 2024 · similarity 0.90)
- BLAST: Block-Level Adaptive Structured Matrices for Efficient Deep Neural Network Inference (NeurIPS 2024 · similarity 0.89)
- Differentiable Learning of Generalized Structured Matrices for Efficient Deep Neural Networks (ICLR 2024 · similarity 0.86)
- Geometry-aware training of factorized layers in tensor Tucker format (NeurIPS 2024 · similarity 0.85)
- ModHiFi: Identifying High Fidelity predictive components for Model Modification (NeurIPS 2025 · similarity 0.84)
- SS1: Accelerating Inference with Fast and Expressive Sketch Structured Transform (NeurIPS 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.