Low-rank compression stops spending the same compute per token: training only routers lifts Llama-2-7B downstream accuracy by 7.6 points
Synopsis
Low-rank compression previously assigned a fixed rank during inference, so every token cost the same compute. LRCC adds one lightweight router per Transformer block that selects among a few nested low-rank paths for each token, training only the routers while the low-rank factors stay frozen. Within the same average active-parameter budget, Llama-2-7B gains 7.6 percentage points in average downstream accuracy over static low-rank compression; at matched batch-size-1 decoding latency, Llama-3.2-1B improves both perplexity and downstream accuracy and Llama-2-7B stays competitive, without specialized CUDA kernels.
Interpretation
A low-rank compressed model can now decide per token how many ranks to activate, instead of using one rank allocation for the whole inference run. Conventional low-rank compression uses a fixed rank allocation during inference, assigning the same amount of compute regardless of the input token. Within the same average active-parameter budget, Llama-2-7B improves average downstream accuracy by 7.6 percentage points over the strongest static baseline, with WikiText-2 perplexity improving as well.
The construction factorizes each linear layer inside a Transformer block into shared nested low-rank factors, then lets one block-level router pick a path per token. Paths are ordered and nested in the Matryoshka style, so lower-cost paths retain subsets of the components used by higher-cost paths, all paths share one parameter base, and a single path index selects among them. Routers are trained with Gumbel-Softmax and a straight-through estimator while pretrained weights and low-rank factors stay frozen; experiments cover Llama-3.2-1B, Llama-2-7B and Qwen3-4B, using 256 WikiText-2 samples of 2,048 tokens for decomposition and router training.
The routers do allocate different compute to different tokens rather than collapsing to a single static rank. If every token were assigned the same path, conditional computation would reduce to static compression. On Llama-2-7B some blocks use almost all available paths and non-word tokens are more often assigned to lower-cost paths; routers trained only on WikiText-2 keep actual activated-parameter ratios close to their targets across several evaluation datasets.
Perspective
The demonstrated scope is language modeling and zero-shot downstream tasks on Llama-3.2-1B, Llama-2-7B and Qwen3-4B, compressing linear layers inside Transformer blocks, with routers trained on 256 WikiText-2 samples of 2,048 tokens. Latency results are limited to batch-size-1 autoregressive decoding, compiled with TorchDynamo and TorchInductor on a DGX Spark without specialized CUDA kernels. Intended for researchers and engineers working on model compression and inference efficiency.
Routers are trained toward an average activated-parameter ratio that is not enforced at inference, so actual deployment cost varies with the dataset, though it stays close to target in the evaluations. Paths are built before router training from a fixed set of nested rank allocations, so path structure and routing policy are not learned jointly. Latency gains are measured only for batch-size-1 decoding, and the prefill implementation still computes the first factor at full rank before masking, so it does not yet exploit all available savings.