基于矩阵字典学习的Transformer权重共享

大型语言模型已彻底革新了AI应用，但其高计算和内存需求阻碍了其广泛部署。现有压缩技术侧重于层内优化（如低秩近似或注意力剪枝），而Transformer的重复层结构暗示了显著的层间冗余——这一维度在KV缓存之外鲜有探索。灵感来自卷积网络中的字典学习，我们提出了一种跨Transformer层结构化权重共享框架。该方法将注意力投影矩阵（Q、K、V、O）分解为共享字典原子，减少注意力模块的参数占用66.7%，同时实现与原模型持平的性能。与需要蒸馏或架构变化的复杂方法不同，MASA（Matrix Atom Sharing in Attention）可作为即插即用方案，通过标准优化器进行训练——表示每个层的权重为共享矩阵原子的线性组合。跨尺度（100M-700M参数）的实验表明，MASA在可比参数预算下，在基准准确率和perplexity方面优于GQA、低秩基线方法以及最近的Repeat-all-over/Sequential sharing。消融研究确认了对字典大小的鲁棒性以及共享表示在捕获跨层统计规律方面的有效性。扩展到视觉Transformer（ViT），MASA在图像分类任务中以66.7%更少的注意力参数实现了相当的性能指标。通过将字典学习策略与Transformer效率结合，MASA为参数高效模型提供了可扩展的蓝图，而不牺牲性能。最后，我们探索了在大型预训练模型上应用MASA以减少参数数量，而不会出现显著性能下降的可能性。