ConSA: 通过可学习分配实现混合注意力中的可控稀疏性

# 通过可学习分配实现混合注意力的可控稀疏性
来源：https://arxiv.org/html/2606.18056
陈遥¹,²，杨引齐³∗，尚君远³，郝祥和³，张思萌¹,²，陈奕龙¹,²，刘廷文¹,²†，王硕欢³，于大海³
¹中国科学院信息工程研究所
²中国科学院大学网络空间安全学院
³百度公司
\{chenyao2023, liutingwen\}@iie.ac.cn \{yangyinqi, shangjunyuan, wangshuohuan\}@baidu.com

###### 摘要

结合全注意力（FA）和滑动窗口注意力（SWA）的混合架构是实现高效 LLM 推理的一种有前景的范式。然而，现有方法通常依赖手工规则或简单的后验启发式进行 FA/SWA 分配，并且对这些设计背后的注意力行为分析有限。我们提出了 ConSA（混合注意力中的可控稀疏性），一个在用户指定的稀疏度目标下学习最优 FA/SWA 分配的框架。ConSA 采用 L0 正则化来学习选择每个注意力单元是使用 FA 还是 SWA 的二元掩码，同时通过增广拉格朗日约束在层粒度或 KV 头粒度上强制达到目标稀疏度。我们在两个 LLM（0.6B 和 1.7B 规模）上评估了 ConSA。学习到的分配始终优于基于规则的基线，其中 KV 头级分配相比层级分配带来了明显的增益。学习到的模式将 SWA 置于底层，并将 FA 集中在连续的中间层块中，这与基于规则方法中的均匀交错模式不同。这种结构在模型规模、稀疏度水平和分配粒度上持续存在，揭示了学习分配背后细粒度的内在注意力行为谱系。

ConSA：通过可学习分配实现混合注意力的可控稀疏性

陈遥¹,²†††表示同等贡献。†表示通讯作者。，杨引齐³∗，尚君远³，郝祥和³，张思萌¹,²，陈奕龙¹,²，刘廷文¹,²†，王硕欢³，于大海³
¹中国科学院信息工程研究所
²中国科学院大学网络空间安全学院
³百度公司
\{chenyao2023, liutingwen\}@iie.ac.cn
\{yangyinqi, shangjunyuan, wangshuohuan\}@baidu.com

## 1 引言

大型语言模型使注意力成本成为部署瓶颈：全注意力（FA）在计算上随序列长度呈二次方缩放，在 KV 缓存上呈线性缩放 [Kwon et al. (2023)](https://arxiv.org/html/2606.18056#bib.bib38)。在高效的替代方案中，滑动窗口注意力（SWA）[Beltagy et al. (2020)](https://arxiv.org/html/2606.18056#bib.bib7) 将每个 token 限制在一个固定的局部窗口内，从而在推理期间减少了注意力计算和每头 KV 缓存。然而，固定窗口丢弃了长程依赖，这可能损害需要全局上下文的任务 [Xiao et al. (2024)](https://arxiv.org/html/2606.18056#bib.bib39)。平衡成本与能力的一个自然策略是在单个架构中结合 FA 和 SWA。

生产模型如 Mistral [Jiang et al. (2023)](https://arxiv.org/html/2606.18056#bib.bib1)、Gemma 2 [Team (2024)](https://arxiv.org/html/2606.18056#bib.bib2) 和 MiMo-V2-Flash [Xiaomi (2026)](https://arxiv.org/html/2606.18056#bib.bib17) 已经通过手工制作的交错模式采用了这种混合设计。然而，这些手动指定的分配并未考虑到原始模型中跨层和跨头的异质性注意力行为 [Xiao et al. (2025)](https://arxiv.org/html/2606.18056#bib.bib15)。LoZA [Zhang et al. (2025)](https://arxiv.org/html/2606.18056#bib.bib3) 用一个轻量级校准阶段取代了手动设计，该阶段通过可学习的标量权重对层进行评分，并将得分低的层转换为局部注意力。然而，当在少量预训练数据上校准时，这种标量分数在层间的区分能力可能有限，使得层选择在不同目标稀疏度水平下不太可靠。此外，先前关于混合注意力的工作 [Xiao et al. (2025)](https://arxiv.org/html/2606.18056#bib.bib15)；[Zhang et al. (2025)](https://arxiv.org/html/2606.18056#bib.bib3)；[Zhao et al. (2026a)](https://arxiv.org/html/2606.18056#bib.bib40) 很少分析跨层和跨头会出现什么样的 FA/SWA 模式，以及这些模式如何与原模型的内在注意力行为相关联。

这些局限性促使我们提出一种可学习的分配方法，该方法在显式的稀疏性目标下优化 FA/SWA 分配，同时还需要对学习到的分配背后内在注意力行为进行更细粒度的分析。我们提出了 ConSA（混合注意力中的可控稀疏性），一个在可控稀疏性下学习混合 FA/SWA 分配的框架。给定一个预训练的 Transformer 和用户指定的目标稀疏度ρ\\rho，ConSA 将混合注意力形式化为一个稀疏约束优化问题：每个注意力单元接收一个二元掩码，该掩码通过 L0 正则化下的硬具体分布 [Louizos et al. (2018)](https://arxiv.org/html/2606.18056#bib.bib28)；[Xia et al. (2024)](https://arxiv.org/html/2606.18056#bib.bib6) 参数化，用于在 FA 和 SWA 之间进行选择。我们设计了一个增广拉格朗日约束来强制达到目标ρ\\rho，使模型能够在层粒度或 KV 头粒度上发现最优分配。掩码参数和模型权重首先在掩码学习阶段联合优化，之后学习到的掩码被二值化并固定，用于继续预训练。

我们进一步分析了在模型规模和稀疏度水平下学习到的分配模式和模型的固有注意力行为。学习到的掩码一致地将 SWA 放在底层，并将 FA 集中在连续的中间层块中，这与基于规则的方法中使用的均匀交错模式不同。检查学习分配下代表性层和头的注意力行为，揭示出多样化的注意力尖峰范围，超越了先前工作中描述的检索与流式二分法 [Xiao et al. (2025)](https://arxiv.org/html/2606.18056#bib.bib15)，并且与 ConSA 学习到的分配良好对齐。

我们的贡献有三点：(1) 我们提出了 ConSA，一个通过 L0 正则化和增广拉格朗日优化在层级和 KV 头级粒度上学习混合 FA/SWA 分配的框架，使用户能够指定任意目标ρ\\rho，该目标在优化过程中被可靠地满足。(2) 在两个模型规模（0.6B 和 1.7B）和多个稀疏度水平上的实验表明，学习到的分配始终优于基于规则的基线，其中 KV 头级分配相比层级分配带来了明显的增益。消融研究进一步证实 L0-拉格朗日公式优于依赖无约束标量门控和后验排序的基于校准的方法。(3) 对学习到的模式的分析揭示了跨模型规模、稀疏度水平和分配粒度一致的 SWA-底部 / FA-中间结构。对内在注意力行为的检查表明，这种结构与多样化的注意力尖峰范围对齐，超越了先前工作中的检索与流式二分法。

## 2 相关工作

#### 高效注意力机制。

全注意力的二次方缩放导致了各种高效的替代方案。滑动窗口注意力（SWA）[Beltagy et al. (2020)](https://arxiv.org/html/2606.18056#bib.bib7)；[Zaheer et al. (2020)](https://arxiv.org/html/2606.18056#bib.bib8)；[Child et al. (2019)](https://arxiv.org/html/2606.18056#bib.bib9) 是一个常见选择，因为它限制了推理期间的计算开销和 KV 缓存占用。其他方法包括线性注意力 [Katharopoulos et al. (2020)](https://arxiv.org/html/2606.18056#bib.bib10)、具有学习模式的稀疏注意力 [Kitaev et al. (2020)](https://arxiv.org/html/2606.18056#bib.bib11)；[Roy et al. (2021)](https://arxiv.org/html/2606.18056#bib.bib12) 以及状态空间模型 [Gu and Dao (2023)](https://arxiv.org/html/2606.18056#bib.bib13)。我们的工作没有引入新的注意力机制，而是专注于如何在模型内分配 FA 和 SWA。

#### 混合注意力架构。

最近的 LLM 通常通过手工模式结合 FA 和 SWA：Mistral [Jiang et al. (2023)](https://arxiv.org/html/2606.18056#bib.bib1) 交替使用 SWA 和 FA 层，Gemma 2 [Team (2024)](https://arxiv.org/html/2606.18056#bib.bib2) 使用与规模相关的交错方式，Command-R 和 Jamba [Lenz et al. (2025)](https://arxiv.org/html/2606.18056#bib.bib14) 采用混合类型。最近两项工作转向了可学习分配：SwiAttn [Zhao et al. (2026b)](https://arxiv.org/html/2606.18056#bib.bib29) 通过每层路由器将 token 路由到 FA 或 SWA，但必须保留统一的 KV 缓存；LoZA [Zhang et al. (2025)](https://arxiv.org/html/2606.18056#bib.bib3) 校准一个每层标量权重，并将排名最低的层以固定的 50% 比例转换为流式稀疏注意力。ConSA 的不同之处在于将 FA/SWA 分配形式化为一个稀疏约束优化问题，通过增广拉格朗日约束强制达到用户指定的目标；详细比较见附录 A (https://arxiv.org/html/2606.18056#A1)。

#### 注意力头分析。

已知注意力头扮演着不同的角色，例如跟踪位置、语法或罕见词 [Voita et al. (2019)](https://arxiv.org/html/2606.18056#bib.bib4)；[Clark et al. (2019)](https://arxiv.org/html/2606.18056#bib.bib5)。更近期的工作识别出检索头（将大量注意力权重分配给整个上下文中的少数关键 token）和流式头（主要关注近期 token 和注意力汇聚点）；这种分类源于对合成长程检索任务的输出偏差分析，并已指导了 KV 缓存压缩 [Xiao et al. (2025)](https://arxiv.org/html/2606.18056#bib.bib15)。ConSA 学习到的分配揭示了跨模型规模、稀疏度水平和分配粒度一致的 SWA-底部 / FA-中间结构。对代表性层和头的分析表明，它们内在的注意力尖峰范围形成了一个超越这种二元分类的更加细粒度的谱系，并且与学习到的 FA/SWA 分配良好对齐。

参见图注图 1：ConSA 概览。左图：两阶段训练流水线。阶段 1 在 11B 个 token 上联合优化模型参数θ\\theta、掩码参数α\\alpha 和拉格朗日乘子 {λ,ϕ}\\\{\\lambda,\\phi\\\}，约束条件 ρ^\(z\)=ρ\\hat\{\\rho\}\(z\)=\\rho 强制执行用户指定的目标稀疏度。阶段 2 对掩码进行二值化，并在固定的 FA/SWA 分配下继续进行 100B 个 token 的预训练。右图：每头分配机制。对于每个 KV 头 (l,i)\(l,i\)，一个硬具体掩码 zl,iz\_\{l,i\} 由可学习的 αl,i\\alpha\_\{l,i\} 参数化，在全注意力（FA）和滑动窗口注意力（SWA）之间进行选择。

## 3 预备知识

我们考虑一个具有 LL 层的 Transformer，每层包含多个键值（KV）头。FA 和 SWA 可以应用于不同的粒度；我们在 KV 头级别形式化两者，这是我们的方法中考虑的最细粒度。

#### 全注意力（FA）。

第 ll 层中第 ii 个 KV 头组的输出为：

Ol,iFA=softmax\(Ql,iKl,i⊤dk\)Vl,i,\\mathbf\{O\}\_\{l,i\}^\{\\mathrm\{FA\}\}=\\mathrm\{softmax\}\\\!\\left\(\\frac\{\\mathbf\{Q\}\_\{l,i\}\\mathbf\{K\}\_\{l,i\}^\{\\top\}\}\{\\sqrt\{d\_\{k\}\}\}\\right\)\\mathbf\{V\}\_\{l,i\}, (1) 其中 Ql,i\\mathbf\{Q\}\_\{l,i\} 表示该组中所有 gg 个查询头的拼接查询，Kl,i,Vl,i∈Rn×dk\\mathbf\{K\}\_\{l,i\},\\mathbf\{V\}\_\{l,i\}\\in\\mathbb\{R\}^\{n\\times d\_\{k\}\} 是长度为 nn 的序列的共享键和值矩阵，头维度为 dkd\_\{k\}。在因果掩码下，FA 允许每个 token 关注所有前面的 token，每个头组的计算量为 O\(n2\)O\(n^\{2\}\)，KV 缓存为 O\(n\)O\(n\)。

#### 滑动窗口注意力（SWA）。

SWA 将每个 token 限制为只关注最近的 ww 个前面的 token，其中 ww 是固定的窗口大小：

Ol,iSWA=softmax\(Ql,i\(Kl,iw\)⊤dk\)Vl,iw,\\mathbf\{O\}\_\{l,i\}^\{\\text\{SWA\}\}=\\mathrm\{softmax\}\\\!\\left\(\\frac\{\\mathbf\{Q\}\_\{l,i\}\(\\mathbf\{K\}\_\{l,i\}^\{w\}\)^\{\\top\}\}\{\\sqrt\{d\_\{k\}\}\}\\right\)\\mathbf\{V\}\_\{l,i\}^\{w\}, (2) 其中 Kl,iw,Vl,iw∈Rw×dk\\mathbf\{K\}\_\{l,i\}^\{w\},\\mathbf\{V\}\_\{l,i\}^\{w\}\\in\\mathbb\{R\}^\{w\\times d\_\{k\}\} 是仅包含窗口内条目的共享键和值矩阵。当 w≪nw\\ll n 时，KV 缓存从 O\(n\)O\(n\) 减少到 O\(w\)O\(w\)，计算成本从每个头组的 O\(n2\)O\(n^\{2\}\) 降低到 O\(nw\)O\(nw\)，相比 FA 带来了显著的效率提升。

## 4 方法

### 4.1 问题形式化

ConSA 将混合注意力的设计形式化为一个稀疏约束分配问题：给定一个预训练的 Transformer，目标是确定每个 KV 头应该执行全注意力（FA）还是滑动窗口注意力（SWA），使得得到的混合模型满足用户指定的目标稀疏度，同时保持语言建模性能。

令 ρ∈[0,1]\\rho\\in[0,1] 表示目标稀疏度比率，定义为分配给 SWA 的 KV 头比例。对于第 ll 层中的第 ii 个 KV 头，我们引入一个二元分配变量 zl,i∈{0,1}z\_\{l,i\}\\in\\\{0,1\\\}，用于在两种注意力类型之间进行选择。每个 KV 头组的输出则是在两者之间进行硬选择：

O^l,i=zl,i⋅Ol,iFA+(1−zl,i)⋅Ol,iSWA。\\hat\{\\mathbf\{O\}\}\_\{l,i\}=z\_\{l,i\}\\cdot\\mathbf\{O\}\_\{l,i\}^\{\\mathrm\{FA\}\}\+(1\-z\_\{l,i\}\)\\cdot\\mathbf\{O\}\_\{l,i\}^\{\\mathrm\{SWA\}\}。 (3)

ConSA 在两种粒度级别上应用此形式化。*头级别*变体将每个 zl,iz\_\{l,i\} 视为独立变量，允许同一层内的不同 KV 头采用不同的注意力类型。*层级*变体约束一层中的所有 KV 头共享一个分配变量，即对于所有 ii 有 zl,i=zlz\_\{l,i\}=z\_\{l\}，这减少了搜索空间的大小。在头级别分配下诱导的稀疏度比率 ρ^\(z\)\\hat\{\\rho\}\(z\) 为

ρ^\(z\)=ρ^head\(z\)=1−1L⋅HKV∑l=1L∑i=1HKVzl,i，\\hat\{\\rho\}\(z\)=\\hat\{\\rho\}\_\{\\mathrm\{head\}\}\(z\)=1\-\\frac\{1\}\{L\\cdot H\_\{\\mathrm\{KV\}\}\}\\sum\_\{l=1\}^\{L\}\\sum\_\{i=1\}^\{H\_\{\\mathrm\{KV\}\}\}z\_\{l,i\}, (4)

在层级分配下，它简化为

ρ^\(z\)=ρ^layer\(z\)=1−1L∑l=1Lzl，\\hat\{\\rho\}\(z\)=\\hat\{\\rho\}\_\{\\mathrm\{layer\}\}\(z\)=1\-\\frac\{1\}\{L\}\\sum\_\{l=1\}^\{L\}z\_\{l\}, (5)

其中 LL 是层数，HKVH\_\{\\mathrm\{KV\}\} 是每层的 KV 头数。整体优化问题为

minθ,zLLM\(θ,z\)s.t.ρ^\(z\)=ρ，\\min\_\{\\theta,z\}\\;\\mathcal\{L\}\_\{\\mathrm\{LM\}\}\(\\theta,z\)\\quad\\text\{s.t.\}\\quad\\hat\{\\rho\}\(z\)=\\rho, (6)

其中 θ\\theta 表示模型参数，LLM\\mathcal\{L\}\_\{\\mathrm\{LM\}\} 表示自回归语言建模损失。

### 4.2 使用硬具体分布进行可微分掩码学习

由于 zl,i∈{0,1}z\_\{l,i\}\\in\\\{0,1\\\} 是二元的，方程 (6) 不可微，无法直接使用梯度优化。为了联合训练 θ\\theta 和 zz，我们使用硬具体分布 [Louizos et al., 2018](https://arxiv.org/html/2606.18056#bib.bib28) 对每个 zl,iz\_\{l,i\} 进行参数化，该分布将非零概率质量分配给 0 和 1，同时在其间保持连续且可微。我们将得到的 zl,iz\_\{l,i\} 称为可学习