PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Functional Subspaces for Neural Network Compression

Massimo Bini, Anders Christensen, Stephan Alaniz, Judah Goldfeder, Ole Winther, Yann LeCun, Ravid Shwartz-Ziv, Zeynep Akata

arXiv:2609.40127v1cs.LGcs.CL

TL;DR

Layer-local compression criteria choose directions without directly accounting for how the network output depends on them. LSP jointly learns projectors against a global objective with frozen pretrained weights; across decoder LLMs, it improves compression results, including at −70%.

  • Problem

    Existing compression methods choose subspaces using local activation, reconstruction, or loss-surrogate criteria rather than having network loss choose which directions to remove.

    Source: §1 · p. 2

  • Method

    LSP jointly learns projectors against dense-model output KL or the original task loss while pretrained weights remain frozen, allocating ranks using measured output-KL cost per saved parameter.

    Source: §1 · p. 2, §2.2 · p. 4

  • Results

    Across four decoder LLMs, one LSP variant has lowest perplexity in all 12 model–ratio settings; at −70%, Llama-2-7B reaches 10.9 perplexity versus 13.3 for the strongest baseline.

    Source: §5 · p. 10

  • Takeaways & Limitations

    Merged low-rank factors reduce model footprint, and shared input factors let attention cache one latent instead of separate full keys and values.

    Source: §1 · p. 2

    Compression loss may be uneven across subpopulations, domains, or rare-event behaviors under-represented in the calibration distribution.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Low-rank compression can reduce transformer memory and compute, but layer-local criteria may discard directions that matter to the network output. LSP instead learns which directions to remove against a global objective.

1 INTRODUCTION

  • Figure 1 contrasts energy-based removal with LSP’s output-sensitivity-based choice of directions to discard.Arrow length encodes spectral energy, orientation encodes output dependence, and faded arrows are removed.Source: Figure 1 · p. 2
  • Existing activation-, reconstruction-, and surrogate-loss criteria assess layers locally, while LSP searches for directions whose removal preserves network behavior.The introduction argues that local criteria do not account for downstream error propagation, particularly under aggressive compression.Source: §1 · p. 2
  • LSP jointly optimizes projectors against the dense model’s output KL or the original training objective while pretrained weights remain frozen.Layers reading the same activations share a projector; ranks are allocated by measured output KL per parameter saved.Source: §1 · p. 2, Figure 2 · p. 3

2 LEARNABLE SUBSPACE PROJECTIONS (LSP)

LSP compresses linear layers by learning orthogonal projectors, jointly optimized while pretrained weights remain frozen. It initializes and allocates ranks using calibration data, then merges trained projectors into low-rank factors.

2.1 LEARNING SUBSPACE PROJECTIONS

  • An orthogonal projector removes a learned subspace from a layer’s input or output, limiting the projected weight’s rank.For input-side projection, the unconstrained parameter is orthonormalized with QR; layers sharing activations can share one projector.Source: §2.1 · p. 3

2.2 WHITENED INITIALIZATION AND MEASURED-KL ALLOCATION

  • Whitened SVD truncation initializes removed directions, while measured output-KL cost per saved parameter allocates rank across units.Allocation measures candidate truncations through the full network with other layers dense; unlike a small-perturbation curvature estimate, it evaluates finite truncations.Source: §2.2 · p. 4

2.3 TRAINING PROCEDURE

  • At fixed ranks, LSP jointly trains projectors with pretrained parameters frozen, using dense-model output distillation by default or the task loss as an alternative.Distillation targets the dense model’s output distribution; task-loss training fits calibration data directly and can specialize when it matches deployment.Source: §2.3 · pp. 4–5
  • After training, projectors merge into standard low-rank factors, with input-tied layers sharing a factor that lets attention cache one latent for keys and values.The common latent replaces storing full keys and values in the KV cache.Source: §2.3 · p. 5

3 EXPERIMENTS

Across language and vision benchmarks, LSP generally outperforms compression baselines, with advantages that grow at higher compression. The experiments also examine inference efficiency and the subspaces LSP removes.

3 EXPERIMENTS

  • The evaluation covers OPT-125M/1.3B, Qwen3-4B, and Llama-2-7B language models, plus a ViT-B/16, with compression up to −70%.Language-model evaluation includes WikiText-2 perplexity and six zero-shot benchmarks; vision evaluation includes source accuracy and linear-probe transfer.Source: §3 · p. 5

3.1 LARGE LANGUAGE MODELS

  • One LSP variant has the lowest WikiText-2 perplexity in all 12 model–ratio settings, and both variants beat every baseline on all four models at −70%.The margin grows with compression; no loss-aware baseline closes the gap at −70%.Source: Table 1 · p. 5, §3.1 · p. 6
  • Both LSP variants lead the strongest training-free baseline in mean zero-shot accuracy by 6.7–10.5 points on Llama-2-7B and 6.2–10.8 points on Qwen3-4B.The comparison covers three compression ratios for each model; distillation is ahead or tied in every setting.Source: Table 2 · p. 6, §3.1 · p. 6

3.2 VISION TRANSFORMERS

  • At −70%, ViT-B/16 scores 85.9 source accuracy against the dense model’s 89.9, while the strongest training-free baseline scores 77.1.On transfer to three unseen datasets, LSP has the best mean accuracy on the diverse calibration pool at all three ratios.Source: §3.2 · p. 7, Table 3 · p. 7

3.3 INFERENCE EFFICIENCY

  • On Llama-2-7B, LSP decodes 1.21×, 1.36×, and 1.56× faster than the dense model at −30%, −50%, and −70%, respectively.It is the fastest factorized model in every measured model–ratio setting; untied factorizations are slower than dense at −30%.Source: §3.3 · p. 7
  • At −70% compression with 128k cached tokens and batch size 8, weights plus cache use 13.5× less memory than dense, versus at most 6.5× for untied factorizations.Tied attention caches one narrow latent rather than separate keys and values.Source: §3.3 · p. 8, Table 5 · p. 8

3.4 WHAT DOES LSP COMPRESS?

  • Measured-KL allocation is nonuniform: attention inputs are compressed more than MLP projections, while learned removal spans the spectrum rather than sharply truncating its tail.At −70%, the shared K/V latent is about one-seventh of full rank; relative to NoLSP, LSP generally removes more leading directions and retains more small singular directions.Source: §3.4 · pp. 8–9, Figure 3 · p. 9

4 RELATED WORK

Prior low-rank methods select directions or allocate ranks using local criteria, while LSP learns removed subspaces against network output and ties projectors across layers sharing an input.

Pruning

  • Structured pruning removes coordinates or changes block widths, whereas LSP removes a continuous subspace and realizes savings as dense low-rank factors.Global objectives with frozen weights have also been used for pruning masks and per-block widths.Source: §4 · p. 9

Low-rank compression

  • LSP optimizes removed subspaces against network output with pretrained weights frozen, rather than relying on local criteria or post-truncation weight fine-tuning.Earlier methods choose directions from activation statistics, layer-output reconstruction, or local loss surrogates; LSP applies global optimization to the removed subspace.Source: §4 · p. 9

Rank allocation and shared bases

  • LSP ties projectors among layers sharing inputs and allocates ranks per compression unit using measured output KL.Prior shared-basis and rank-allocation approaches retain closed-form bases.Source: §4 · p. 9

5 CONCLUSION

LSP learns removed subspaces against network output rather than treating activation or layer-reconstruction redundancy as functional redundancy. Across language and vision evaluations, it performs strongly as compression increases and also improves inference efficiency.

Conclusion

  • Across four decoder LLMs, one LSP objective has the lowest perplexity in all 12 model–ratio settings, while both variants achieve the best low-rank zero-shot accuracy.At −70% compression on Llama-2-7B, LSP reaches 10.9 perplexity versus 13.3 for the strongest baseline and 222.8 without learning.Source: §5 · p. 10
  • At −70% compression on Llama-2-7B, LSP reaches 10.9 perplexity versus 13.3 for the strongest baseline.The conclusion reports that the performance margin grows with compression ratio.Source: §5 · p. 10
  • On ViT-B/16, LSP is most robust to calibration shift and transfers best when calibrated on diverse unlabeled data.The conclusion also reports that its advantage on transfer is largest on the diverse pool.Source: §5 · p. 10
  • The merged factorization decodes up to 1.56× faster than the dense model, and its shared latent cuts combined weights-and-KV-cache memory by 13.5× at 128k tokens.The memory reduction is reported as twice the savings of untied factorizations.Source: §5 · p. 10

AI DISCLOSURE STATEMENT

The authors disclose using generative AI tools for several research and writing tasks, with stated checks on the resulting work.

AI DISCLOSURE STATEMENT

  • Generative AI assisted with methodology feedback, implementation, mathematical claims, writing, analysis, and interpretation; authors report checking the relevant outputs.They state that they did not use these tools to develop theoretical models, translate, reformat datasets, or generate synthetic data.Source: AI DISCLOSURE STATEMENT · p. 10

ETHICS STATEMENT

LSP’s ethical discussion notes that compression may affect underrepresented behaviors unevenly and recommends evaluating compressed models against dense-model fairness, robustness, and safety metrics. The accompanying technical material justifies the projector and describes its initialization, allocation, training, and export.

ETHICS STATEMENT

  • Compression loss may be uneven across underrepresented subpopulations, domains, or rare events, so practitioners should evaluate fairness, robustness, and safety before deployment.The paper notes that aggregate degradation is small in its experiments, but calibration data may underrepresent these cases.Source: ETHICS STATEMENT · p. 10

A.1 OPTIMALITY OF THE ORTHOGONAL PROJECTOR

  • For a fixed removed subspace, the orthogonal projector is the unique map that annihilates it while minimizing Frobenius distance from the identity.The removed directions are chosen separately, by whitened initialization and subsequent loss optimization.Source: §A.1 · p. 16

A.2 PRACTICAL DESIGN CHOICES

  • Training ramps projection strength from identity and uses per-direction dropout; activation-space projection avoids materializing a dense projected weight.For tied input groups, sharing saves (|G| − 1) dinr parameters, which can instead support a higher rank, while constraining members to one input subspace.Source: §A.2 · pp. 16–17

A.3 WHITENED INITIALIZATION AND ITS EXCESS ERROR

  • Whitened initialization drops trailing directions; for input-tied layers, the optimal shared starting subspace comes from the leading right singular vectors of the row-stacked whitened weights.The initialized approximation’s error comprises the optimal rank-r error and an excess term; the alternative recast can amplify error for ill-conditioned removed directions.Source: §A.3 · p. 17

A.4 MEASURED-KL ALLOCATION

  • Rank allocation measures each unit’s output KL over candidate truncations, then greedily chooses moves with the smallest marginal cost per parameter saved until reaching the target budget.The allocation uses running-maximum KL curves and interpolation; sensitive units may remain dense.Source: §A.4 · p. 18

A.5 LSP ALGORITHM

  • The pipeline gathers calibration statistics, initializes and allocates projectors, jointly trains them with pretrained weights frozen, selects by validation, then merges factors into a compressed model.Training uses either KL distillation or the original task objective; export uses the evaluation projector and a thin SVD.Source: §A.5 · p. 19

A.6 MERGE DETAILS

  • At export, tied groups share a factor, and input-tied K/V can use one latent cache to reconstruct keys and values.The factorization is exact before numerical truncation; dropping singular values below the merge tolerance can reduce rank and add approximation error.Source: §A.6 · pp. 19–20

A.7 LIMITATIONS

The paper identifies limits in measured-KL allocation, compression cost, and the efficiency-performance trade-off created by K/V tying.

A.7 LIMITATIONS

  • Measured-KL allocation is a local surrogate: isolated unit costs omit interactions and provide no global-optimality guarantee.The paper reports that uniform ranks can still be preferable in some cases.Source: §A.7 · p. 20
  • LSP requires KL measurements and projection training, although measurements can be reused across ratios and objectives and inference savings can repay the one-off compression cost.Source: §A.7 · p. 20
  • Input-side K/V tying supports latent-cache savings, while Qwen3-4B’s best-performing output-side K/V configuration gains less inference efficiency.The paper frames this as a performance-efficiency trade-off; input-side tying is an alternative when improved inference efficiency is preferred.Source: §A.7 · p. 20

B IMPLEMENTATION DETAILS

Implementation uses calibration and held-out validation data for selection, with common compression settings and matched evaluation procedures across methods.

B IMPLEMENTATION DETAILS

  • Main accuracy runs use measured-KL allocation, whitened initialization, a linear-warmup cosine schedule, Adam without weight decay, and a 0.05 merge tolerance.Weights load in bf16; compression arithmetic uses fp32 for Qwen3-4B and Llama-2-7B and fp64 for the OPT models and ViT-B/16.Source: B IMPLEMENTATION DETAILS · p. 20
  • KL allocation measurements use 128K calibration tokens for LLMs and 4,096 ViT images; isolated-unit costs understate joint KL by 1.5–14× across the reported models.The joint-KL excess generally rises with compression ratio; Qwen3-4B uses output-side K/V tying and leaves those units mostly dense.Source: B IMPLEMENTATION DETAILS · p. 20
  • Validation selects the reported epoch and the best learning-rate setting per model and ratio; test data is not used during selection.LSP seed sensitivity is reported as mean and standard deviation over three full-pipeline seeds.Source: B IMPLEMENTATION DETAILS · pp. 20–21, Table 6 · p. 21
  • Compression ratios count removed parameters in targeted linear layers after merging, with shared factors counted once; realized savings are within 0.25 points of nominal.Evaluation reuses common calibration and evaluation setups across methods; perplexity uses non-overlapping 2048-token windows and zero-shot accuracy uses the EleutherAI harness.Source: B IMPLEMENTATION DETAILS · p. 21

C INITIALIZATION, TYING AND ALLOCATION

Ablations compare whitened initialization, tying, and measured-KL allocation against alternatives. The reported comparisons favor measured-KL allocation overall, while K/V routing exposes a performance-efficiency trade-off.

C INITIALIZATION, TYING AND ALLOCATION

  • After training, whitened initialization yields about 2% lower perplexity than plain initialization at every OPT-125M ratio and remains ahead on Llama-2-7B.Before training, neither initialization dominates across models and ratios.Source: Table 8 · p. 22, C INITIALIZATION, TYING AND ALLOCATION · p. 22
  • On Llama-2-7B, the whitened start leads under both tying choices by 1.9–3.0 Avg6 points and 20–33% in perplexity; tying raises Avg6 at every ratio.Perplexity under tying is mixed beyond −30%.Source: C INITIALIZATION, TYING AND ALLOCATION · p. 22, Table 9 · p. 23
  • Measured-KL allocation beats the best of three alternatives on OPT-1.3B at every ratio, by 0.2–0.9 perplexity and 0.7–1.0 Avg6 points.The compared alternatives are uniform retained parameters, uniform removed rank, and a calibration-loss diagonal-Fisher budget.Source: C INITIALIZATION, TYING AND ALLOCATION · p. 22, Table 10 · p. 23
  • Measured allocation leads in mean zero-shot accuracy in five of six settings, by up to 3.8 points, and trails by 0.4 points on Llama-2-7B at −50%.It leads in 26 of 36 benchmark cells, including ARC-e and PIQA in every setting; WinoGrande favors uniform allocation in four of six.Source: C INITIALIZATION, TYING AND ALLOCATION · pp. 22–23, Table 11 · p. 23
  • With input tying matched, measured allocation leads by 3.0 Avg6 points at −20% and trails by 0.1 and 0.4 at −40% and −60%, respectively.The uniform control assigns every compressed matrix the same retained fraction before accounting for ties.Source: C INITIALIZATION, TYING AND ALLOCATION · pp. 22, 24
  • On Qwen3-4B, output-side K/V routing leads at every ratio by 0.5–1.6 Avg6 points and 2–5% in perplexity, but leaves K/V nearly uncompressed.Its cache is as wide as the dense model’s at −20% and −40%; under measured allocation, K/V truncation costs about an order of magnitude more KL per parameter saved than other unit types.Source: C INITIALIZATION, TYING AND ALLOCATION · pp. 23–24, Table 12 · p. 24

D WHAT LSP COMPRESSES, ACROSS MODELS

Across models, LSP allocates compression differently by architecture and depth; its removals span the singular spectrum rather than simply cutting the tail. Deeper blocks are generally spared more, and its removal profile differs only modestly from NoLSP.

Where the budget goes

What the removal looks like, block by block

  • LSP removal spans the spectrum, deepens with compression, and reaches the first singular direction; the OPT-1.3B up projection loses 2.5–4.4 amplitude in direction 1 of block 3.The effect is strongest in OPT-1.3B’s up projection and Llama-2-7B’s tied Q/K/V, where direction 1 loses 0.7–2.3 amplitude.Source: D WHAT LSP COMPRESSES, ACROSS MODELS · pp. 25–26, Figure 5 · p. 26
  • Deeper blocks are spared more, while LSP and NoLSP removal differ by at most 2.5% in the shown Llama-2-7B blocks and 5.8% in OPT-1.3B’s block-21 down projection.LSP removes more of the leading quarter and keeps more of the trailing half in 171/180 OPT-1.3B units and 244/248 Llama-2-7B units at −50% and −70%.Source: D WHAT LSP COMPRESSES, ACROSS MODELS · p. 25

E INFERENCE EFFICIENCY

On Llama-2-7B, LSP improves decode throughput over dense and untied factorization, but the advantage varies with context length. Allocation and K/V routing expose different speed, cache, context-capacity, and accuracy trade-offs.

Allocation at deployment

  • Measured-KL allocation decodes 10%, 17% and 6% faster than uniform at −30%, −50% and −70%, but its cache is 38% wider at −30% and 8% and 22% narrower at −50% and −70%.At each ratio, both allocations hold the same weights and issue the same decode FLOPs.Source: E INFERENCE EFFICIENCY · p. 27, Table 14 · p. 28

Llama-2-7B inference efficiency

  • 1.21×, 1.36× and 1.56× faster than dense at −30%, −50% and −70%; at 512 tokens, LSP is 15–26% ahead of the best untied factorization.The advantage over dense falls to 1.11× at 16k tokens, where cache reading dominates.Source: Table 13 · p. 27, E INFERENCE EFFICIENCY · p. 27

K/V routing at deployment

  • No Qwen3-4B configuration wins all three axes: uniform with output-side K/V holds the longest context, measured input tying decodes fastest, and measured output-side K/V is most accurate.Input tying is 2–4% faster than output-side routing in all six pairs.Source: E INFERENCE EFFICIENCY · pp. 27–28, Table 15 · p. 28

F COMPRESSION COST

LSP’s compression cost is reported alongside accuracy and inference efficiency, and the paper compares it with alternative methods and calibration pools. In ViT transfer tests, a diverse calibration pool outperforms a size-matched CIFAR-100-only pool.

F COMPRESSION COST

  • 61.2 against 10.9 perplexity on Llama-2-7B at −70%: Dobi-SVD remains far behind LSP after 20 epochs on the same 1024 samples.At higher budgets, one of LSP’s two objectives gives the lowest perplexity at every model–ratio setting, within 0.5–10 wall-clock hours.Source: F COMPRESSION COST · p. 29, Figure 6 · p. 29

G TRANSFER ON THE GROWING CALIBRATION POOL

Related papers

Loading 2609.40127v1…