LRCC: Generalizing Low-Rank Compression with Conditional Computation

Venue

AXIOM Workshop @ The Fortieth Annual Conference on Neural Information Processing Systems

Introduction

Low-rank compression reduces the cost of pretrained language models by replacing linear transformations with products of smaller matrices. However, conventional methods use a fixed rank allocation during inference, assigning the same amount of compute regardless of the input token.

Conditional computation offers a different possibility: adapt the computational graph to the input. Applying it to pretrained models can require additional training and compute modules. Yet low-rank decompositions already provide a useful structure for this purpose: smaller models can share a subset of the components used by larger ones.

We introduce Low-Rank Conditional Computation (LRCC), which trains one lightweight router per Transformer block to select among a small set of nested low-rank paths. The low-rank factors remain frozen; only the routers are optimized.

Our study asks two questions: can token-dependent rank allocation improve predictive performance within the same active-parameter budget, and do these gains translate into better quality at matched decoding latency? We distinguish these questions because activating fewer parameters does not by itself guarantee lower latency.

From a fixed rank to token-dependent paths

Conventional low-rank compression retains a fixed number of components in each linear layer. Layer-specific ranks are chosen to meet a global budget, using a heuristic or learned allocation. Every token then uses that same allocation.

LRCC builds on this structure in three steps.

First, decompose the linear layers. We use input activations collected on a calibration dataset to find low-rank approximations that minimize reconstruction error in each layer's outputs. These activation-aware decompositions are nested: a lower-rank approximation uses a prefix of the components retained by a higher-rank one. Retaining more components cannot increase the reconstruction error on the calibration activations.

Second, construct a small set of paths. Routing over every possible rank in every layer would introduce many decisions. A Llama-2-7B linear layer can have up to 4,096 rank choices, and each Transformer block contains seven such layers. Instead, we use the Swift-SVD rank-allocation procedure to construct a small set of submodels at different costs.

Within each block, a path specifies the ranks activated in all its linear layers, covering both attention and MLP projections. Lower-cost paths retain subsets of the components used by higher-cost paths, so all paths share the same parameter base.

Third, train a block-level router. For each token entering a block, a learnable linear router selects one path. That single choice determines the ranks used throughout the block. The choice is made separately at each block, allowing different tokens to activate different amounts of computation.

Training only the routers

During training, all pretrained model parameters and low-rank factors remain frozen. We optimize the routers using the standard causal language-modeling loss together with a compression loss that penalizes exceeding an average target cost.

The cost is expressed as an active-parameter ratio: the parameters activated in the Transformer blocks' linear layers relative to their original dense weights. Embeddings and the language-model output head are excluded. This measures the parameter usage of the selected computation, rather than the total memory occupied by the model.

We also introduce target-cost annealing. Training begins at a high target cost and progressively moves towards lower costs, saving a router checkpoint for each target. An adaptive weight strengthens the compression penalty when the current cost exceeds the target. This produces a family of router checkpoints in one progressive training run.

We use 256 WikiText-2 training sequences of 2,048 tokens for both decomposition and router training. Before deployment, we estimate each block's average cost on validation samples. If a factorized block would be more expensive than its original dense implementation, we retain it in dense form.

Better quality within an active-parameter budget

We evaluate LRCC on Llama-2-7B, Llama-3.2-1B, and Qwen3-4B, measuring WikiText-2 perplexity and zero-shot accuracy on seven downstream tasks: ARC-Easy, ARC-Challenge, BoolQ, HellaSwag, OpenBookQA, PIQA, and WinoGrande.

Our static baselines include COALA, SVD-LLM, Swift-SVD, and combinations that use Swift-SVD's rank allocation. No method fine-tunes the pretrained weights after compression; LRCC additionally trains its routers. SVD-LLM's optional weight fine-tuning stage is omitted.

The table below summarizes selected Llama-2-7B results. LRCC-5 and LRCC-9 use five and nine paths per block, respectively.

Average accuracy is the mean over the seven downstream tasks; uncertainty terms are as reported in the paper. Results are within the indicated budgets. When a router exceeded a budget on a task, we used a lower-target checkpoint.

At the 90% budget, LRCC-9 improves average accuracy from 55.1% for the best static baseline to 62.7%, a gain of 7.6 percentage points. At 70%, LRCC-5 improves it from 45.4% to 49.2%. Both variants also improve perplexity at these budgets. For reference, the uncompressed model reaches 66.6% average accuracy and 5.49 perplexity.

Across a wider range of costs, LRCC improves the quality–cost trade-off in almost all plotted settings, with particularly large gains on downstream tasks. The Qwen3-4B experiments extend the quality gains to another model family.

We hypothesize that aggressive static truncation removes subspaces important for particular tasks. LRCC can retain such components for selected tokens while meeting an average budget, providing a possible explanation for these gains.

Do the gains survive at matched latency?

The routers and path selection introduce overhead, so improvements in active-parameter efficiency must also be evaluated through measured runtime.

We measure batch-size-1 autoregressive decoding on an NVIDIA GB10 Grace Blackwell, compiling LRCC and COALA+Swift with identical TorchDynamo and TorchInductor settings. Both models process the same ground-truth token sequences, with five repetitions per measurement. Compilation and warm-up are excluded from the timed region, and no specialized CUDA kernels are used.

At matched latency, LRCC improves perplexity and average downstream accuracy on Llama-3.2-1B. For Llama-2-7B, it improves downstream performance at higher latencies, while the WikiText-2 results are mixed.

At the same nominal target cost, LRCC is slower than COALA+Swift on Llama-2-7B. Its practical benefit is therefore a competitive quality–latency trade-off, with the outcome depending on the model and workload.

A small set of paths is sufficient

LRCC needs enough routing choices to allocate computation flexibly, but more paths also increase implementation costs.

We compare models with 2, 3, 5, 9, 10, 17, and 33 paths per block. With only two paths, perplexity increases sharply; three paths also remain noticeably worse, particularly under more aggressive compression. Increasing the number beyond five provides comparatively little improvement in most settings.

Compilation becomes substantially more expensive as the number of paths increases. Overall, five paths provide a good trade-off between predictive and computational performance, motivating their use in our main trade-off curves.

The routers retain token-dependent choices

A router could, in principle, learn to select the same path for every token. We therefore examine whether LRCC actually uses its available choices, measuring both the average cost and how broadly path selections are distributed within each block.

The analysis shows a backbone of blocks retained in dense form, alongside blocks that route tokens through different low-rank paths. As the target cost decreases, more blocks exhibit conditional routing. Some use almost all available paths, without an explicit path-balancing loss.

We also find that non-word tokens, defined as tokens containing no alphanumeric characters, are more frequently assigned to lower-cost paths. This provides evidence that compute allocation varies with token type.

Although the routers are trained only on WikiText-2, actual deployment costs remain close to their targets across the evaluated datasets. The target is an average learned during training, however, and is not strictly enforced at inference.

What remains open

LRCC shows that shared low-rank hierarchies provide a practical basis for adding conditional computation to pretrained Transformers while training only lightweight routers.

Several directions remain open. Our latency evidence concerns batch-size-1 decoding; optimized prefill and larger batches require further investigation. We have not evaluated combining routing with fine-tuned low-rank factors. Finally, the paths are fixed before router training: jointly learning their structure and the routing policy could further improve how computation is allocated across inputs.