transformer结构

60 阅读3分钟

架构总览

  最经典的transformer的网络设计采用Encoder-Decoder架构,Encoder处理源序列,编码整个序列供Decoder查询,一次推理过程通常只需进行一次encode。Decoder通过Crosss-Attention机制关注Encoder的输出,逐词生成目标序列

Imput Embedding (词嵌入)

  Embedding 把输入的 token id 序列 (B, L) 映射成 稠密向量 (B, L, d_model)。实现上是查表:权重矩阵形状 (vocab_size, d_model),每个 token id 取出对应行。数学上等价于 one-hot 向量乘权重矩阵,但用查表实现更高效。

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

Generator

  Generator 是 Decoder 最后的输出层,负责把 Transformer 的隐状态映射成词表上的分数。具体来说,它把输入 (B, L, d_model) 投影到 (B, L, vocab_size),得到每个位置、每个 token 的 logits。

class Generator(nn.Module):
  "Define standard linear + softmax generation step."

  def __init__(self, d_model, vocab):
      super(Generator, self).__init__()
      self.proj = nn.Linear(d_model, vocab)

  def forward(self, x):
      return log_softmax(self.proj(x), dim=-1)

  实际操作是先经过一个线性层 nn.Linear(d_model, vocab_size),将维度从 (B, L, d_model) 变为 (B, L, vocab_size);之后再在最后一维(vocab 维)做 softmax,得到 token 的概率分布。

  为了压缩参数量,通常让 Generator 使用 Embedding 权重 W(形状 (vocab, d_model))的转置 W^T(形状 (d_model, vocab))作为线性层权重。共享权重时,Generator 的线性层需要设置 bias=False,否则偏置无法与 Embedding 共享。

Generator 的实际计算是 logits = h @ W^T,其中:

  • input 是隐状态,形状 (B, L, d_model);
  • W 是 Embedding 权重,形状 (vocab, d_model);
  • W^T 的形状是 (d_model, vocab),每一列对应一个 token 的 embedding 向量。

因此:

logits[i,j]=input[i,:]⋅W[:,j]\text{logits}[i,j] = input[i, :] \cdot W[:,j]

即隐状态向量 input[i, :] 与 token j 的 embedding 向量做内积。内积越大,说明隐状态向量input[i,:] 与 token j 的语义越匹配。最后对 logits在最后一维做 softmax,就得到每个 token 的概率分布。

共享权重有效的直觉是:如果 Embedding 学到的语义空间能表示 token 的含义,那么在判断“当前上下文最匹配哪个 token”时,就可以复用同一个语义空间。形状对称只是能共享的前提,语义空间的复用才是共享有效的原因。

LayerNorm

  LayerNorm是对每一个样本独立做归一化,用gamma和beta两个可学习的参数缩放平移分布:

μ=1D∑i=1Dxi,σ2=1D∑i=1D(xi−μ)2\mu = \frac{1}{D} \sum_{i=1}^{D} x_i, \quad \sigma^2 = \frac{1}{D} \sum_{i=1}^{D} (x_i - \mu)^2

然后归一化:

x^i=xi−μσ2+ϵ\hat{x}_i = \frac{x_i - \mu}{\sqrt{\sigma^2 + \epsilon}}

最后用两个可学习参数缩放和平移:

yi=γix^i+βiy_i = \gamma_i \hat{x}_i + \beta_i

其中 gamma 和 beta 形状都是 (D,),每个特征维度有独立的参数。输出形状仍是 (B, L, D),维度不变。

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."

    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

LayerNorm 的核心特点是:统计量只依赖当前样本自己,和 batch 无关。因此 batch=1 也能正常工作,训练和推理行为一致,适合变长序列和小 batch 场景,这也是 Transformer 选择它而不是 BatchNorm 的原因。

Feed Forward

FeedForward(FFN)是 Transformer block 里的第二个子层,紧跟在 self-attention 之后。它由两层线性层和一个激活函数组成,对每个位置独立作用。

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."

    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(self.w_1(x).relu()))

结构如下:

FFN(x)=W2⋅Activation(W1x+b1)+b2\text{FFN}(x) = W_2 \cdot \text{Activation}(W_1 x + b_1) + b_2
  • 输入 x 形状 (B, L, d_model);
  • W_1 把维度从 d_model 升到 d_ff;
  • 激活函数通常是 ReLU 或 GELU;
  • W_2 把维度从 d_ff 降回 d_model;
  • 输出形状仍是 (B, L, d_model)。

通常 d_ff = 4 * d_model,先升维再降维。它只做逐位置的非线性变换,不跨位置交互,跨位置交互由 self-attention 负责。