动机

多头注意力是 Transformer 的核心组件，它通过引入多组 QKV 投影来捕获不同的特征子空间，从而在机器翻译、问答等任务中取得巨大成功。研究表明，注意力头的数量对 Transformer 性能至关重要：在保证每个头的隐藏维度充分大的前提下，注意力头数越多可以使模型效果越好。但问题在于，直接增加头数或维度往往伴随着模型参数量和计算开销的剧增，这在训练和部署中代价高昂

在这里插入图片描述

目前也有一些注意力架构旨在提高计算效率，例如共享部分 K 和 V 的 MQA、GQA；使用矩阵分解的 MLA、MFA、TPA 等。但这些方法主要关注降低内存/计算成本，而非提升注意力的表达能力

于是作者希望在不显著增加参数的前提下，设计一种新的注意力架构，实现近似于使用了更多注意力头和更高每头维度的性能提升

本文方法

本文提出 SAS（Simulated Attention Score，模拟注意力分数），核心思想是在注意力计算中引入额外的映射层，将低维的头表示投射到更高维空间，以此“虚拟地”增大注意力头数和每头的隐藏维度

一、扩展注意力头

对于查询Q，其特征维度为 [B, T, H, D]，分别表示 batch_size，序列长度，头数和隐藏维度。为了扩充 H，需要把其他维度拉平，得到张量 Q_0，维度为 [B * T * D, H] ；然后使用一个 H * H’ 的线性变换得到 Q_1，维度为 [B * T * D, H’]，其中 H’ > H；Q_1 过一个 ReLU 引入非线性；最后再过一个 H’ * H’ 的线性层，并加上 Q_1 的残差连接

在这里插入图片描述

于是我们获得了更多的注意力头，其中残差连接的引入可以稳定训练；值得注意的是，原始头数 H 和扩展后的头数 H’ 都远小于每头的特征维度 D，所以这个两层 MLP 的参数开销相对整模型来说可以忽略不计

除了使用 MLP 来扩展维度，作者还尝试了卷积方案。具体地，将查询 Q 的维度整理成 [B * T, H, D]，类似于多通道特征图，然后使用卷积变换将 H 扩展成 H’，同样地，H’ > H，最后再过第二层卷积以及残差连接

在这里插入图片描述

类似地，在 K、V 中都应用上述扩展流程

二、扩展注意力维度

直觉上，每个注意力头内部特征维度 D 越大，其能够捕获的子空间信息越丰富。因此作者进一步在 Q 和 K 上也引入了类似的维度扩展映射。这里之所以不对 V 进行扩展，是因为 V

直接决定了注意力模块的输出张量隐藏维度，扩大 V 的每头维度到 D 会导致后续前馈层的参数量大幅增加，违背了不显著增加计算量的初衷

在这里插入图片描述

三、注意力聚合

在标准多头注意力中，会将所有头的输出向量拼接，再通过一个输出投影矩阵 O 映射回模型的隐藏维度。然而，由于 SAS 对注意力头数进行了扩增，若仍按传统方式拼接势必导致输出维度变大，进而导致 O 的参数量大大增加（H * hidden 变为 H’ * hidden）。为此，作者提出了参数高效注意力聚合机制，旨在不增加输出层参数规模的情况下完成对多头输出的整合

实现过程非常简单：假设注意力头数扩展了 r 倍，即 r * H = H’，那么便把所有头划分成 r 组，每组都按照原本的计算流程与 O 相乘，得到 r 组输出结果，最后取平均作为注意力模块的最终输出传向前馈层

在这里插入图片描述