MASA 把稀疏注意力选择改为矩阵近似,在三种稀疏框架上以零额外预算提升长上下文准确率
核心概要
作者指出稀疏注意力现有选择规则只保留注意力矩阵中数值大或质量高的条目,而 Transformer 实际是通过矩阵乘法使用该矩阵,因此提出 MASA:用矩阵乘积近似误差的闭式分数替换原有的注意力质量排序,在不改变稀疏核、候选单元与预算的前提下接入 MInference、SeerAttention、FlexPrefill 三种框架,在 RULER、InfiniteBench、LongBench 上取得一致的准确率提升且延迟开销可忽略。
Figure 1: A minimal two-column counterexample. Under the same one-column budget, conventional mass-based ranking selects u 1 u_{1} , whose mass is 66 × 66\times larger but incurs 98 % 98\% relative matrix-product error. MASA selects u 2 u_{2} and reduces the error to 3 % 3\% . Here, 𝑿 {\bm{X}} is an illustrative probe matrix.
arXiv深度剖析
论文把稀疏注意力的选择目标重新表述为矩阵近似:在固定预算下应保留使矩阵乘积 A·V 误差最小的稀疏单元,而不是保留注意力矩阵中数值最大的条目。 此前方法(MInference、SeerAttention、FlexPrefill 等)的选择规则最终都是幅度驱动,即保留大标量条目或高质量区域;论文指出这种逐元素近似只在注意力矩阵在 token 位置上近似单位阵时才等价于输出误差最小,而长上下文下该条件不成立。 论文给出推导:展开矩阵乘积误差后,逐元素幅度排序仅对应单位 Gram 特例;并给出一个两列的最小实例,其中质量更低的候选因更好地保持矩阵乘积而获得更高 MASA 分数。
MASA 给出可插拔的闭式打分:对候选稀疏单元按保留该单元后平方矩阵乘积误差的下降量排序,对垂直、斜线单元退化为逐对分数之和,对块单元保留块内条目交互项。 该方法不改变基线的稀疏单元定义、稀疏核与预算规则,只替换排序分数,因此与底层注意力核和 KV 缓存设计正交;打分在基线已有的 profile 分辨率上完成,不需要新的全分辨率稠密注意力图。 论文给出逐对分数的闭式形式,说明分数同时依赖探针向量的方向与范数以及其与稠密矩阵作用的对齐;并给出避免物化贡献张量的矩阵乘法实现,以及软范数平衡与范数裁剪两种探针构造。
在三种稀疏注意力基线、三个长上下文基准与 Llama/Qwen 两类骨干上,MASA 在全部八个匹配平均分上均有提升,且长上下文处收益更明显。 所有对比保持模型、稀疏预算、稀疏核、候选单元族与数据处理不变,只替换排序分数,因此隔离了选择目标本身的影响;例如 RULER 上 SeerAttention 由低于稠密参考变为高于稠密参考,FlexPrefill-Llama 在 RULER 上的增益超过其原有相对全注意力的优势。 RULER 32K/64K/128K 的四个配置中 12 个长度级匹配分数有 11 个提升,且每个配置在 32K–128K 的平均增益都高于 4K–16K;附录 E 在四个框架–骨干组合上直接测得归一化平方输出误差下降,报告的配对区间均高于零。
替换排序分数带来的额外计算在实测中可忽略:MASA 只在既有排序阶段评估下降分数,不新增注意力组件。 这使准确率提升不以牺牲稀疏注意力的效率优势为代价,在 128K 这一稀疏注意力最需要的区间开销仍然可忽略。 论文报告以稠密 FlashAttention-2 归一化的核级延迟,覆盖 32K、64K、128K 三个序列长度与四个框架–骨干组合。
启示与展望
MASA 的作用范围是在给定稀疏注意力框架内修正选择分数,它不搜索新的稀疏核、新的稀疏单元族或新的预算调度;论文也把这一点作为使对比干净的边界。其实现与实验实例化在稀疏预填充场景,即解码前对长提示计算注意力的阶段,基线为 MInference、SeerAttention 与 FlexPrefill。适用对象是希望在不改动现有稀疏核与预算的前提下提升长上下文准确率的工程与研究实践者;探针构造(原始值、软范数平衡、范数裁剪)是留给使用者的实现选择,论文按基准选取聚合分数最高的已测配置。
读者仍需关注:探针选择(原始值、软范数平衡、范数裁剪及其阈值)在不同框架与基准上的最优配置并不一致,论文按基准选取已测配置中聚合分数最高者,因此实际部署时的探针调参空间值得留意。附录 E 的误差测量使用 64K RULER 输入、每模型 208 条提示、八个预设层,并说明所报区间描述的是这些选定设置下的比较而非独立留出评估。此外,正文表格中若干数值在加载文本中缺失,只能依据附录 D 与附录 F 的完整表格核对具体分数;若需精确复现某一配置的数值,应以附录完整表格为准。
