Transformer 架构

相关源文件

• docs/transformers/basic/autoregressive_experiment.html
• docs/transformers/feedback/experiment.html
• docs/transformers/mha.html
• docs/transformers/models.html
• docs/transformers/positional_encoding.html
• labml_nn/normalization/deep_norm/experiment.ipynb
• labml_nn/transformers/basic/autoregressive_experiment.ipynb
• labml_nn/transformers/basic/autoregressive_experiment.py
• labml_nn/transformers/feedback/__init__.py
• labml_nn/transformers/feedback/experiment.py
• labml_nn/transformers/label_smoothing_loss.py
• labml_nn/transformers/mha.py
• labml_nn/transformers/models.py
• labml_nn/transformers/positional_encoding.py
• labml_nn/transformers/relative_mha.py
• labml_nn/transformers/utils.py

本页概述了 labml_nn 仓库中实现的 Transformer 架构。它涵盖了核心 Transformer 组件、不同变体，以及这些组件如何相互作用以形成用于自然语言处理任务的完整模型。

本文档侧重于 Transformer 的实现细节，而非理论基础。有关训练 Transformer 模型的信息，请参阅训练基础设施部分 训练基础设施。

核心 Transformer 架构

Transformer 架构在论文“Attention Is All You Need”中提出，它依赖于自注意力机制，而非循环或卷积。此仓库中的实现遵循原始架构并包含一些变体。

来源： labml_nn/transformers/mha.py69-206 labml_nn/transformers/models.py24-207

多头注意力

多头注意力是 Transformer 架构的核心组件，它使模型能够同时关注输入序列的不同部分。

该实现包含以下组件：

1. PrepareForMultiHeadAttention：准备用于注意力计算的查询、键和值向量
2. MultiHeadAttention：实现多头注意力机制的主类

实现细节

MultiHeadAttention(heads, d_model, dropout_prob=0.1, bias=True)
├── query = PrepareForMultiHeadAttention(d_model, heads, d_k, bias)
├── key = PrepareForMultiHeadAttention(d_model, heads, d_k, bias)
├── value = PrepareForMultiHeadAttention(d_model, heads, d_k, bias)
└── forward(query, key, value, mask=None)
    ├── Compute attention scores: Q·K^T / sqrt(d_k)
    ├── Apply mask (if provided)
    ├── Apply softmax to get attention weights
    └── Compute weighted values and project back

主要参数

• heads: 注意力头数
• d_model：模型维度
• d_k：每个头的维度 (d_model / heads)
• dropout_prob：Dropout 概率

来源： labml_nn/transformers/mha.py33-66 labml_nn/transformers/mha.py69-206

Transformer 层

Transformer 架构由编码器和解码器层组成。每个层都有其自己的组件集。

编码器

编码器由多个相同的层组成，每个层包含：

1. 多头自注意力机制
2. 逐位前馈网络

解码器

解码器与编码器类似，但包含一个额外的注意力层：

1. 掩码多头自注意力机制
2. 在编码器输出上的多头注意力
3. 逐位前馈网络

来源： labml_nn/transformers/models.py60-128 labml_nn/transformers/models.py131-150 labml_nn/transformers/models.py153-172

位置编码

由于 Transformer 不包含循环或卷积，它们需要位置编码来整合序列顺序信息。

该仓库实现了两种类型的位置编码：

1. 固定正弦位置编码：使用不同频率的正弦和余弦函数

   PE_{pos, 2i} = sin(pos / 10000^{2i/d_model})
   PE_{pos, 2i+1} = cos(pos / 10000^{2i/d_model})
   

2. 学习位置编码：在训练过程中学习的嵌入

来源： labml_nn/transformers/positional_encoding.py30-61 labml_nn/transformers/models.py24-39 labml_nn/transformers/models.py42-57

Transformer 变体

该仓库实现了几种具有不同架构修改的 Transformer 变体。

基本 Transformer

基本 Transformer 遵循论文“Attention Is All You Need”中描述的架构，包含编码器和解码器组件。

来源： labml_nn/transformers/models.py193-224

自回归 Transformer

自回归 Transformer 专为文本生成任务设计，仅使用 Transformer 的编码器部分并带有因果掩码。

来源： labml_nn/transformers/basic/autoregressive_experiment.py26-60

反馈 Transformer

反馈 Transformer 通过允许每个层关注前一时间步所有层的输出来为 Transformer 架构添加循环。

反馈 Transformer 有两种实现方式：

1. 原始反馈 Transformer：存储所有层输出的加权和
2. 更新的反馈 Transformer（带共享 KV）：预计算并缓存键和值以获得更好的性能

来源： labml_nn/transformers/feedback/__init__.py198-250 labml_nn/transformers/feedback/__init__.py444-530

训练与配置

该仓库提供了用于训练 Transformer 模型的配置和实用工具。

常见配置参数：

参数	描述	示例值
tokenizer	分词器类型	'character'
seq_len	上下文/序列长度	512
batch_size	训练批次大小	16
d_model	模型维度	256
n_heads	注意力头的数量	16
d_ff	前馈网络维度	1024
optimizer	优化器类型	'Noam'
epochs	训练轮数	32

来源： labml_nn/transformers/basic/autoregressive_experiment.ipynb113-190 labml_nn/transformers/feedback/experiment.py113-129

模型初始化与生成

使用 Transformer 模型的典型工作流程包括：

1. 初始化模型组件（嵌入、注意力、层）
2. 构建编码器/解码器或自回归模型
3. 创建注意力掩码（例如，用于自回归生成的 subsequent_mask）
4. 通过模型传递输入
5. 生成输出（通常使用 top-k 或采样）

自回归 Transformer 的模型初始化示例：

来源： labml_nn/transformers/basic/autoregressive_experiment.py96-105 labml_nn/transformers/utils.py13-18

结论

labml_nn 仓库中的 Transformer 实现提供了模块化和可定制的组件，可以组合起来创建各种 Transformer 架构。该代码设计为易于理解和扩展，使其适用于研究和教育目的。

对于实际应用，用户可以根据其特定的任务要求和性能限制选择合适的 Transformer 变体。