arXiv:2605.12879cs.LG2026-05

提出ASAP方法,让注意力机制推理更快且不降精度。

ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection

论文配图:ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection
图 1 · 摘自论文原文
  • 用切片对偶投影替代迭代计算,推理时无需重复优化
  • 训练阶段用Sinkhorn,推理时速度比教师模型快5.3倍
  • 适合需要快速推理的部署场景,尤其适配冻结层任务

双重随机化注意力作为基于传输的行软最大值注意力替代方案,近年来在变压器变体中被用于减少注意力塌陷和秩崩溃,同时提升性能。标准方法为Sinkhorn缩放,虽训练效率高,但在每次推理前仍需重复矩阵缩放。切片传输注意力消除了在线迭代,但其软排序近似会为每个切片生成稠密张量,显著增加训练资源消耗。本文提出ASAP:通过切片对偶投影实现可摊销的双重随机化注意力,先以Sinkhorn训练双随机层,再将推理中的迭代缩放循环替换为固定的切片对偶算子。该方法学习从精确的一维Kantorovich势能到Sinkhorn查询侧对偶的轻量级参数映射,并通过双向熵性c-变换重建注意力计划。在语言与视觉基准测试中,ASAP保持低成本训练设置,性能接近最新基线。在主要的冻结层基准上,ASAP比训练好的Sinkhorn教师模型快5.3倍且准确率相当;在下游替换任务中,无需重训练即可恢复教师模型大部分性能。

原文摘要 · Abstract (English)

Doubly-stochastic attention has emerged as a transport-based alternative to row-softmax attention, with recent Transformer variants using it to reduce attention sinks and rank collapse while improving performance. In this family, the standard approach is Sinkhorn scaling, which trains more efficiently but still repeats matrix scaling in every inference forward pass. Sliced-transport attention removes the online iteration, but its soft sorting approximation materializes dense tensors for each slice, requiring substantially more training resources than Sinkhorn attention. We introduce ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection, a train-then-compile method that trains the doubly-stochastic layer with Sinkhorn, then replaces the iterative scaling loop at inference with a fixed sliced-dual operator. It learns a lightweight parametric map from exact one-dimensional Kantorovich potentials to the Sinkhorn query-side dual, then reconstructs the attention plan with a two-sided entropic c-transform. Across language and vision benchmarks, ASAP keeps the cheaper training setup and remains highly competitive with recent baselines. In the main frozen-layer benchmark, ASAP is 5.3 faster than the trained Sinkhorn teacher while matching its accuracy; in downstream replacements, ASAP recovers most of the teacher performance without any retraining.

注意力机制推理加速深度学习优化Transformer

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。