Research map / Architectures and efficiency / RNNs and normalization
Searching for Efficient Linear Layers over a Continuous Space of 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.
Searches a continuous space of Einsum-expressible structured linear layers, characterizes them by parameter sharing, rank, and compute intensity, and proposes BTT-MoE, a per-layer sparse mixture of full-rank BTT experts that gives ~5x compute savings over dense GPT-2.
Problem
Dense linear layers dominate compute in large models, but prior structured-matrix alternatives were shown to help only in multi-epoch or fixed-model-size settings, not under compute-optimal scaling laws.
Approach
Parameterize all two-factor Einsum linear operators by a 7-dim continuous vector theta, derive three taxonomy exponents (omega for parameter sharing, psi for rank, nu for compute intensity), and show that omega=0 and psi=1 (full-rank BTT) govern compute-optimal scaling; then sparsify the BTT rank sum into a per-layer MoE.
Key results
- BTT-MoE with 16 experts achieves a compute multiplier of 5.3 (±0.3) versus 4.1 (±0.3) for standard MoE.
- Full-rank structures without parameter sharing (omega=0, psi=1) scale similarly to dense GPT-2 in compute-optimal training, while larger omega or smaller psi lead to worse scaling.
- Exploiting attention head structure improves compute efficiency by an average of 17% over naive BTT replacement.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace dense linear layers with BTT-MoE using 16 experts | GPT-2 · GPT-2 language modeling | compute multiplier: 5.3_{\pm0.3} | Helped |
| Use BTT-MoE rather than standard MoE over entire FFNs | GPT-2 · GPT-2 language modeling | compute multiplier: 5.3_{\pm0.3} compared to 4.1_{\pm0.3} | Helped |
| Use attention-head-aware BTT axes | BTT layers in GPT-2 · GPT-2 language modeling | compute efficiency: 17\% | Helped |
| Use low-rank matrices as the sparse MoE experts instead of BTT (ablation) | Structured MoE · GPT-2 language modeling | compute multiplier: \sim2\times | Hurt |
| Use dense matrices as the sparse MoE experts instead of BTT (ablation) | Structured MoE · GPT-2 language modeling | compute multiplier: \sim2\times | Hurt |
| Use Einsums with parameter sharing (larger ω) | GPT-2 with structured linear layers · OpenWebText | compute-optimal scaling performance | Hurt |
| Use full-rank Einsums instead of low-rank Einsums | GPT-2 with structured linear layers · OpenWebText | compute-optimal scaling performance | Helped |
| Vary compute intensity ν among full-rank, non-sharing structures | GPT-2 with structured linear layers · OpenWebText | compute-optimal scaling performance | No clear effect |
| Use BTT with ν=0.75 rather than Monarch with ν=0.5 | Autoregressive pixel modeling · CIFAR-5M | compute-optimal scaling performance | Helped |
| Use μP initialization and width/structure-scaled learning rates instead of unit-variance initialization and a global learning rate | Einsum layers · CIFAR-10 | test error | Helped |
Surprising findings
- Varying the compute intensity nu has minimal effect on scaling laws among full-rank no-sharing structures, so dense-like structures do not win because they are dense.
- Most Einsums outperform dense at small scales on CIFAR-5M, though the authors hypothesize this is due to larger embedding dimensions.
Limitations
- Scaling experiments use small GPT-2 models (up to 76M parameters) with a reduced 96-token vocabulary and sequence length 128; the original-vocabulary setup shows less clean power-law behavior.
- Existing structured matrices (low-rank, TT, Monarch, BTT) do not significantly outperform dense under compute-optimal training.
- Theory of the taxonomy and mu-P scaling for Einsums relies on asymptotic arguments that may not hold at finite width.
Experimental setup
Datasets: OpenWebText CIFAR-5M
Benchmarks: OpenWebText language modeling loss CIFAR-5M pixel modeling synthetic MLP regression
Metrics: cross-entropy loss compute multiplier MSE
Baselines: dense GPT-2 standard MoE low-rank-MoE dense-MoE Monarch Block Tensor-Train
Builds on: Block Tensor-Train Monarch muP Tensor-Train load-balancing MoE
Closest papers by meaning
- Compute Better Spent: Replacing Dense Layers with Structured Matrices (ICML 2024 · similarity 0.90)
- Composing Linear Layers from Irreducibles (NeurIPS 2025 · similarity 0.86)
- BLAST: Block-Level Adaptive Structured Matrices for Efficient Deep Neural Network Inference (NeurIPS 2024 · similarity 0.86)
- Geometry-aware training of factorized layers in tensor Tucker format (NeurIPS 2024 · similarity 0.86)
- Differentiable Learning of Generalized Structured Matrices for Efficient Deep Neural Networks (ICLR 2024 · similarity 0.86)
- Compress to Impress: Efficient LLM Adaptation Using a Single Gradient Step on 100 Samples (NeurIPS 2025 · similarity 0.85)
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.