让注意力头数随深度递增,提升模型捕捉复杂模式的能力。
Prism Transformer: Progressive Head Schedules for Hierarchical Attention Processing

- 早期层用更少但更宽的注意力头,捕捉局部复杂模式。
- 深层用更多窄头分解语义特征,实现从局部到全局的表征层次。
- 无需额外参数或计算量,即可显著提升零样本任务表现。
多头注意力通常在每层均分隐藏维度,导致所有层的子空间维度相同,形成结构性瓶颈:早期层因维度受限,难以捕捉复杂的高维上下文模式。为此,我们提出Prism Transformer,采用渐进式头数调度机制,逐层增加注意力头数,构建局部到全局的表征层次。早期层使用少量超宽头捕捉复杂局部组合模式,深层则通过大量窄头分解为专业化语言特征。该结构保持参数与计算量不变,不增加训练或推理开销。在124M、354M和757M三个规模模型上,普里姆斯变压器均优于均匀基线,验证集损失持续下降,并在多个零样本基准(包括PIQA、HellaSwag、ARC-Easy、WinoGrande)上取得一致提升。结果表明,非均匀子空间分配可释放标准Transformer预算下的潜在容量,更高效利用模型能力。
原文摘要 · Abstract (English)
Multi-head attention conventionally partitions the hidden dimension equally across all heads at every layer, enforcing an identical representational subspace dimension (dh = dmodel/h) throughout the models depth. In this work, we identify this uniform allocation as a fundamental structural bottleneck: due to their restricted dimensional space, early-layer heads are unable to faithfully capture complex, high-dimensional contextual patterns. To resolve this, we introduce the Prism Transformer, a novel architectural paradigm that replaces the static, uniform head configuration with a progressive head schedule. By monotonically increasing the head count across layers, the Prism Transformer naturally establishes a local-to-global representational hierarchy: early layers leverage fewer, exceptionally wide heads to capture complex, local compositional patterns, while deep layers deploy many, narrow heads to decompose these patterns into specialized linguistic features. Crucially, this structural shift is parameter-neutral, compute-neutral, and introduces zero training or inference overhead, preserving identical weight matrices and FLOP budgets as the standard Transformer. Across three model scales (124M, 354M, and 757M), the Prism Transformer consistently outperforms uniform baselines, achieving consistent reductions in validation loss alongside consistent gains on downstream zero-shot benchmarks (including PIQA, HellaSwag, ARC-Easy, and WinoGrande). Our findings demonstrate that non-uniform subspace allocation unlocks latent capacity within the standard Transformer budget, enabling more effective use of model capacity.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。