UniAttn：通过 Softmax 统一降低后训练大语言模型的推理成本
计算与语言
2026-01-23 v2
摘要
后训练对于将大语言模型（LLMs）适配到实际应用至关重要。部署后训练模型面临显著挑战，原因在于巨大的内存开销和明显的推理延迟。现有工作已识别出 LLMs 中存在大量冗余，并提出了高效的架构，即层内 KV 共享和跨层 KV 共享。然而，这些方法仍然导致较高的推理时间开销，对于后训练预训练 LLMs 而言并非最优。本文中，我们发现 Softmax 操作是 LLM 推理的主要瓶颈，并揭示其在后训练期间实际上高度冗余。我们提出了注意力中的 Softmax 统一（UniAttn），这是一种新颖的后训练方法，通过统一 Transformer 模块间的 Softmax 激活来降低 LLM 推理成本。此外，UniAttn 采用线性投影来补偿 Softmax 统一引入的误差。实验表明，UniAttn 在显著降低推理成本的同时，达到了标准后训练的性能，在后训练期间优于现有的高效架构。
引用
@article{arxiv.2502.00439,
title = {UniAttn: Reducing Inference Costs via Softmax Unification for Post-Training LLMs},
author = {Yizhe Xiong and Wei Huang and Xin Ye and Hui Chen and Zijia Lin and Haoran Lian and Zhenpeng Su and Jungong Han and Guiguang Ding},
journal= {arXiv preprint arXiv:2502.00439},
year = {2026}
}
备注
8 pages, 6 figures. Preprint, under review