通过统一注意力中的Softmax,降低大模型推理开销。
UniAttn: Reducing Inference Costs via Softmax Unification for Post-Training LLMs
- 将多层注意力的Softmax激活统一,减少计算冗余。
- 保持与标准微调相当的性能,推理延迟显著下降。
- 适合需要高效部署的大模型应用开发者。
后训练是使大语言模型适配真实应用场景的关键步骤。然而,部署后训练模型面临显著的内存开销和明显的推理延迟问题。现有方法虽识别出模型中存在大量冗余,提出层内和跨层键值对共享等高效架构,但仍导致较高的推理时间开销,对后训练的预训练大模型而言仍不理想。本文发现,Softmax操作是大模型推理的主要瓶颈,且在后训练阶段高度冗余。为此,我们提出一种新型后训练方法——注意力软最大化统一(UniAttn),通过在Transformer块间统一Softmax激活来降低推理成本。同时,UniAttn采用线性投影补偿因统一带来的误差。实验表明,UniAttn在保持与标准后训练相当性能的同时,显著降低推理开销,优于现有高效架构在后训练场景的表现。
原文摘要 · Abstract (English)
Post-training is essential for adapting Large Language Models (LLMs) to real-world applications. Deploying post-trained models faces significant challenges due to substantial memory overhead and noticeable inference latency. Existing work has identified significant redundancies in LLMs and proposed efficient architectures, namely intra-layer KV sharing and cross-layer KV sharing. However, these methods still result in high inference time overhead, remaining suboptimal for post-training pre-trained LLMs. In this paper, we identify that the \texttt{Softmax} operation is a primary bottleneck for LLM inference and discover that it is actually highly redundant during post-training. We propose Softmax \textbf{Uni}fication in \textbf{Att}e\textbf{n}tion (\textbf{UniAttn}), a novel post-training method that unifies Softmax activations across transformer blocks to reduce LLM inference costs. Additionally, UniAttn adopts a linear projection to compensate for the errors induced by Softmax unification. Experiments show that UniAttn matches the performance of standard post-training while significantly reducing inference costs, outperforming existing efficient architectures during post-training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。