学习神经网络的函数子空间以实现压缩

中文摘要 AI 辅助

现代Transformer在提供强大能力的同时，也带来了巨大的内存和计算需求。低秩权重分解在保持矩阵稠密（从而在标准硬件上高效）的同时，减少了这两方面需求。然而，现有方法通过局部闭式准则（如激活能量、逐层重建误差或损失的二次近似）来选择从每个权重矩阵中移除的子空间。这些准则忽略了误差在网络中的传播方式，因此在高度压缩时，误差会随深度累积，导致性能崩溃。我们引入了可学习子空间投影（LSP），该方法端到端地学习要丢弃的子空间。每个线性层，或读取相同激活的绑定层组，被分配一个正交投影器。所有投影器在预训练权重保持冻结的情况下，针对全局目标（即与稠密模型输出分布的KL散度或模型原始训练损失）进行联合优化。投影器从白化SVD截断初始化，并根据每个投影器每节省参数所产生的输出KL来分配秩。训练后，投影器合并为标准低秩因子，每个绑定组共享一个因子。在注意力机制中，这也使得模型能够缓存一个窄潜变量以替代完整的键和值。在多种LLM（OPT-125M/1.3B、Qwen3-4B、Llama-2-7B）和ViT-B/16上，LSP优于基线方法，且其优势随压缩程度增加而扩大。在-70%压缩率下，LSP使Llama-2-7B达到10.9的WikiText-2困惑度和42.2%的平均零样本准确率，而最强基线分别为13.3和36.0%。分解后的模型在小批量大小下解码速度比稠密模型快达1.6倍，且在128k-token上下文下，缓存共享潜变量将权重和KV缓存的总内存缩小13.5倍，而解绑基线分解最多仅缩小6.5倍。

英文摘要

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.

发表机构

• Helmholtz Munich（亥姆霍兹慕尼黑研究中心）
• Technical University of Munich（慕尼黑工业大学）
• MCML（慕尼黑机器学习中心）
• Orbital Industries（轨道工业公司）
• LTCI, Télécom Paris, Institut Polytechnique de Paris（巴黎高等电信学院LTCI实验室，巴黎综合理工学院）
• Columbia University（哥伦比亚大学）
• University of Copenhagen（哥本哈根大学）
• Technical University of Denmark（丹麦技术大学）
• New York University（纽约大学）

机构由 AI 辅助整理，请以论文原文为准。