TensorLLM: Tensorising Multi-Head Attention for Enhanced Reasoning and Compression in LLMs

本文提出了一种基于多头张量化和Tucker分解的框架，通过强制共享高维子空间对大型语言模型的多头注意力权重进行结构化去噪和压缩，显著提升推理能力并实现高达247倍的压缩率。

Large Language Model, Transformer, Reasoning, Efficiency, Multimodality

Yuxuan Gu, Wuyang Zhou, Giorgos Iacovides, Danilo Mandic

Imperial College London

Generated by grok-3

Background Problem

大型语言模型（LLMs）基于Transformer架构，在自然语言处理（NLP）中取得了巨大成功，但其过参数化特性导致推理和部署成本高昂。现有研究表明，通过结构化去噪和压缩模型权重，可以在不显著牺牲性能的情况下提升效率和推理能力。然而，当前方法（如LASER和TRAWL）主要聚焦于前馈网络（FFN）块的去噪，对Transformer核心组件——多头注意力（MHA）块的去噪效果不佳。论文试图解决这一问题，探索如何通过利用MHA的领域知识（即同一层内注意力头共享相似抽象层次但具有不同特化功能），提升LLMs的推理能力并实现MHA权重的高效压缩。

Method

论文提出了一种新颖的框架，通过多头张量化（Multi-Head Tensorisation）和Tucker分解对MHA权重进行结构化去噪和压缩。具体方法如下：

• 核心思想：基于MHA的领域知识，假设同一层内多个注意力头的权重共享一个高维子空间，通过强制共享因子矩阵的方式进行去噪和压缩，以提升推理能力。
• 实现步骤：
  1. 多头张量化：将MHA的查询（Q）、键（K）、值（V）和输出（O）权重矩阵从2D格式拆分为每个注意力头的子矩阵，并堆叠成3D张量（维度为$d_{model} \times d_v \times 4$），最终形成包含所有头权重的4D张量。
  2. 共享因子矩阵的Tucker分解：对每个注意力头的3D张量应用Tucker分解，但强制所有头共享一组相同的因子矩阵（$U^{(1)}, U^{(2)}, U^{(3)}$），仅核心张量（$G_i$）因头而异。这种设计确保各注意力头在同一高维子空间内编码不同信息。
  3. 去噪与压缩：通过最小化原始张量与分解后张量之间的Frobenius范数损失（$\frac{1}{2} \| \mathcal{W}_{all} - \mathcal{G}_{all} \times_1 \mathbf{U}^{(1)} \times_2 \mathbf{U}^{(2)} \times_3 \mathbf{U}^{(3)} \|_F^2$），实现权重去噪，并通过设置较低的多线性秩（$R_n \ll I_n$）实现参数压缩。
• 关键点与批判：方法创新在于利用MHA的直觉设计共享子空间，但缺乏理论依据证明为何共享因子矩阵能提升推理能力。此外，Tucker分解的多线性秩作为超参数，其选择对结果影响较大，但论文未提供系统性选择策略，可能导致方法在不同模型或数据集上的不稳定性。

Experiment

实验在三个LLM模型（RoBERTa 125M、GPT-J 6B、LLaMA2 7B）上进行，涵盖编码器和解码器架构，并使用四个推理基准数据集（HotPotQA、FEVER、Bios Profession、BigBench-WikidataQA）评估性能。

• 实验设置：方法以层选择方式应用（每次仅对单层Transformer进行去噪和压缩），与现有方法（如LASER）进行公平比较。实验还包括与FFN去噪方法的联合测试，以及对MHA权重单独和联合张量化的消融研究。使用NVIDIA A100 GPU进行计算。
• 结果分析：
  1. 推理能力提升：在所有模型和数据集上，方法均提升了测试准确率（如GPT-J在BigBench-WikidataQA上从51.8%提升至68.81%），表明去噪对推理能力有正向影响。
  2. 压缩效果：MHA权重压缩率最高达247.3倍（GPT-J在HotPotQA上），显示出显著的参数效率提升。
  3. 联合应用：与LASER（FFN去噪）结合后，方法在大多数情况下进一步提升准确率（如LLaMA2在Bios Profession上从86.7%提升至87.07%）。
  4. 消融研究：联合张量化所有MHA权重（Q、K、V、O）优于单独处理，验证了方法设计的合理性。
• 批判性评价：尽管准确率提升明显，但某些情况下损失值（Loss）增加（如LLaMA2在HotPotQA上从3.15增至9.80），可能表明去噪过程对模型预测分布的负面影响未被充分考虑。实验仅限于单层应用，未探讨多层去噪的累积效应或层间交互影响。此外，超参数（如多线性秩）对不同数据集的影响未详细分析，限制了结果的泛化性解释。总体而言，实验设置较为合理，但不够全面，未能完全揭示方法在实际复杂场景中的潜力与局限。

Further Thoughts

论文提出的共享因子矩阵设计为MHA权重去噪提供了一种新颖视角，但其理论基础仍显薄弱。未来研究可以探索共享子空间的具体机制，例如通过可视化或分析因子矩阵的语义含义，验证其是否真正捕捉了注意力头的协同作用。此外，方法在高压缩率下的潜在性能权衡值得深入研究，尤其是在长上下文任务或多模态模型中，压缩可能对模型的泛化能力产生更大影响。另一个有趣的方向是将此方法与其他参数高效微调技术（如Low-Rank Adaptation）结合，探索去噪与微调的协同效应，尤其是在资源受限环境下的应用场景。最后，考虑到作者提到的超参数敏感性问题，可以引入自动化超参数搜索或元学习方法，以提升方法的通用性和实用性。