你所不知道的潜藏在Transformer 中 LayerNorm 执行顺序问题

这几天一篇关于Transformer 架构图与其代码实现不一致的问题在网上异常火热，发现错误的是一位知名机器学习与 AI 研究者、初创公司 Lightning AI 的首席 AI 教育家 Sebastian Raschka。他指出，原始 Transformer 论文中的架构图有误，将层归一化（LN）放置在了残差块之间，而这与代码不一致。今天我们就这个问题，进行一下剖析。

关于《Attention is all you need》 这篇paper 中 Transformer 的架构图如下：

Transformer 架构图

对应图的代码实现

官方代码实现

代码地址：https://github.com/tensorflow/tensor2tensor/commit/f5c9b17e617ea9179b7d84d36b1e8162cb369f25#diff-76e2b94ef16871bdbf46bf04dfe7f1477bafb884748f08197c9cf1b10a4dd78e.

现在大家架构图代码实现不一样的争论如下：

Sebastian Raschka 认为架构图中执行：x + LayerNorm(Sublayer(x))

论文对应代码执行: LayeNorm(x + Sublayer(x))

这么看确实存在问题，架构图表达的和实际代码执行的确实不一样的，实际Transformer 架构图，想表达的是一种Post-LayerNorm 操作，即如下图的（a）部分，也即代码的呈现方式。

(a)后置LayerNorm (b)前置layerNorm操作

后置LayerNorm ：LayerNorm(x + SubLayer(x))

前置LayerNorm： x + SubLayer(LayerNorm(x))

论文On Layer Normalization in the Transformer Architecture On Layer Normalization in the Transformer Architecture

这里SubLayer 在encoder 中表示 self-attention或者Feed Forwad Net

在decoder 可以表示为 self-attention或者encoder-decoder attention，Feed Forwad

Sebastian Raschka 认为的执行方式即不是Pre-LayerNorm也不是Post-LayerNorm，这里争议的根源就是对Transformer 架构图的理解，如果将架构图中理解为先对self-attention 进行norm 然后再add，确实不一致；但是笔者的理解是架构图先add subLayer的结果，然后再对整体进行LayerNorm（目前尚未看到原始第一版代码）后续google 将不会再更新架构图.

话说回来， 这个Transformer架构图与代码不一致问题和Pre-LayerNorm 与Post-LayerNorm 的优劣问题是两件不同的事情，原来的探讨将这两件事情放在一起来讲，令人费解。

现在大多数大模型上基本采用Pre-LayerNorm 方式，Pre-LayerNorm 更好优化，能够减小训练时间。