arXiv:2509.22840cs.LG2025-09被引 2

解释多头注意力为何能提升模型容量,关键在减少嵌入重叠干扰。

A Capacity-Based Rationale for Multi-Head Attention

  • 通过关系图识别任务分析注意力通道容量,量化编码能力。
  • 固定预算下,多头设计使可区分的关联数提升,与理论预测一致。
  • 适合关注注意力机制原理、模型容量的读者,尤其对结构设计敏感者。

我们研究自注意力中键-查询通道的容量:在固定预算下,单层能可靠编码多少个不同的词元-词元关系?引入关系图识别任务,其中键-查询通道需编码有向图,并在给定上下文(顶点子集)时恢复每个顶点的邻居。资源以总键维度 $D_K = h imes d_k$ 衡量。在可处理的多头模型中,我们通过显式构造证明了信息论下界与上界匹配,表明在 $d_{\text{model}}$ 维嵌入中恢复 $m'$ 个关系所需的 $D_K$ 随 $m'/d_{\text{model}}$ 增长(含对数因子),并给出缩放softmax注意力的相应保证。该分析为多头注意力提供了新的容量依据:即使在排列图(所有查询仅关注单一目标)中,将固定 $D_K$ 分配到多个头仍能通过减少嵌入叠加干扰提升容量。受控实验验证理论预测,显示出明确的容量相变,且多头优势在加入softmax归一化、值路由及使用冻结GPT-2嵌入训练完整Transformer块后依然存在。

原文摘要 · Abstract (English)

We study the capacity of the self-attention key-query channel: for a fixed budget, how many distinct token-token relations can a single layer reliably encode? We introduce Relational Graph Recognition, where the key-query channel encodes a directed graph and, given a context (a subset of the vertices), must recover the neighbors of each vertex in the context. We measure resources by the total key dimension $D_K = h\,d_k$. In a tractable multi-head model, we prove matching information-theoretic lower bounds and upper bounds via explicit constructions showing that recovering a graph with $m'$ relations in $d_{\text{model}}$-dimensional embeddings requires $D_K$ to grow essentially as $m'/d_{\text{model}}$ up to logarithmic factors, and we obtain corresponding guarantees for scaled-softmax attention. This analysis yields a new, capacity-based rationale for multi-head attention: even in permutation graphs, where all queries attend to a single target, splitting a fixed $D_K$ budget into multiple heads increases capacity by reducing interference from embedding superposition. Controlled experiments mirror the theory, revealing sharp phase transitions at the predicted capacity, and the multi-head advantage persists when adding softmax normalization, value routing, and a full Transformer block trained with frozen GPT-2 embeddings.

注意力机制模型容量多头注意力

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