Asymmetric Attention Heads: 面向Transformer注意力的结构化头级上下文分配

# 非对称注意力头：基于结构化逐头上下文分配的Transformer注意力  
来源：https://arxiv.org/html/2608.19203  

###### 摘要  
标准多头注意力（MHA）为每个头提供相同的完整因果上下文范围，尽管各头可能承担不同的上下文角色。一些头可能主要依赖局部的词汇或句法上下文，而另一些头可能依赖更长距离的关系，如实体交互、语篇链接或状态变化。我们提出**非对称注意力头（AAH）**，一种将上下文长度视为每个头或每组显式分配变量的逐头上下文分配框架。AAH通过特征衍生统计量对头进行分组，建立层级结构，并在保持标准扁平MHA输出接口的同时分配因果局部窗口。在种子0的4096 token实验中，几种AAH式的局部分配变体实现了比纯全注意力更低的验证损失。短预算消融实验表明，稳定的局部分配和头窗口分配结构至关重要，而固定/局部控制可以与自适应层级竞争。我们将AAH解释为一种用于质量与分析的结构化逐头上下文分配机制，并报告**注意力覆盖比率（ACR）**作为选定窗口路由诊断指标。  

(a) 标准MHA  
输入隐藏状态  
共享Q/K/V投影  
头  
头  
头  
头  
完整因果范围（每个头）  
标准因果注意力  
扁平头拼接  
输出投影  
标准块输出  
均匀全范围执行  

(b) AAH-v3  
输入隐藏状态  
共享Q/K/V投影  
控制分支  
Q/K/V头特征  
EMA平滑  
层级  
头  
头  
头  
头  
组  
组  
组  
上层层级  
宽关节兄弟节点评分器  
最终逐头窗口索引  
执行分支  
按选定窗口分组头  
稠密掩码  
FlexAttention  
FlashAttention  
散射输出至头顺序  
扁平头拼接  
输出投影  
标准块输出  
窗口执行  
选定窗口  

图1：标准MHA与AAH-v3对比。标准MHA在扁平输出合并前为所有头应用相同的完整因果范围。AAH-v3保持共享Q/K/V投影和扁平Transformer接口，但添加控制分支为每个头分配窗口；执行分支按选定窗口对头分组，可使用稠密掩码、FlexAttention或FlashAttention执行。  

## 1 引言  
Transformer依赖多头注意力（MHA）[1]，其中每个头通常以统一的注意力模式和完整注意力范围执行。此设计简单稳健，但也隐藏了一个结构性问题：不同头可能不需要相同数量的上下文才能做出有用贡献。这引出了本文的核心问题：**Transformer能否在保持（或有时提升）语言建模质量的同时，为各头不均匀地分配上下文窗口？** 我们将这一设计方向称为**非对称注意力头（AAH）**。注意力头仍是标准MHA接口的一部分，但被允许接收不等长的上下文窗口。AAH-v3通过选择离散的因果局部窗口来控制跨头和头组的MHA执行，同时保持标准Transformer块接口和扁平输出合并。该方法改变了**逐头上下文分配策略**，而非Transformer块的输出拓扑。  

我们的评估范围是有意限制的。我们将**注意力覆盖比率（ACR）**[1]视为选定窗口覆盖的代理指标，而非实测FLOPs度量。稠密掩码路径并未建立跳跃后端工作；因此主要行省略了后端工作核算列。真正的GPU FLOPs仅通过Nsight推导的GPU浮点运算计数器测量，当前Nsight证据不支持本文实现实现了FLOPs减少的声明。AAH在狭义上是后端无关的，因为它为每个头选择上下文窗口，然后将这些窗口交给执行路径。当前证据表明此路由结构会影响质量和选定窗口诊断，但未显示实测GPU FLOPs降低。因此，我们将AAH主要视为用于质量和可解释性的逐头上下文分配机制，系统分析作为负向或边界证据报告，而非系统性能声明。  

早期内部变体探索了静态和初步动态不对称，但未纳入本文主要叙述。AAH-v3是本文研究的唯一主要方法。当前声明是结构性的：AAH暴露了非均匀的逐头上文分配，报告的实验测试了该结构在固定4096 token设置下如何影响验证质量和路由诊断。  

我们的贡献如下：  
1. **方法**：我们提出AAH-v3作为MHA级跨头上下文分配机制，使用EMA平滑的头特征、混合层级构建、增强的控制器输入、宽关节兄弟节点评分、自上而下的父约束和分组因果局部注意力。  
2. **结构化分配诊断**：我们报告ACR作为选定窗口路由诊断指标，同时将其与实测GPU FLOPs区分。  
3. **质量与结构证据**：我们报告种子0的4096 token训练/推理结果以及3000步筛选消融，表明几种结构化/局部分配策略实现了比纯全注意力更低的验证损失，而固定/随机控制警示不要将增益仅归因于自适应层级。  
4. **声明边界**：我们报告Nsight和FLOPs实验室诊断，表明当前实现实现未建立实测GPU-FLOPs节约；最终确认需要多种子和匹配的长预算控制。  

## 2 动机与问题诊断  
### 2.1 统一的头执行作为上下文分配僵化  
标准MHA为层中的每个头分配相似的注意力模式和完整因果范围。这很稳健，但当头在专业化、冗余性和依赖上下文的效用上存在差异时，可能结构性僵化。潜在故障模式不仅是成本：统一的全范围规则也可能阻止模型表达稳定的头特定上下文角色。AAH直接针对此诊断：头级上下文预算应可不均匀分配，同时保持标准Transformer块语义。因此，目标不是通用稀疏化或已证实的GPU-FLOPs减少，而是**在多头接口内进行质量感知的上下文预算重分配**。  

### 2.2 为什么需要AAH  
早期静态不对称确认了不均匀分配是可能的，但静态策略在上下文和训练阶段中过于僵化。初步动态不对称提高了灵活性，但对不稳定性和控制噪声敏感。AAH-v3通过使用增强的层级感知特征、关节兄弟节点评分和约束解决决策的分组动态控制来解决此问题。此设计在保持外部Transformer接口不变的同时，在跨头间暴露稳定的MHA级上下文分配策略。  

### 2.3 路由诊断与硬件性能分析  
第二个诊断是方法论上的：选定窗口覆盖不等于实测硬件工作。掩码路径、内核行为、内存流量、启动开销和块调度可能使路由诊断与Nsight测量的GPU浮点操作脱钩。AAH在路由代理下减少了选定注意力覆盖，但当前实现实现并未将此代理转化为更低的实测GPU FLOPs。因此，我们将ACR视为结构诊断，并将AAH主要评估为用于质量和可解释性的逐头上下文分配机制。  

## 3 相关工作  
### 3.1 标准MHA与头冗余  
Transformer引入了缩放点积注意力和多头注意力作为核心序列混合机制[1]。标准MHA在层内所有头上分配相似的执行模式和注意力范围，但经验头剪枝和头专业化分析表明，头在推理时并非同等有用、同等专业化或同等必要[2,3]。AAH基于此观察，但并未剪枝头；它保持标准MHA接口，并询问是否可以不均匀分配执行范围同时保持质量。  

### 3.2 稀疏、局部与路由注意力  
大量工作通过改变注意力模式本身来降低注意力成本，包括稀疏分解、局部/全局稀疏模式和基于内容的路由[4,5,6,7]。AAH定位不同。它**不是**稀疏注意力替代品，也**不是**改变拓扑的输出设计。相反，AAH在标准Transformer块语义内提供**头级执行控制**：输出仍来自通常的扁平头拼接和输出投影，而执行分支为头或头组分配不同的因果局部窗口范围。  

### 3.3 K/V共享与分组查询注意力  
多查询注意力（MQA）和分组查询注意力（GQA）通过在查询头间共享或减少K/V头来降低解码器推理成本[9,10]。AAH-v3是互补而非相同的：它保持完整的注意力头集，并改变每个头或组的选择执行范围。因此，AAH针对注意力区域和有效注意力元素，而MQA/GQA主要针对K/V缓存带宽和K/V头共享。  

### 3.4 自适应计算与路由  
自适应计算方法根据输入调整计算量或位置，从自适应循环计算时间到路由稀疏专家模型[11,12]。路由Transformer类似地学习内容依赖的稀疏注意力模式[7]。AAH共享条件分配的广泛目标，但路由单元是MHA头或头组，决策空间是层级约束下的离散因果窗口范围集。在报告的硬策略运行中，执行的窗口策略最好理解为特征条件化和统计驱动，而非已证实的可微学习路由器。自适应注意力跨度学习每个头的最大注意力跨度，并使用可微掩码机制和跨度长度惩罚来减少长上下文Transformer的内存和计算[8]。AAH-v3在精神上相关，因为它在头上分配不等长注意力跨度，但不同之处在于使用单独的层级感知控制分支、组级决策、父约束和分组局部执行，同时保持标准的扁平MHA输出接口。  

### 3.5 系统感知的注意力执行与分析  
IO感知的精确注意力内核（如FlashAttention和FlashAttention-2）表明，注意力运行时强烈依赖于内存流量、分块、并行性、工作划分和内核实现，而不仅是名义注意力元素[13,14]。可编程的局部注意力后端（如PyTorch FlexAttention）提供了在后端可执行形式中表达稀疏或局部注意力模式的互补路径[15]。AAH不是内核优化方法。FlashAttention/FlexAttention和AAH解决栈中的不同层：Flash/Flex优化或表达注意力执行，而AAH选择不对称的逐头上文窗口。因此我们报告ACR作为路由诊断指标，并将GPU-FLOPs声明保留给Nsight推导的浮点运算计数器。  

## 4 背景：标准多头注意力  
设 $X \in \mathbb{R}^{T \times d}$ 表示输入序列表示，其中 $T$ 是序列长度或最大因果范围，$d$ 是模型维度。对于每个注意力头 $h \in \{1, \dots, H\}$，其中 $H$ 是头数，  
$$ Q_h = X W_h^Q, \quad K_h = X W_h^K, \quad V_h = X W_h^V. $$  
其中 $Q_h$、$K_h$ 和 $V_h$ 是头 $h$ 的查询、键和值投影，$W_h^Q$、$W_h^K$ 和 $W_h^V$ 是对应的投影矩阵。头输出 $O_h$ 为  
$$ O_h = \operatorname{softmax}\!\left(\frac{Q_h K_h^\top}{\sqrt{d_h}} + M\right) V_h, $$  
其中 $d_h$ 是头维度，$M$ 是因果掩码。最终的多头注意力输出 $O$ 为  
$$ O = \operatorname{Concat}(O_1, \dots, O_H) W^O, $$  
其中 $W^O$ 是输出投影矩阵。  

此公式有两个重要特性。首先，所有头并行操作于相同层输入。其次，所有头在序列上使用相同的全宽注意力模式。换言之，标准MHA在头上分配**均匀的注意力计算**。如果每个头关注完整因果上下文，主导的注意力分数计算与 $\sum_{h=1}^H T^2 d_h$ 成比例（忽略常数因子和非注意力项）。这种对序列长度的二次依赖使注意力成为研究执行感知上下文分配的自然目标。  

## 5 方法：AAH-v3  
### 5.1 方法概述  
AAH-v3通过在头上和头组上进行层级上下文分配来增强标准多头注意力。它保持Transformer块约定：模型仍形成 $Q/K/V$ 投影，在因果掩码下计算缩放点积注意力，拼接头输出，并应用通常的输出投影。改变的是MHA上下文策略：不同的头可使用由层级感知控制器选择的因果局部窗口。在报告的硬策略实现中，控制器/评分器状态从种子AAH状态固定，并基于特征统计操作；语言建模损失训练执行的Transformer路径，而非硬窗口选择的可微路由目标。  

概念上，AAH-v3有两个主要部分。**头分组**总结当前 $Q/K/V$ 行为和最近注意力诊断，随时间平滑这些信号，形成0级头组，并构建或重用决定哪些头被共同控制的上层层级。**窗口控制**则使用该分组层级来选择并执行逐头上下文窗口。其选择阶段构建增强的控制器输入，应用关节兄弟节点评分，实施父约束，并将组决策映射回头。其执行阶段按选定窗口对头分组，并在通常的散射、扁平合并和输出投影之前运行所选注意力后端。  

从早期草稿到当前版本的中心控制器转变是从独立组评分转向**关节兄弟节点**评分。