论文

Prism Transformer:用于分层注意力处理的渐进式头部调度

Prism Transformer: Progressive Head Schedules for Hierarchical Attention Processing

模型架构Transformer

摘要

多头注意力通常会在每一层的所有头之间均匀地划分隐藏维度,从而在整个模型深度上强制执行相同的表征子空间维度(dh = dmodel/h)。在这项工作中,我们将这种统一分配视为一个基本的结构瓶颈:由于其有限的维度空间,早期层的头部无法忠实地捕获复杂的高维上下文模式。为了解决这个问题,我们引入了 Prism Transformer,这是一种新颖的架构范例,用渐进式头部时间表取代了静态、统一的头部配置。通过单调增加各层的头数,Prism Transformer 自然地建立了本地到全局的表示层次结构:早期层利用较少且异常宽的头来捕获复杂的局部组成模式,而深层则部署许多狭窄的头来将这些模式分解为专门的语言特征。至关重要的是,这种结构转变是参数中立、计算中立的,并且引入了零训练或推理开销,保留了与标准 Transformer 相同的权重矩阵和 FLOP 预算。在三个模型规模(124M、354M 和 757M)中,Prism Transformer 始终优于统一基线,实现了验证损失的持续减少,同时下游零样本基准(包括 PIQA、HellaSwag、ARC-Easy 和 WinoGrande)的持续收益。我们的研究结果表明,非均匀子空间分配可释放标准 Transformer 预算内的潜在容量,从而更有效地利用模型容量。