在线伪均值偏移注意力（PASA）用于鲁棒的低精度大语言模型推理：算法与数值分析
机器学习
2025-03-05 v1 人工智能
数值分析
性能
数值分析
摘要
对于长序列推理任务（如文本或图像/视频生成），在大模型中注意力计算极其耗时。为加速这一过程，我们基于 Flash Attention 开发了一种低精度、数学等价的算法 PASA。PASA 引入了两种新技术：在线伪均值偏移和全局恢复。这些技术使 Flash Attention 全程可以使用半精度计算，而不会引发溢出不稳定或不可接受的数值精度损失。该算法通过减少数据移动和增加计算 FLOPs，增强了在内存受限的 AI 硬件架构（如昇腾神经网络处理器 NPU）上的性能。通过设计随机基准和真实大模型对该算法进行了验证。我们发现，注意力输入数据的大偏差和大幅度值是导致两类大模型（Qwen2-7B 语言模型和 Stable-Video-Diffusion 多模态模型）中半精度数值溢出（>65504）的关键因素。具体而言，溢出源于序列维度上的大偏差以及 Stable-Video-Diffusion 模型中查询矩阵与键矩阵在头维度上的共振机制。共振机制定义为查询矩阵与键矩阵之间的相位一致或 180 度相位差，它会显著放大注意力分数矩阵中的元素值。此问题同样适用于 Qwen 模型。此外，通过均方根误差（RMSE）以及将最终生成的文本和视频与高精度注意力生成的结果进行比较来评估数值精度。
引用
@article{arxiv.2503.01873,
title = {Online Pseudo-average Shifting Attention(PASA) for Robust Low-precision LLM Inference: Algorithms and Numerical Analysis},
author = {Long Cheng and Qichen Liao and Fan Wu and Junlin Mu and Tengfei Han and Zhe Qiu and Lianqiang Li and Tianyi Liu and Fangzheng Miao and Keming Gao and Liang Wang and Zhen Zhang and Qiande Yin},
journal= {arXiv preprint arXiv:2503.01873},
year = {2025}
}
备注
21 Pages, 14 figures, conference paper