提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

Transformer原理

传统循环模型的劣势

• 传统循环模型，比如 RNN，通常按照输入/输出序列中的 token 位置一步步计算。
• 对于序列中第 t 个位置，模型会生成一个隐藏状态 h_t。
• h_t 的计算依赖两部分：
  • 上一个隐藏状态 h_{t-1}
  • 当前第 t 个位置的输入
• 所以 RNN 的计算形式可以理解为：
• h_t = f(h_{t-1}, x_t)
• 由于当前步骤依赖前一步结果，RNN 必须按顺序计算：
• h_1 -> h_2 -> h_3 -> … -> h_t
• 这种顺序依赖导致一个问题：
  同一个训练样本内部不能很好地并行计算。
• 当序列变长时，这个问题会更严重，因为模型必须一步一步处理长序列。
  

整体模型架构

• Transformer 用的是经典的 Encoder-Decoder 架构。
• Encoder 负责把输入序列：
  x1, x2, …, xn
  编码成连续表示：
  z1, z2, …, zn
• Decoder 根据 encoder 的输出 z，一步步生成目标序列：
  y1, y2, …, ym
• Decoder 是 auto-regressive 的：(之前的输出会作为当前的输入)
  • 生成当前位置时，只能看之前已经生成的 token
  • 不能提前看到未来 token
• Transformer 的核心特点是：
  • 不使用 RNN
  • 不使用 CNN
  • 主要依赖 self-attention 和 feed-forward network
    • self-attention，是一种将单个序列的不同位置联系在一起的注意机制，以便计算该序列的表示。

Encoder 结构

• Encoder 由 N = 6 层相同的 layer 堆叠而成。
• 每一层有两个子层：
  • Multi-Head Self-Attention
  • Position-wise Feed-Forward Network
• 每个子层外面都有：
  • residual connection
  • layer normalization
• 形式可以写成：
  LayerNorm(x + Sublayer(x))
• 模型中大多数表示维度都是：
  d_model = 512

Decoder 结构

• Decoder 也由 N = 6 层相同的 layer 堆叠而成。
• 每一层有三个子层：
  • Masked Multi-Head Self-Attention
  • Encoder-Decoder Attention
  • Position-wise Feed-Forward Network
• Decoder 也使用 residual connection 和 layer normalization。
• Decoder 的 self-attention 要加 mask：
  • 防止当前位置看到未来位置
  • 保证生成第 i 个 token 时，只能依赖第 i 之前的输出

Attention 部分

• Attention 可以理解为：
  • 给定一个 query，去一组 key-value 里找相关信息。
  • 输出是 value 的加权和。
  • 权重由 query 和 key 的相似度决定。
• Transformer 使用的是：
  Scaled Dot-Product Attention
• 公式是：
  

除以 sqrt(d_k) 是为了防止点积数值太大，导致 softmax 梯度太小。
Multi-Head Attention

• 不是只做一次 attention，而是并行做多次 attention。
• 每个 attention head 可以关注不同类型的信息。
• 多个 head 的结果会拼接起来，再做一次线性变换。
• 好处是模型可以同时学习不同位置、不同语义空间里的关系。

Feed-Forward Network

• 每个位置都会经过同一个前馈网络。
• 它对每个 token 独立作用，不混合不同位置的信息。
• 结构大致是：
  Linear -> ReLU -> Linear

Embedding 和 Softmax

• 输入 token 先被转成 embedding。
• Decoder 输出最终通过 linear + softmax 预测下一个 token。
• 输入 embedding 和输出 softmax 的权重共享。

Positional Encoding

• 因为 Transformer 没有 RNN/CNN，本身不知道 token 的顺序。
• 所以作者加入 positional encoding 来表示位置信息。
• positional encoding 会和 token embedding 相加。
• 作者使用了正弦、余弦函数构造位置编码：
  • 偶数维用 sin
  • 奇数维用 cos
• 这样模型可以利用相对位置和绝对位置信息。

Self-Attention

算法流程

怎么计算两个向量之间的关联性呢，也就是如何计算的α？
以下就是常见的计算α两种方式。例如左边的：输入向量分别乘上W(q)和W(k) 矩阵,得到q和k两个向量。然后再将q和k向量进行点乘得到α

分别计算a1与a2、a3和a4的关联性。

计算出a1和每个向量的关联性之后，会做一个softmax得到α’。然后会根据得到的α’，得出哪些向量跟a1是最有关联的。接着会根据关联性来抽取重要的信息

怎么抽取重要的信息呢？
把输入的每个向量都乘上一个W(v)得到新的向量，也就是下图的v1-v4。然后再将v1和v4都去乘上对应的α’，然后将他们加起来，就能得到考虑全局信息的b向量了。
如果某一个向量得到的分数越高，比如a1和a2的关联性很强也就是a’1,2的值很大，那么得到的b1的值就可能比较接近于v2。也就是哪个向量的attention的分数越大，那么这个向量的v就会dominate抽出来的结果。

后续会计算 q2与k1，k2，k3，k4、q3与k1，k2，k3，k4、q4与k1，k2，k3，k4的计算结果

矩阵计算

why Self-Attention

Encoder and Decoder Stacks

为啥使用layer normalization而不是batch normalization

核心原因：Transformer 处理的是变长序列，而且训练/推理时 batch 形态不稳定，LayerNorm 更适合序列建模。

1. BatchNorm 依赖 batch 统计量，不适合变长序列
   Batch Normalization是对一个 batch 里的样本做归一化。它需要统计当前 batch 的均值和方差：mean / variance over batch

但 NLP 任务里的 batch 很麻烦：

• 句子长度不同
• padding 数量不同
• 每个 batch 的 token 分布差异大
• batch size 可能因为显存限制而变化
  这会导致 BatchNorm 的统计量不稳定。
  而 LayerNorm 是对单个样本内部的 hidden dimension 做归一化，不依赖其他样本：
  mean / variance over hidden features
  所以它不关心 batch size，也不太受句长变化影响。
  可以记成：
  BatchNorm：跨样本归一化，依赖 batch，适合 CNN
  LayerNorm：样本内归一化，不依赖 batch，适合 Transformer/RNN

Post-Norm vs Pre-Norm

LayerNorm 放在子层的前面还是后面，是一个看似微小但影响深远的设计选择。
Post-Norm（原始论文方案）
$$output=LayerNorm(x+SubLayer(x))$$
先做子层计算和残差加法，再做归一化。
Pre-Norm（当前大模型的主流选择）
$$output=x+SubLayer(LayerNorm(x))$$
先做归一化，再做子层计算，最后做残差加法。

两者的关键区别在于残差路径是否经过 LayerNorm：

• Post-Norm 中，残差路径上有 LayerNorm，梯度回传时会被 LayerNorm 的导数”调制”
• Pre-Norm 中，残差路径是一条”干净”的直通通道，梯度可以无损地回传

为什么大模型普遍用 Pre-Norm

实际工程中，当模型规模达到数十亿参数以上时，训练稳定性远比理论上 0.x% 的效果差异重要——训练一次的成本可能是数百万美元，任何导致 loss spike 或 diverge 的风险都不可接受。这就是 Pre-Norm 成为默认选择的原因。

decoder为啥使用Masked Multi-Head Self-Attention

训练时不能让模型看到未来 token，否则它会作弊。
以机器翻译为例，假设目标句子是：
I love you
Decoder 在生成第 1 个词 I 的时候，不能提前看到 love 和 you。
生成第 2 个词 love 的时候，只能看到前面的 I，不能看到后面的 you。
所以 decoder 的 self-attention 必须加 mask。

• 普通 self-attention 会让每个位置都看到整个序列：

I      可以看 I, love, you
love   可以看 I, love, you
you    可以看 I, love, you

这对 encoder 是没问题的，因为 encoder 本来就是读完整输入句子。
但 decoder 是生成器，必须保持自回归性质：
P(y) = P(y1) P(y2 | y1) P(y3 | y1, y2)
所以 decoder 应该是：

I      只能看 I
love   只能看 I, love
you    只能看 I, love, you

如何加mask呢
Decoder 的 mask 是在 attention score 上把未来位置加成 -inf(非常大的负数)，softmax 后这些位置权重为 0，从而保证自回归生成不能偷看未来。

为什么还是 Multi-Head？
因为 decoder 也需要从已生成的上下文里捕捉不同关系：

• 有的 head 关注前一个词
• 有的 head 关注主谓关系
• 有的 head 关注长距离依赖
• 有的 head 关注短语边界

Attention

Attention 的基本思想是：给定一个 query，去一组 key-value pairs 里查找相关信息，最后输出 value 的加权和。权重由 query 和 key 的相似度决定。
可以类比成查资料：

Query：我现在想找什么信息
Key：每条信息的索引/标签
Value：真正要取出来的内容

模型会先计算 query 和每个 key 的相关性，再根据相关性对 value 做加权求和。

Scaled Dot-Product Attention

Transformer 使用的 attention 形式叫 Scaled Dot-Product Attention。
公式是：

含义是：

• Q 是 queries
• K 是 keys
• V 是 values
• S=QK^T：用 Q 和 K 的内积来衡量每对 token 之间的”匹配度”：
  • 得到的S是一个N×N 的矩阵，S[i][j]表示第i个 token 对第j个 token 的关注程度（原始分数）。
• A=softmax(Sscaled)把相似度变成权重
  • A[i][j] 表示：第 i个 token 分配给第 j 个 token 的注意力权重，每行之和为 1。
• 最后乘以 V，最终每个 token 得到一个 d维向量，其中融合了它”应该关注”的所有其他 token 的信息。
  这里最关键的是除以：sqrt(d_k)
  原因是：如果 key/query 的维度 d_k 很大，点积结果可能会变得很大(因为是d_k 个分量相加)，导致softmax进入梯度很小的区域，训练变慢或不稳定。
  所以作者用 sqrt(d_k) 做缩放，让数值更稳定。

Multi-Head Attention

作者发现，只做一次 attention 不如做多次 attention 并行效果好，于是提出 Multi-Head Attention。
它的做法是：

1. 把 Q, K, V 分别投影到多个不同的子空间
2. 每个子空间独立做一次 attention
3. 把多个 attention head 的结果拼接起来
4. 再经过一次线性变换得到最终输出
   
   直观理解：
   一个 attention head 只能从一个角度看句子，多头 attention 可以让模型从多个角度同时看。
   比如：

• 一个 head 关注相邻词
• 一个 head 关注主谓关系
• 一个 head 关注代词指代
• 一个 head 关注长距离依赖
  论文里 base model 使用：

h = 8
d_model = 512
d_k = d_v = 64

也就是说，512 维表示被分成 8 个 head，每个 head 处理 64 维。
多头机制的参数组成：

W_Q:(4096,4096),即4096×4096=16M参数
W_K:(4096,4096),16M参数
W_V:(4096,4096),16M参数
W_O:(4096,4096),16M参数
合计:4·dmodel = 64M参数

Attention 在 Transformer 中的三种用法

- 第一种是 encoder-decoder attention
这里 query 来自 decoder，key 和 value 来自 encoder 输出：
作用是让 decoder 在生成目标句子时，能够关注输入句子的相关部分。
这和传统机器翻译里的 attention 类似。

- 第二种是 encoder self-attention
encoder 中的 Q, K, V 都来自 encoder 的上一层输出。
作用是让输入序列中每个位置都能关注输入序列中的其他位置。
比如翻译一句话时，某个词可以直接看见句子里和它相关的其他词。

- 第三种是 decoder masked self-attention
decoder 中的 Q, K, V 都来自 decoder 自己，但要加 mask。
作用是保证生成第 i 个位置时，只能看到第 i 个位置之前的信息，不能看到未来 token。
这是为了保持自回归生成的因果性，防止训练时“偷看答案”。

Position-wise Feed-Forward Networks

这个模块叫 position-wise，意思是：它对每个位置的 token 单独处理。
假设一句话有 5 个 token：
x1, x2, x3, x4, x5
FFN 会分别作用在每个位置上：
FFN(x1), FFN(x2), FFN(x3), FFN(x4), FFN(x5)
但是它不会在 FFN 里让 x1 和 x2 互相交流。
也就是说：

• attention 负责 token 之间的信息交互
• FFN 负责对每个 token 自己的表示做非线性变换和特征加工
  可以这么理解：

Attention：让每个词看其他词
FFN：消化刚刚看到的信息

为什么需要 FFN？
如果只有 attention，模型主要是在做信息加权混合。
但神经网络还需要非线性变换能力，否则表达能力不够。
FFN 的作用就是增加模型的非线性表达能力，让模型能对 attention 得到的信息进一步加工。

为什么FFN是模型参数的大头

粗略地看，FFN 的参数量大约是 Attention 的 2 倍。在整个 Transformer Block 中，FFN 贡献了约 2/3 的参数。

Embeddings and Softmax

在 Transformer 里，输入和输出 token 本来都是离散的词或子词，比如：
I love you
模型不能直接处理字符串，所以要先把每个 token 转成向量，这一步就是 embedding。
比如：

I      -> [0.2, -0.1, ...]
love   -> [0.5,  0.7, ...]
you    -> [-0.3, 0.4, ...]

论文中 embedding 的维度和模型主维度一致：d_model = 512
所以每个 token 会被映射成一个 512 维向量。
Decoder 最后会输出一个隐藏向量。
但我们最终要的是“下一个词是什么”的概率分布。
所以需要经过两步：
hidden vector -> linear layer -> softmax
具体来说：

1. Linear layer：把 hidden vector 映射到词表大小的 logits
2. Softmax：把 logits 转成每个词的概率
   如果词表大小是 30,000，那么 decoder 每个位置最后会输出：30000 个分数
   softmax 后变成：每个词作为下一个 token 的概率
   然后模型选择概率最高的词，或者用 beam search 等方法生成句子。
   作者把 embedding 乘以了：sqrt(d_model)
   embedding = embedding * sqrt(d_model)
   原因是 embedding 的数值尺度可能比较小，而后面要和 positional encoding 相加。
   乘以 sqrt(d_model) 可以调整 embedding 的尺度，让它和位置编码的数值规模更匹配。

positional encoding

Transformer 怎么表示词的顺序呢？
这是 Transformer 里非常关键的一点，因为它没有 RNN，也没有 CNN。
RNN 是按时间步一个词一个词读，所以天然知道顺序；CNN 有卷积窗口，也能感知局部顺序。
但 self-attention 本身是“集合式”的：如果没有额外信息，它只知道有哪些 token，不知道谁在前谁在后。
比如：

I love you
You love I

如果只看 token 集合，它们包含的词差不多，但顺序完全不同，语义也不同。
所以 Transformer 必须额外加入位置信息。

作者的做法是给每个位置生成一个 positional encoding，然后和 token embedding 相加：
input = token_embedding + positional_encoding
也就是说，每个 token 的输入向量里同时包含两类信息：
词本身的语义信息 + 这个词所在的位置
例如

第 1 个词: embedding("I") + PE(1)
第 2 个词: embedding("love") + PE(2)
第 3 个词: embedding("you") + PE(3)

这样 attention 在计算时，就能区分不同位置。
论文使用的是基于 sin 和 cos 的固定位置编码。

• pos 表示 token 在序列中的位置
• i 表示向量的维度索引
• d_model 是模型维度，比如 512
• 偶数维用 sin
• 奇数维用 cos

直观理解：
每个位置会被编码成一个由很多不同频率的正弦、余弦波组成的向量。

低维和高维对应不同频率：

• 有些维度变化很快，适合区分近距离位置
• 有些维度变化很慢，适合表达远距离位置

为什么用 sin/cos？
一个重要原因是它可以帮助模型学习相对位置关系。
论文里提到，对于任意固定偏移 k，PE(pos + k) 可以表示成 PE(pos) 的线性函数。
这意味着模型更容易学到：

当前位置后面第 k 个词
当前位置前面第 k 个词

这种相对位置信息。
比如模型不只知道“这个词在第 5 个位置”，还更容易学到“这个词和另一个词相隔 3 个位置”。

Pytorch从零搭建Transformer

import numpy as np
import torch
import torch.nn as nn

掩码机制：三种Mask的实现

掩码是实现中最容易出错的地方。一共有三种，对应第 2 节中描述的三种注意力机制。
Encoder 长度 Mask：同一 batch 中不同样本经过 padding 进行长度补齐，需要屏蔽 padding 位置的 Attention 权重。做法是将对应位置的 score 填为 -1e4，过 softmax 后权重趋近于 0。

'''
    生成 Encoder自注意力的长度掩码。
    Returns: shape (b, max_len, max_len),True 表示该位置需要被屏蔽

    b: batch size
    max_len: batch 中 padding 后的最大长度
    feat_lens: 每个样本的真实长度,shape通常是 (b)
    device: mask 放在哪个设备上,比如 cpu / cuda
'''
device1 = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def get_len_mask(b:int, max_len: int, feat_lens: torch.Tensor, device: torch.device) -> torch.Tensor:
    #先创建全 1 的 mask,每个样本都有一个 max_len × max_len 的 attention mask。
    #1 表示需要 mask / 不允许看
    attn_mask = torch.ones((b, max_len, max_len),  device = device)
    #对第 i 个样本，把前 feat_lens[i] 列设成 0。
    for i in range(b):
        attn_mask[i, :, :feat_lens[i]] = 0
    #把0和1转换为False和True
    return attn_mask.to(torch.bool) 

Decoder Causal Mask（因果掩码）：训练时整个目标序列并行输入，但预测位置 i 只能看到位置 0到 i，通过上三角掩码实现。

def get_subsequent_mask(b: int, max_len: int, device: torch.device) -> torch.Tensor:
    """生成 Decoder 自注意力的因果掩码（上三角矩阵）。
    Returns: shape (b, max_len, max_len)，True 表示未来位置（需屏蔽）
    """
    return torch.triu(torch.ones((b, max_len, max_len), device=device), diagonal=1).to(torch.bool)

Encoder-Decoder Cross-Attention Mask：Decoder 每层的 Cross-Attention 中，K/V 来自 Encoder，需屏蔽 Encoder 输出中 padding 的部分。

'''
     decoder 看 encoder 输出时，遮住 encoder 侧的 padding 部分。
    Returns: shape (b, max_label_len, max_feat_len)，True 表示 Encoder padding 位置
'''
def get_enc_dec_mask(b: int, max_feat_len: int, feat_lens: torch.Tensor, max_label_len: int, device:torch.device) ->torch.Tensor:
    attn_mask = torch.zeros((b, max_label_len, max_feat_len), device=device)
    for i in range(b):
        #对第 i 个样本，把 encoder 真实长度之后的位置设为 1
        attn_mask[i, :, feat_lens[i]:] = 1
    return attn_mask.to(torch.bool)

Mutil-Head Attention

实现分四步：线性投影 → 拆分多头 → Scaled Dot-Product Attention → 合并输出。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_k, d_v, d_model, num_heads, p = 0.):
        '''
        Args:
            d_k: 每个注意力头的key/Query维度
            d_v: 每个注意力头的Value维度
            d_model: 输入/输出的总维度
            num_heads: 注意力头数
        '''
        super(MultiHeadAttention, self).__init__()
        self.d_k = d_k
        self.d_v = d_v
        self.d_model = d_model
        self.num_heads = num_heads
        self.dropout = nn.Dropout(p)

        #四个线性投影，Q、K、V投影和输出投影
        self.W_Q = nn.Linear(d_model, d_k * num_heads)
        self.W_K = nn.Linear(d_model, d_k * num_heads)
        self.W_V = nn.Linear(d_model, d_v * num_heads)
        self.W_out = nn.Linear(d_v * num_heads, d_model)

        # 权重初始化
        nn.init.normal_(self.W_Q.weight, mean = 0, std = np.sqrt(2.0 / (d_model + d_k)))
        nn.init.normal_(self.W_K.weight, mean = 0, std = np.sqrt(2.0 / (d_model + d_k)))
        nn.init.normal_(self.W_V.weight, mean = 0, std = np.sqrt(2.0 / (d_model + d_v)))
        nn.init.normal_(self.W_out.weight, mean = 0, std = np.sqrt(2.0 / (d_model + d_v)))

    def forward(self, Q, K, V, attn_mask):
        '''
        args:
            Q: (batch, q_len, d_model)
            K: (batch, k_len, d_model)
            V: (batch, v_len, d_model), k_len = v_len
            attn_mask: (batch, q_len, k_len) 或 None
        return:
            output: (batch, q_len, d_model)
        '''
        #取第 0 维大小，也就是 batch size
        N = Q.size(0)
        q_len, k_len = Q.size(1), K.size(1)
        d_k, d_v, num_heads = self.d_k, self.d_v, self.num_heads
        #Step1: 线性投影+拆分多头
        #(N, q_len, d_model) -> (N, num_heads, q_len, d_k)
        Q = self.W_Q(Q).view(N, -1, num_heads, d_k).transpose(1, 2)
        K = self.W_K(K).view(N, -1, num_heads, d_k).transpose(1, 2)
        V = self.W_V(V).view(N, -1, num_heads, d_v).transpose(1, 2)
        #Step2: 广播Mask到head维度
        if attn_mask is not None:
            assert attn_mask.size() == (N, q_len, k_len)
            #unsqueeze(1):在第 1 维插入一个新维度变成 (N, 1, q_len, k_len)
            #repeat(1, num_heads, 1, 1)：把 head 维复制 num_heads次变成(N, num_heads, q_len, k_len)
            attn_mask = attn_mask.unsqueeze(1).repeat(1, num_heads, 1, 1).bool()
        # Step 3: Scaled Dot-Product Attention: softmax(QK^T / sqrt(d_k))
        #Q:(N, num_heads, q_len, d_k)
        # K.transpose(-1,-2):  (N, num_heads, d_k, k_len)
        #矩阵乘法后：scores:  (N, num_heads, q_len, k_len)表示每个query token和每个key token的相似度分数
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)
        if attn_mask is not None:
            #凡是 attn_mask == True 的位置，都把 scores 填成 -1e4
            scores.masked_fill_(attn_mask, -1e4)
            #对最后一维做 softmax,最后一维是 k_len，也就是每个query对所有 key 的分数
        attns = torch.softmax(scores, dim = -1)
        #对 attention 权重做 dropout。
        # 训练时会随机把一部分注意力连接置 0，用来正则化，降低过拟合。
        attns = self.dropout(attns)
        #Step4: 加权求和 + 合并多头 + 输出投影
        #output: (N, num_heads, q_len, d_v),每个 query 位置，根据 attention 权重，对所有 value 向量做加权平均
        output = torch.matmul(attns, V)
        #合并多头
        #(N, num_heads, q_len, d_v) -> (N, q_len, num_heads, d_v)
        #.contiguous()让内存连续。 
        # (N, q_len, num_heads, d_v)->(N, q_len, num_heads * d_v)
        output = output.transpose(1, 2).contiguous().reshape(N, -1, d_v * num_heads)
        #(N, q_len, d_v * num_heads)-> (N, q_len, d_model)
        output = self.W_out(output)
        return output

正弦位置编码

'''
    生成正弦/余弦位置编码矩阵（不可学习，冻结权重）。
    Returns: (seq_len, d_model)
'''
def pos_sinusoid_embedding(seq_len: int, d_model: int) -> torch.Tensor:
    embeddings = torch.zeros((seq_len, d_model))
    for i in range(d_model):
        f =  torch.sin if i % 2 == 0 else torch.cos
        #torch.arange(0, seq_len)生成所有位置：[0, 1, 2, ..., seq_len - 1]
        embeddings[:, i] = f(torch.arange(0, seq_len) / np.power(1e4, 2 * (i // 2) /d_model))
    return embeddings.float()
'''
    pos_sinusoid_embedding(seq_len=3, d_model=4)
    pos 0: [sin(0), cos(0), sin(0), cos(0)]
    pos 1: [sin(1), cos(1), sin(1/100), cos(1/100)]
    pos 2: [sin(2), cos(2), sin(2/100), cos(2/100)]
    d_model = 4
    第 2/3 维的分母是 10000^(2/4) = 100
'''

Point-Wise Feed-Forward Network

$$FFN(x)=ReLU(xW1+b1)W2+b$$

class FFN(nn.Module):
    def __init__(self, d_model: int, d_ff: int, p: float = 0.):
        #d_ff通常为4 * d_model
        super(FFN, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU(inplace = True)
        self.dropout = nn.Dropout(p=p)
        
    def forward(self, X):
        #X: (N, seq_len, d_model)
        out = self.fc1(X)
        out = self.relu(out)
        out = self.fc2(out)
        return self.dropout(out)

EncoderLayer与Encoder

EncoderLayer 包括Multi-Head Self-Attention + FFN，每个子层外包残差连接和 LayerNorm。

class EncoderLayer(nn.Module):
    def __init__(self, dim: int, n: int, dff: int, dropout_posffn:float, dropout_attn: float):
        assert dim % n == 0
        super(EncoderLayer, self).__init__()
        hdim = dim // n #每个注意力头的维度
        
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.multi_head_attn = MultiHeadAttention(hdim, hdim, dim, n, dropout_attn)
        self.ffn= FFN(dim, dff, p = dropout_posffn)

    def forward(self, enc_in, attn_mask):
        #self-attention: Q=K=V=enc_in
        residual = enc_in
        #对应MultiHeadAttention中的 def forward(self, Q, K, V, attn_mask):
        context = self.multi_head_attn(enc_in, enc_in, enc_in, attn_mask)
        out = self.norm1(residual + context)

        #FFN
        residual = out
        out = self.ffn(out)
        out = self.norm2(residual + out)
        return out

Encoder

class Encoder(nn.Module):
    def __init__(self, dropout_emb, dropout_posffn, dropout_attn,
                num_layers, enc_dim, num_heads, dff, tgt_len):
        super(Encoder, self).__init__()
        self.tgt_len = tgt_len
        #固定的正弦位置编码(freeze=True表示不参与梯度更新)
        self.pos_emb = nn.Embedding.from_pretrained(pos_sinusoid_embedding(tgt_len, enc_dim), freeze = True)
        self.emb_dropout = nn.Dropout(dropout_emb)
        self.layers = nn.ModuleList(
            [EncoderLayer(enc_dim, num_heads, dff, dropout_posffn, dropout_attn) for _ in range(num_layers)]
        )

    def forward(self, X, X_lens, mask = None):
        #X:(batch, seq_len, d_model)
        seq_len  = X.size(1)
        out = X + self.pos_emb(torch.arange(seq_len, device = X.device))
        out = self.emb_dropout(out)
        for layer in self.layers:
            out = layer(out, mask)
        return out

DecoderLayer和Decoder

DecoderLayer包含三个子层——Masked Self-Attention → Cross-Attention → FFN。

class DecoderLayer(nn.Module):
    def __init__(self, dim: int, n: int, dff: int, dropout_posffn: float, dropout_attn:float):
        assert  dim % n == 0
        super(DecoderLayer, self).__init__()
        hdim = dim // n
        
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.norm3 = nn.LayerNorm(dim)
        self.ffn = FFN(dim, dff, p = dropout_posffn)
        self.dec_attn = MultiHeadAttention(hdim, hdim, dim, n, dropout_attn) #Masked Self Attention
        self.enc_dec_attn = MultiHeadAttention(hdim, hdim, dim, n, dropout_attn) #Cross Attention

    def forward(self, dec_in, enc_out, dec_mask, dec_enc_mask):
        '''
            dec_in:       (batch, dec_len, d_model)  Decoder 当前层输入
            enc_out:      (batch, enc_len, d_model)  Encoder 最终输出
            dec_mask:     Causal Mask (batch, dec_len, dec_len)
            dec_enc_mask: Cross-Attn Mask (batch, dec_len, enc_len)
        '''

        #1.Masked Self-Attention
        residual = dec_in
        context = self.dec_attn(dec_in, dec_in, dec_in, dec_mask)
        dec_out  = self.norm1(residual + context)
        
        #cross attention
        residual = dec_out
        context = self.enc_dec_attn(dec_out, enc_out, enc_out, dec_enc_mask)
        dec_out = self.norm2(residual + context)
        
        #FFN
        residual = dec_out
        out = self.ffn(dec_out)
        dec_out = self.norm3(residual + out)
        return dec_out

Decoder

class Decoder(nn.Module):
    def __init__(self, dropout_emb, dropout_posffn, dropout_attn,
                 num_layers, dec_dim, num_heads, dff, tgt_len, tgt_vocab_size):
        super(Decoder, self).__init__()
        #Word Embedding将token ID映射为d_model维向量
        self.tgt_emb = nn.Embedding(tgt_vocab_size, dec_dim)
        self.dropout_emb = nn.Dropout(dropout_emb)
        self.pos_emb = nn.Embedding.from_pretrained(pos_sinusoid_embedding(tgt_len, dec_dim), freeze = True)
        self.layers = nn.ModuleList([DecoderLayer(dec_dim, num_heads, dff, dropout_posffn, dropout_attn) for _ in range(num_layers)])
        
    def forward(self, labels, enc_out, dec_mask, dec_enc_mask):
        #labels: (batch, dec_len) token ID序列
        tgt_emb = self.tgt_emb(labels)
        pos_emb = self.pos_emb(torch.arange(labels.size(1), device = labels.device))
        dec_out = self.dropout_emb(tgt_emb + pos_emb)
        for layer in self.layers:
            dec_out = layer(dec_out, enc_out, dec_mask, dec_enc_mask)
        return dec_out

整合Transformer

class Transformer(nn.Module):
    def __init__(self, frontend: nn.Module, encoder: Encoder,
                 decoder: Decoder, dec_out_dim: int, vocab: int):
        '''
            frontend:    输入特征变换（如从 fbank_dim → d_model 的线性层）
            dec_out_dim: Decoder 输出维度（= d_model）
            vocab:       目标词表大小
        '''
        super().__init__()
        self.frontend = frontend
        self.encoder = encoder
        self.decoder = decoder
        self.linear = nn.Linear(dec_out_dim, vocab)
    
    def forward(self, X: torch.Tensor, X_lens: torch.Tensor, labels: torch.Tensor):
        '''
            X:      (batch, enc_len, fbank_dim)   输入特征序列
            X_lens: (batch,)                      每个样本的实际输入长度
            labels: (batch, dec_len)              目标 token ID 序列
            Returns: logits (batch, dec_len, vocab_size)
        '''
        X_lens, labels = X_lens.long(), labels.long()
        b, device = X.size(0), X.device
        
        out = self.frontend(X)
        max_feat_len = out.size(1)
        enc_mask = get_len_mask(b, max_feat_len, X_lens, device)
        enc_out = self.encoder(out, X_lens, enc_mask)
        
        max_label_len = labels.size(1)
        dec_mask = get_subsequent_mask(b, max_label_len, device)
        dec_enc_mask = get_enc_dec_mask(b, max_feat_len, X_lens, max_label_len, device)
        dec_out = self.decoder(labels, enc_out, dec_mask, dec_enc_mask)
        
        return self.linear(dec_out)

测试

class Transformer(nn.Module):
    def __init__(self, frontend: nn.Module, encoder: Encoder,
                 decoder: Decoder, dec_out_dim: int, vocab: int):
        '''
            frontend:    输入特征变换（如从 fbank_dim → d_model 的线性层）
            dec_out_dim: Decoder 输出维度（= d_model）
            vocab:       目标词表大小
        '''
        super().__init__()
        self.frontend = frontend
        self.encoder = encoder
        self.decoder = decoder
        self.linear = nn.Linear(dec_out_dim, vocab)
    
    def forward(self, X: torch.Tensor, X_lens: torch.Tensor, labels: torch.Tensor):
        '''
            X:      (batch, enc_len, fbank_dim)   输入特征序列
            X_lens: (batch,)                      每个样本的实际输入长度
            labels: (batch, dec_len)              目标 token ID 序列
            Returns: logits (batch, dec_len, vocab_size)
        '''
        X_lens, labels = X_lens.long(), labels.long()
        b, device = X.size(0), X.device
        
        out = self.frontend(X)
        max_feat_len = out.size(1)
        enc_mask = get_len_mask(b, max_feat_len, X_lens, device)
        enc_out = self.encoder(out, X_lens, enc_mask)
        
        max_label_len = labels.size(1)
        dec_mask = get_subsequent_mask(b, max_label_len, device)
        dec_enc_mask = get_enc_dec_mask(b, max_feat_len, X_lens, max_label_len, device)
        dec_out = self.decoder(labels, enc_out, dec_mask, dec_enc_mask)
        
        return self.linear(dec_out)