Learning Functional Subspaces for Neural Network Compression

AI summary

Overview

Research area: Model compression for Transformers — specifically low-rank weight factorization of large language models and vision transformers.

Technical level: Intermediate. The core idea is intuitive (learn which directions to delete), but the paper assumes familiarity with singular value decomposition, orthogonal projectors, and the Transformer attention/MLP structure.

Scope: The paper proposes Learnable Subspace Projections (LSP), a method that learns end-to-end which subspaces to discard from each linear layer of a frozen pretrained network, and evaluates it on four LLMs (OPT-125M, OPT-1.3B, Qwen3-4B, Llama-2-7B) and ViT-B/16 against pruning, activation-based, reconstruction-based, and loss-aware baselines.

What This Paper Is About

Low-rank compression replaces a weight matrix with two smaller matrices, but existing methods decide which directions to throw away using local criteria — activation energy, layer-wise reconstruction error, or a quadratic approximation of the loss — that never look at how the error propagates through the rest of the network. Because these criteria ignore downstream effects, errors compound with depth and performance collapses at high compression ratios. LSP instead makes the subspace itself a learned parameter, optimizing all projectors jointly against a global objective (the KL divergence to the dense model's outputs, or the original training loss) while keeping the pretrained weights frozen.

Key Contributions

1. Learned subspace projections with tied groups. Each targeted linear layer is composed with an orthogonal projector P = I − UUᵀ, and layers that read the same activations (Q/K/V, gate/up) share one projector. All projectors are optimized jointly across the network against a global objective, with pretrained weights frozen. An activation-space parameterization with fused QR orthonormalization makes this tractable at billion-parameter scale (the paper reports O(d_in·k + k²) per layer versus O(d_in·k²) for an explicit Gram–Schmidt loop).

2. Whitened initialization and measured-KL rank allocation. Projectors are initialized from a whitened SVD truncation, recast as an orthogonal projector and extended to tied groups. The number of directions each unit removes is decided by applying each projector in isolation, measuring the KL divergence it induces on the model output, and allocating budget by KL cost per saved parameter. The untrained model at this initialization — called NoLSP — serves as a control isolating the effect of learning.

3. Efficient low-rank inference with no deployment overhead. After training, projectors merge exactly into plain low-rank factors W P = (W U⟂) U⟂ᵀ = BA, so the deployed model contains no LSP-specific operations. Because each input-tied group shares the input factor A, attention can cache one narrow latent z = Ax instead of separate keys and values.

4. A label-free objective that exploits unlabeled calibration data. The default distillation objective needs no labels, so any unlabeled pool can be used, and the paper shows this improves transfer to unseen distributions relative to baselines.

Main Findings

• LSP wins on every LLM perplexity setting, and the margin grows with compression. On WikiText-2, one of the two LSP variants has the lowest perplexity in all 12 model–ratio settings across −30/50/70% compression. At −70%, LSP brings Llama-2-7B to 10.9 perplexity and 42.2% mean zero-shot accuracy, versus 13.3 and 36.0% for the strongest baseline.

• Dense baselines degrade far more steeply. At −70% on Llama-2-7B, SVD-LLM reaches 13.3 perplexity and NoLSP (the untrained initialization at the same ranks) reaches 222.8; learning the directions is what lowers it to 10.9. ASVD exceeds 10⁴ perplexity on Llama-2-7B at −50% and >10⁵ at −70%. At −70% the training-free SVD methods (SliceGPT 86.3, SVD-LLM(W) 214.9, Swift-SVD 203.6) deteriorate sharply.

• Task loss helps at low compression and on smaller models; distillation helps at high compression and on larger models. LSP_T can fall below dense test perplexity at −30% on both OPT models (24.5 and 12.9 versus dense 27.9 and 14.7), which the paper attributes to specialization to the WikiText-2 calibration domain.

• Zero-shot accuracy margins are consistent. Across six benchmarks — OpenBookQA, ARC-easy, WinoGrande, HellaSwag, PIQA and MathQA — both LSP variants lead the strongest training-free baseline in mean accuracy at every ratio, by 6.7–10.5 points on Llama-2-7B (Alpaca calibration) and 6.2–10.8 points on Qwen3-4B (C4 calibration). The two objectives are within two points of each other, with distillation ahead or tied in every setting.

• On ViT-B/16, LSP is most robust to calibration shift. With in-domain (single CIFAR-100) calibration at −70%, LSP scores 85.9 source accuracy against the dense model's 89.9, where the strongest training-free baseline falls to 77.1. Moving to a diverse calibration pool, LSP has the largest transfer gain at every ratio (+3.7, +7.1, +6.6 at −30/50/70%), and reaches 84.6 source / 39.6 transfer at −70%.

• Fewer, larger matrix multiplies make LSP the fastest factorized model measured. On Llama-2-7B it decodes faster than the dense model at every ratio: 1.21×, 1.36×, 1.56× at −30, −50, −70%, whereas every untied factorization is slower than dense at −30%. Decode throughput at 512 tokens jumps from 141 tok/s (dense) to 221 tok/s (LSP); at 16k context from 36 to 40 tok/s.

• The latent cache shrinks KV memory substantially. At −70%, LSP stores 556 KV floats per token per layer against 8192 for dense and 1228 for SVD-LLM/Swift-SVD — a 14.7× reduction versus dense and 2.2–2.6× versus untied factorizations. Weights plus KV cache at 128k tokens and batch 8 take 38.9 GiB versus 524.6 GiB for dense and 80.9 GiB for SVD-LLM/Swift-SVD, a 13.5× reduction versus at most 6.5× for untied baselines. A 95.5 GiB GPU holds about 320k cached tokens per sequence, against at most 146k.

• Allocation is strongly non-uniform. Q/K/V groups are compressed far more than MLP projections, even though one direction removed from gate/up saves more parameters (26,112) than one removed from Q/K/V (16,384). At −70%, the shared K/V latent is about a seventh of full rank. Later blocks retain more rank than earlier ones at every ratio, and at −30%, 53 of 128 units stay dense because their allocated rank exceeds the break-even rank at which two factors are smaller than the dense weight.

• Learning corrects the weight spectrum, most strongly along the leading direction. Removal has no sharp spectral cutoff; it spans the whole spectrum, deepest on the leading directions in Q/K/V. Relative to NoLSP at identical allocated ranks, LSP removes more of its high singular directions in exchange for small ones, and direction 1 changes most in absolute terms.

Methodology in Plain English

The method rests on one observation: deleting k directions from a weight matrix is equivalent to multiplying it by an orthogonal projector P = I − UUᵀ. So instead of asking "which directions have low energy?", the authors make the projector a learnable parameter and let the network's own output loss decide.

Concretely: for each targeted layer (or group of layers reading the same activations), they parameterize an unconstrained matrix V, recover an orthonormal U via a QR decomposition on every forward pass, and apply the projector to the frozen pretrained weight. All the V matrices across the network are then trained jointly against either the KL divergence to the dense model's output distribution or the model's original training loss. The dense model itself acts as the teacher (with projectors disabled), so no second copy of the weights is needed.

Three practical devices make this stable: computing in activation space (W(x − U(Uᵀx))) instead of materializing WP; a warm-up parameter α that ramps from 0 to 1 and a random mask that keeps some "removed" directions temporarily, acting as dropout; and an orthogonality penalty on V's columns.

Before training, ranks and starting subspaces are chosen without gradients. A whitened SVD truncation provides an ordered basis, so deciding how many directions to drop is just choosing how many trailing vectors to discard. The number itself comes from measurement: each unit's projector is applied in isolation, the KL divergence it causes on the model output is measured on a grid of removal fractions, and budget is allocated to whichever unit offers the smallest marginal KL cost per parameter saved. After training, projectors fold back into ordinary low-rank factors.

Why This Matters

Impact on research. The paper reframes low-rank compression from a spectral question ("which directions carry little energy?") to a functional one ("which directions does the output depend on least?"). It demonstrates that the weight spectrum ranks directions poorly — whitened truncation discards trailing directions the output depends on — and that end-to-end learning corrects this. It also shows that a label-free KL objective opens a new axis for compression: using unlabeled data of diverse composition to better preserve the dense model's general abilities.

Real-world applications:

• Serving large language models on constrained hardware. At −70%, Llama-2-7B fits with far longer contexts on a single GPU — about 320k cached tokens versus at most 146k for untied factorizations.
• Long-context inference. The shared latent cache targets the KV cache, which the paper identifies as the memory bottleneck that grows with context length.
• On-device or edge deployment at small batch sizes, where decoding is memory-bound and fewer, larger matrix multiplies translate into higher throughput — LSP decodes faster than the dense model at every measured ratio.
• Compression with unlabeled data. Because the default objective needs no labels, models can be compressed on whatever text or image pool is available, which the ViT experiments show also improves transfer to unseen datasets.

Industry relevance. The deployed artifact is a standard low-rank factorization with no LSP-specific operations, matching other factorizations on weights and FLOPs per token — so adoption does not require novel inference kernels except for the optional latent cache. The measured gains are in the two quantities that dominate serving costs: GPU memory and decode latency.

Future Directions

• Extending the latent cache to a full deployment path. The paper computes cache size and context capacity for the latent cache but measures throughput with a standard full-width KV cache, and notes the latent cache serves capacity rather than speed because keys must be rebuilt under rotary embeddings. A dedicated attention implementation is left as future work.

• Jointly allocating rank and learning subspaces. Ranks are fixed before training begins; the paper notes that joint training handles the combined effect of independently measured units, but whether ranks should be revisited during or after training is left open.

• Broader architectures and modalities. Only decoder-only LLMs and ViT-B/16 are evaluated. The paper claims the recipe transfers to any linear layer without architecture-specific surgery, but this is not tested beyond those model families.

• Understanding why the leading singular direction changes most and whether that pattern generalizes. The analysis shows learning shifts removal toward high singular directions and away from trailing ones, most strongly along direction 1, but the mechanism behind this redistribution is not explained.

Target Audience

Researchers and engineers working on model compression, efficient inference, and Transformer architecture. It is most useful to those already familiar with low-rank factorization methods (SVD-LLM, ASVD, SliceGPT, LLM-Surgeon) who want to understand what a global, learned objective buys over local criteria, and to practitioners deciding how to fit large models into limited GPU memory or long-context serving budgets. Readers without a background in linear algebra will find the mathematical formulation heavy, though the central intuition — learn which directions to delete by asking the network what it depends on — is stated plainly.

Authors’ abstract

Modern transformers pair impressive capabilities with substantial memory and compute demands. Low-rank weight factorization reduces both while keeping the matrices dense, and thus efficient on standard hardware. Existing methods, however, choose the subspace to remove from each weight matrix with local closed-form criteria: activation energy, layer-wise reconstruction error, or a quadratic approximation of the loss. These criteria ignore how errors propagate through the network, so at high compression the errors compound with depth and performance collapses. We introduce Learnable Subspace Projections (LSP), which instead learns the subspaces to discard end-to-end. Each linear layer, or tied group of layers that read the same activations, is assigned an orthogonal projector. All projectors are optimized jointly against a global objective--the KL divergence to the dense model's output distribution or the model's original training loss--while the pretrained weights remain frozen. Projectors are initialized from a whitened SVD truncation, and ranks are allocated by the output KL each projector induces per parameter saved. After training, the projectors merge into standard low-rank factors, with each tied group sharing one factor. In attention, this also lets the model cache one narrow latent in place of full keys and values. Across LLMs (OPT-125M/1.3B, Qwen3-4B, Llama-2-7B) and ViT-B/16, LSP outperforms baselines, and its advantage widens as compression increases. At -70% compression, LSP brings Llama-2-7B to 10.9 WikiText-2 perplexity and 42.2% mean zero-shot accuracy, versus 13.3 and 36.0% for the strongest baseline. The factorized model decodes up to 1.6x faster than the dense model at small batch sizes, and aching the shared latent shrinks the combined memory of weights and KV cache by 13.5x at a 128k-token context, versus at most 6.5x for untied baseline factorizations.