Learning Functional Subspaces for Neural Network Compression

Learning Functional Subspaces for Neural Network Compression. By Massimo Bini and seven others, from Helmholtz Munich, the Technical University of Munich, and seven other institutions. Published on the thirtieth of September. No code repository is linked to it.

Modern transformers are capable, and they are demanding in memory and in compute. Low-rank factorisation of the weights reduces both, while keeping the matrices dense, which keeps them efficient on standard hardware. The question is which subspace to remove from each weight matrix, and existing methods answer it locally, with closed-form criteria. Some use activation energy. Some use layer-wise reconstruction error. Some use a quadratic approximation of the loss. All of these look at one layer at a time and ignore how errors propagate through the network. At high compression, that blind spot shows: errors compound with depth, and performance collapses.

Learnable Subspace Projections, or LSP, learn the subspaces to discard end-to-end instead. Each linear layer is assigned an orthogonal projector, or, where several layers read the same activations, the tied group shares one. All the projectors are optimised jointly against a single global objective, while the pretrained weights stay frozen. That objective is either the KL divergence to the dense model's output distribution, or the model's original training loss.

The projectors start from a whitened SVD truncation. Ranks are allocated by how much output KL each projector induces per parameter saved, so capacity goes where removing it would cost the most. Once training is done, the projectors merge into standard low-rank factors, with each tied group sharing one factor. In attention, that sharing has a further benefit: the model can cache one narrow latent in place of the full keys and values.

Tests cover several LLMs, OPT-125M and its larger sibling, Qwen3-4B and Llama-2-7B, as well as the vision transformer ViT-B sixteen. LSP outperforms the baselines, and its advantage widens as compression increases. At seventy percent compression, LSP brings Llama-2-7B to a WikiText-2 perplexity of ten point nine, and a mean zero-shot accuracy of forty-two point two percent.

The strongest baseline, at the same compression, reaches a perplexity of thirteen point three and an accuracy of thirty-six percent. The factorised model decodes up to one point six times faster than the dense model at small batch sizes.

Caching the shared latent shrinks the combined memory of weights and KV cache by thirteen point five times at a context of one hundred and twenty-eight thousand tokens. Untied baseline factorisations manage at most six point five times.

The shift is from local to global. Rather than asking which directions each layer can lose on its own terms, LSP asks which directions the whole network can lose, measured at the output. That is why its lead grows exactly where local criteria were failing, at high compression, and the tying of layers turns the same idea into a smaller cache as well as smaller weights.

Each step of the recipe serves that global view. Projectors start from a whitened SVD truncation, so training begins from a sensible point rather than from scratch. Ranks follow the output KL per parameter saved, which is again a measure taken at the output rather than inside one layer. And after training, nothing exotic is left behind. The projectors fold into ordinary low-rank factors that run efficiently on standard hardware, which was the reason for choosing low-rank factorisation in the first place. The gains hold across language models of several sizes and a vision transformer.