Compute Better Spent: Replacing Dense Layers with Structured Matrices

TL;DR

Qiu et al. systematically study whether structured matrices (low-rank, Kronecker, Monarch, Tensor-Train, and a newly introduced Block Tensor-Train / BTT family) can replace dense linear layers and achieve better scaling exponents on training loss vs compute. The main methodological contribution is a µP-style, structure-aware recipe for initialization and Adam learning rates that treats each dense sub-component of a structured matrix as its own µP layer; without this, structured layers systematically underperform. The empirical headline is that structures where parameter count equals FLOPs (dense, low-rank, Monarch, BTT) achieve clean power-law scaling, while parameter-sharing structures (Kronecker, TT) scale strictly worse. BTT improves the scaling exponent over dense on CIFAR-10/100 MLPs/ViTs and matches dense ViT-S/32 on ImageNet-1k with 3.8× less compute; small GPT-2 BTT models also beat dense at matched compute, but only after a weight-normalization fix is applied to prevent slow activation growth from diverging training.

Key claims

Naively using dense-layer learning rates for structured layers causes significant underperformance; structured layers often need learning rate multipliers that grow with width [§3.3, Fig. 1(c), Fig. 2].
The structure-aware recipe sets each dense sub-component’s init std and Adam LR using µP applied to that component’s shape alone, with a cross-structure constant $c_s = 1/k$ (k = number of learnable dense components) to preserve overall layer-update magnitude [§3.2, Eq. 3, Table 2].
With the structure-aware recipe, feature-update RMS stays stable across widths and structures, and the LR landscape is approximately preserved so dense-tuned base LRs transfer to other structures [§3.3, Fig. 2, Fig. 2(b)].
Scaling exponents of training error vs compute are structure-dependent, not purely task-determined — contradicting the data-manifold-dimension prediction that architecture only shifts the prefactor [§4, Fig. 1(b), Fig. 4].
Parameter-sharing structures (Kronecker, TT) scale worse than non-sharing structures (dense, low-rank, Monarch, BTT) at matched FLOPs; the “parameters = FLOPs” condition is a shared feature of the best-scaling structures [§4, Fig. 1(b), Fig. 4].
BTT with 2 cores and rank 1 (equivalent to Monarch with √d blocks) achieves the largest fitted scaling exponent and consistently outperforms all other structures on CIFAR-10/100 MLPs and ViTs [§4, Fig. 1(b), Fig. 4].
Lower-rank BTT and higher-block-count Monarch (lower compute-per-dimension) are more compute-efficient than dense but less activation-memory-efficient; the optimal compute-per-dimension $s$ is well below the dense $s=d$ and is task-dependent [§4, Fig. 5, Fig. 6].
On ImageNet-1k, BTT matches dense ViT-S/32 performance with 3.8× less compute [Abstract, §5].
Training small GPT-2 with BTT layers suffers slow unbounded activation growth that eventually diverges; a weight-normalization reparameterization (BTT cores clamped to initialization RMS × learnable scalar, Salimans & Kingma 2016 style) eliminates the divergence without sacrificing expressivity [§5.1, Fig. 7].

Method

Each structured matrix is decomposed as $M \cdot B_k \cdots M \cdot B_1 \cdot M$ , where the $M$ are fixed norm-preserving transforms (reshapes, permutations) and the $B_i$ are learnable dense tensors with batch-like dimensions. Treat each $B_i$ as a µP dense layer in its own shape: initialize entries with std $\sigma_i = 1/\sqrt{\text{fan\_in}_i}$ and set Adam LR as $\eta_i = \eta_{\text{base}} \cdot c_s / \text{fan\_in}_i$ , with the shared constant $c_s = 1/k$ chosen so the leading-order sum of per-component contributions to the layer-output update stays $O(1)$ (Eq. 3). The result is a width-invariant LR landscape where the dense-tuned base LR transfers to BTT, Monarch, Kronecker, TT, and low-rank without additional tuning.

BTT itself is defined as a two-core generalization of Monarch where the parameter-sharing along block dimensions in the Tensor-Train decomposition is removed: $(M \cdot x)_{ij} = \sum_k L_{i,j,k} \cdot R_{k,j,i} \cdot x$ for input/output reshaped as matrices (Eq. 2). BTT with rank $r$ interpolates between Monarch (r=1) and dense ( $r \geq \sqrt{d}$ ), with intermediate ranks parameterizing a family of structured matrices that are strictly more expressive than Monarch at the same FLOP budget per dimension.

For GPT-2 training with BTT, a weight-normalization reparameterization $\tilde L = (s_L / \text{RMS}(L)) \cdot L$ (and symmetric for $R$ ) with scalar $s_L, s_R$ learnable eliminates the slow activation drift observed in the plain parameterization.

Results

CIFAR-10/100 MLP and ViT training error vs compute follow approximate power laws for each structure, with BTT having the largest fitted exponent and parameter-sharing structures (Kronecker, TT) the smallest [Fig. 1(b), Fig. 4].
BTT matches dense ViT-S/32 ImageNet-1k performance at 3.8× less compute [Abstract, §5].
Small GPT-2 trained on OpenWebText: BTT layers outperform dense at matched compute, but only after the weight-normalization fix; without it, pre-LN activation RMS grows without bound and training eventually diverges [Fig. 7].
Structure-aware LR scaling is necessary for optimal performance in heterogeneous architectures (e.g., ViT with BTT FFN + dense attention projections): the naive single-base-LR recipe forces a bad compromise even after grid-searching the base LR [Fig. 3].
On the compute–memory tradeoff: lowering compute-per-dimension (smaller BTT rank, more Monarch blocks) buys compute efficiency at the cost of activation memory, since activation storage scales with width $d$ which grows as compute-per-dimension shrinks [Fig. 6].

Why it’s interesting

This is a cleanly-run µP-style ablation on the architecture axis orthogonal to the µP work the wiki already tracks: where Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning extends µP into the LoRA-rank joint limit and Controlled LLM Training on Spectral Sphere closes the loop by constraining weight and update spectra, this paper extends µP into the structure of the weight matrix itself and finds that the structure-aware LR prescription is load-bearing for even measuring the scaling exponent correctly. It also connects directly to the ongoing MoE / Monarch / sparse-attention line on Hyperparameter scaling laws and to MonarchRT: Efficient Attention for Real-Time Video Generation‘s use of Monarch factorization in video DiT attention — raising the question of whether MonarchRT would see a larger gain if trained with the structure-aware LR recipe here rather than reusing dense-tuned LRs. The main caveat (and the point Kyle flagged on filing) is scale: the clean scaling-exponent-advantage result lives in CIFAR MLPs/ViTs; the ImageNet and GPT-2 results are at modest scale and the GPT-2 case already required a stability patch to work at all.