[控场AI]
· 4 分钟阅读· 2,308 字

GAP架构:注意力与MLP合并的Transformer实验

GAP架构:注意力与MLP合并的Transformer实验

实验架构GAP将注意力与MLP合并为门控块,参数省约20%但精度略降。

文章介绍了 Reddit 开发者提出的实验性架构 GAP(Gated Attention Projection),其核心是把 Transformer 中分离的多头注意力子层与前馈网络子层融合进同一个门控块,以共享投影覆盖两类计算,目标是压缩参数量。作者在 WikiText-103 上用 1 亿计分 token 做对照:主模型 GAP20 参数量比标准 Transformer 少 19.91%,但最终语言模型准确率落后;更激进的配置可减 34.31% 参数,但表现下滑明显。torch.compile 让 GAP20 在 CPU 基准上训练吞吐提升约 19.9%,但这只反映编译加速。作者明确将其定位为实验性探索,公开了代码、可配置尺寸、速度基准与论文,并不声称整体优于 Transformer。

Transformer 的标准结构——多头注意力后接前馈网络(FFN)——在过去几年几乎成了默认配方。但也有研究者开始尝试更激进的改动:把这两个子层合并成一个带门控的模块。Reddit 上一位开发者分享的实验性架构 GAP(Gated Attention Projection)正是沿着这个方向做的探索。

reddit 帖子截图

GAP 的核心思路:一个门控块替代两个子层

GAP 的基本主张很直接:不再让注意力层和 MLP 层各自独立堆叠,而是把注意力计算与 MLP 风格的投影处理融合进同一个门控块内。门控机制在这里承担了原本由两层之间的残差连接和归一化隐式完成的部分调节职责。

据作者描述,这个设计的出发点是参数效率——标准 Transformer 中注意力和 FFN 各占相当比例的参数,如果能用一个共享投影的门控结构覆盖两类计算,就有机会压缩总参数量。需要强调的是,作者自己在帖子中明确表态:这是一个实验性架构,并不声称它整体上优于 Transformer。

门控机制(Gating) 在神经网络中通常指用一个可学习的标量或向量对另一路信号做逐元素乘性调制,形式可写作 (y = g \odot x),其中门控值 (g) 由 sigmoid 或 softmax 等函数从输入计算而来。它的作用是让网络动态决定信息通过的比例,从而获得类似“软开关”的选择性。门控在 LSTM 的遗忘门、GRU 的更新门中已有长期应用,近年也被引入 Transformer 变体,例如门控线性单元(GLU)及其派生结构在 FFN 中替代部分激活函数。GAP 把门控放在注意力与投影融合的位置,思路是让同一个乘性调节同时服务于两类计算,而不是像标准结构那样依赖两层之间的残差连接与 LayerNorm 做隐式调节。

实测结果:参数省了,精度没赢

作者用 WikiText-103 做了对照实验,每个模型使用 1 亿个计分训练 token。主要发现可以分几层看:

参数压缩幅度可观。 一个更小的配置能把参数量降低 34.31%,但这个版本相对标准 Transformer 的表现下滑明显,所以作者没有采用它。最终测试的主模型 GAP20 把省下的参数中的约 14% 还给模型用于增加宽度,净效果是比标准 Transformer 少 19.91% 的参数——这也是“GAP20”这个名字的来源。

语言模型准确率上标准 Transformer 更好。 这是作者如实报告的结果,没有回避。GAP20 在最终语言模型精度上落后于对照组,说明参数效率的提升并没有直接转化为建模能力的优势。

训练吞吐有亮点。 在 CPU 基准测试中,用 torch.compile 编译 GAP20 后,训练吞吐相比未编译版本提升了约 19.9%。这个数字只反映编译带来的加速,不是 GAP 相对标准 Transformer 的速度优势,解读时需要区分清楚。

WikiText-103 是语言建模领域常用的中等规模语料,取自维基百科,约 1 亿词级别,常用于评测模型在长程依赖上的建模能力。所谓“计分训练 token”指参与损失计算的 token 数量,通常不等于语料总词数,因为部分 token 可能被掩码或用于特殊处理。以 1 亿计分 token 作为统一预算做对照,是控制训练计算量的常见做法,便于在相同训练量下比较不同架构。需要留意,这一规模的实验属于中等规模验证,结论对超参数、学习率调度和模型宽深比都比较敏感,因此单次对照的结果更适合作为方向性参考,而非定论。

torch.compile 是 PyTorch 2.0 引入的编译前端,通过对计算图做算子融合、内核生成等优化来提升训练与推理速度。它带来的加速来源于编译器对既有计算图的改写,与模型架构本身的设计无关;因此文中约 19.9% 的吞吐提升应理解为“GAP20 在编译后相比未编译的自身更快”,而不是 GAP 相对标准 Transformer 的速度优势。跨架构比较速度时,必须保证编译状态、硬件和批大小一致,否则数字不可直接对比。

架构代码与可配置尺寸

作者在 GitHub 仓库(rickey1990/Gated-Attention-Projection)中放出了架构代码、可配置的模型尺寸、速度基准以及一份研究论文。对想复现或二次实验的人来说,这套材料的完整度算是比较高的——既有可运行实现,也有写作层面的动机说明。

从工程角度看,门控融合结构的一个潜在吸引力在于它对编译器和硬件的友好度:单个块内的计算路径更集中,融合机会更多,这与 torch.compile 带来的吞吐改善可能是相互关联的。不过这只是推测,作者并未在帖子中论证这层因果关系。

这个实验说明了什么

参数减少约 20% 而精度略降,这个 trade-off 是否值得,取决于部署场景对模型体积和推理成本的敏感程度。在边缘设备或需要严格控制显存的场景下,GAP 这类结构的思路有参考价值;但在追求基准精度的训练任务中,标准 Transformer 仍是更稳妥的基线。

作者公开征求对设计、数学推导和改进方向的反馈,这种开放态度本身值得肯定。架构创新很少一次到位,GAP 当前更像是一个有明确数据支撑的假设检验,而非终局方案。

你如何看待注意力与 MLP 的融合设计?如果你的场景对参数量敏感,会考虑尝试这类结构吗?

#GAP架构#Transformer架构#门控注意力#注意力与MLP融合#参数效率#torch.compile#语言模型实验#WikiText-103
分享:

相关推荐