这是理解《Attention Is All You Need》的最佳方式。它通过用PyTorch构建研究内容，帮助你理解基础概念。我强烈推荐它作为软件工程师的AI基础资源。https://nlp.seas.harvard.edu/annotated-transformer/…

带注释的Transformer

来源：https://nlp.seas.harvard.edu/annotated-transformer/

Attention Is All You Need (https://arxiv.org/abs/1706.03762)

• v2022: Austin Huang, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, and Stella Biderman.
• 原文 (https://nlp.seas.harvard.edu/2018/04/03/attention.html)：Sasha Rush (http://rush-nlp.com/).

Transformer在过去五年里一直备受关注。这篇文章以逐行实现的形式呈现了论文的注释版本。它重新排序并删除了原始论文中的部分章节，并在全文中添加了注释。本文档本身是一个可工作的笔记本，应该是一个完全可用的实现。代码可在此处获取 (https://github.com/harvardnlp/annotated-transformer/)。

目录

• 序言 (https://nlp.seas.harvard.edu/annotated-transformer/#prelims)
• 背景 (https://nlp.seas.harvard.edu/annotated-transformer/#background)
• 第一部分：模型架构 (https://nlp.seas.harvard.edu/annotated-transformer/#part-1-model-architecture)
  • 模型架构 (https://nlp.seas.harvard.edu/annotated-transformer/#model-architecture)
  • 编码器和解码器堆栈 (https://nlp.seas.harvard.edu/annotated-transformer/#encoder-and-decoder-stacks)
  • 位置前馈网络 (https://nlp.seas.harvard.edu/annotated-transformer/#position-wise-feed-forward-networks)
  • 嵌入和Softmax (https://nlp.seas.harvard.edu/annotated-transformer/#embeddings-and-softmax)
  • 位置编码 (https://nlp.seas.harvard.edu/annotated-transformer/#positional-encoding)
  • 完整模型 (https://nlp.seas.harvard.edu/annotated-transformer/#full-model)
  • 推理 (https://nlp.seas.harvard.edu/annotated-transformer/#inference)
• 第二部分：模型训练 (https://nlp.seas.harvard.edu/annotated-transformer/#part-2-model-training)
  • 训练 (https://nlp.seas.harvard.edu/annotated-transformer/#training)
  • 批次和掩码 (https://nlp.seas.harvard.edu/annotated-transformer/#batches-and-masking)
  • 训练循环 (https://nlp.seas.harvard.edu/annotated-transformer/#training-loop)
  • 训练数据和批处理 (https://nlp.seas.harvard.edu/annotated-transformer/#training-data-and-batching)
  • 硬件和调度 (https://nlp.seas.harvard.edu/annotated-transformer/#hardware-and-schedule)
  • 优化器 (https://nlp.seas.harvard.edu/annotated-transformer/#optimizer)
  • 正则化 (https://nlp.seas.harvard.edu/annotated-transformer/#regularization)
  • 第一个实例 (https://nlp.seas.harvard.edu/annotated-transformer/#a-first-example)
    • 合成数据 (https://nlp.seas.harvard.edu/annotated-transformer/#synthetic-data)
    • 损失计算 (https://nlp.seas.harvard.edu/annotated-transformer/#loss-computation)
    • 贪婪解码 (https://nlp.seas.harvard.edu/annotated-transformer/#greedy-decoding)
• 第三部分：一个真实世界的实例 (https://nlp.seas.harvard.edu/annotated-transformer/#part-3-a-real-world-example)
  • 数据加载 (https://nlp.seas.harvard.edu/annotated-transformer/#data-loading)
  • 迭代器 (https://nlp.seas.harvard.edu/annotated-transformer/#iterators)
  • 训练系统 (https://nlp.seas.harvard.edu/annotated-transformer/#training-the-system)
  • 附加组件：BPE、搜索、平均 (https://nlp.seas.harvard.edu/annotated-transformer/#additional-components-bpe-search-averaging)
  • 结果 (https://nlp.seas.harvard.edu/annotated-transformer/#results)
    • 注意力可视化 (https://nlp.seas.harvard.edu/annotated-transformer/#attention-visualization)
    • 编码器自注意力 (https://nlp.seas.harvard.edu/annotated-transformer/#encoder-self-attention)
    • 解码器自注意力 (https://nlp.seas.harvard.edu/annotated-transformer/#decoder-self-attention)
    • 解码器源注意力 (https://nlp.seas.harvard.edu/annotated-transformer/#decoder-src-attention)
  • 结论 (https://nlp.seas.harvard.edu/annotated-transformer/#conclusion)

序言

跳过 (https://nlp.seas.harvard.edu/annotated-transformer/#background)

# !pip install -r requirements.txt

# # 取消注释用于colab
# # # !pip install -q torchdata==0.3.0 torchtext==0.12 spacy==3.2 altair GPUtil
# # !python -m spacy download de_core_news_sm
# # !python -m spacy download en_core_web_sm

import os
from os.path import exists
import torch
import torch.nn as nn
from torch.nn.functional import log_softmax, pad
import math
import copy
import time
from torch.optim.lr_scheduler import LambdaLR
import pandas as pd
import altair as alt
from torchtext.data.functional import to_map_style_dataset
from torch.utils.data import DataLoader
from torchtext.vocab import build_vocab_from_iterator
import torchtext.datasets as datasets
import spacy
import GPUtil
import warnings
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

# 设置为 False 以跳过笔记本执行（例如用于调试）
warnings.filterwarnings("ignore")
RUN_EXAMPLES = True

# 一些在整个笔记本中使用的便利辅助函数
def is_interactive_notebook():
    return __name__ == "__main__"

def show_example(fn, args=[]):
    if __name__ == "__main__" and RUN_EXAMPLES:
        return fn(*args)

def execute_example(fn, args=[]):
    if __name__ == "__main__" and RUN_EXAMPLES:
        fn(*args)

class DummyOptimizer(torch.optim.Optimizer):
    def __init__(self):
        self.param_groups = [{"lr": 0}]
        None
    def step(self):
        None
    def zero_grad(self, set_to_none=False):
        None

class DummyScheduler:
    def step(self):
        None

我的注释采用块引用格式。主要文本均来自论文本身。

背景

减少顺序计算的目标也是扩展神经GPU、ByteNet和ConvS2S的基础，所有这些模型都使用卷积神经网络作为基本构建块，并行计算所有输入和输出位置的隐藏表示。在这些模型中，关联两个任意输入或输出位置信号所需操作的数量随着位置之间的距离而增长，对于ConvS2S呈线性增长，对于ByteNet呈对数增长。这使得学习远距离位置之间的依赖关系变得更加困难。在Transformer中，这被减少为常数数量的操作，尽管代价是由于平均注意力加权位置而导致有效分辨率降低，我们通过多头注意力来抵消这一影响。

自注意力，有时称为内部注意力，是一种注意力机制，它关联单个序列的不同位置，以计算该序列的表示。自注意力已成功用于多种任务，包括阅读理解、抽象摘要、文本蕴含和学习任务无关的句子表示。端到端记忆网络基于循环注意力机制而非序列对齐的循环，并且在简单语言问答和语言建模任务上表现良好。然而，据我们所知，Transformer是第一个完全依赖自注意力来计算其输入和输出表示，而不使用序列对齐的RNN或卷积的转换模型。

第一部分：模型架构

模型架构

大多数具有竞争力的神经序列转换模型都具有编码器-解码器结构 (https://arxiv.org/abs/1409.0473)。这里，编码器将输入符号表示序列 (x₁, …, xₙ) 映射到连续表示序列 z = (z₁, …, zₙ)。给定 z，解码器然后逐个元素地生成输出符号序列 (y₁, …, yₘ)。每一步，模型都是自回归的 (https://arxiv.org/abs/1308.0850)，在生成下一个符号时，将先前生成的符号作为额外输入。

class EncoderDecoder(nn.Module):
    """
    标准的编码器-解码器架构。这是此模型及许多其他模型的基础。
    """
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def forward(self, src, tgt, src_mask, tgt_mask):
        "接收并处理带掩码的源和目标序列。"
        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)

    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

class Generator(nn.Module):
    "定义标准的线性 + softmax 生成步骤。"
    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)

Transformer 遵循这个整体架构，在编码器和解码器中均使用堆叠的自注意力和逐点全连接层，分别如图1的左半部分和右半部分所示。

编码器和解码器堆栈

编码器

编码器由 N=6 个相同层的堆栈组成。

def clones(module, N):
    "产生 N 个相同的层。"
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class Encoder(nn.Module):
    "核心编码器是 N 层的堆栈"
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        "依次将输入（和掩码）传递给每一层。"
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

我们在两个子层中的每一个周围都使用了残差连接 (https://arxiv.org/abs/1512.03385)，然后进行层归一化 (https://arxiv.org/abs/1607.06450)。

class LayerNorm(nn.Module):
    "构建一个层归一化模块（详情见引用）。"
    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(x + Sublayer(x))，其中 Sublayer(x) 是子层本身实现的函数。我们对每个子层的输出应用 dropout (http://jmlr.org/papers/v15/srivastava14a.html)，然后将其添加到子层输入并进行归一化。为了促进这些残差连接，模型中的所有子层以及嵌入层都产生维度 d_model = 512 的输出。

class SublayerConnection(nn.Module):
    """
    残差连接后跟层归一化。
    注意，为了代码简洁，归一化放在前面而不是最后。
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "对任何相同大小的子层应用残差连接。"
        return x + self.dropout(sublayer(self.norm(x)))

每一层有两个子层。第一个是多头自注意力机制，第二个是简单的、逐位置全连接的前馈网络。

class EncoderLayer(nn.Module):
    "编码器由自注意力和前馈网络组成（定义如下）"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size

    def forward(self, x, mask):
        "按照图1（左）的连接顺序。"
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

解码器

解码器也由 N=6 个相同层的堆栈组成。

class Decoder(nn.Module):
    "通用的带掩码的 N 层解码器。"
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

除了每个编码器层中的两个子层之外，解码器还插入第三个子层，它对编码器堆栈的输出执行多头注意力。与编码器类似，我们在每个子层周围使用残差连接，然后进行层归一化。

class DecoderLayer(nn.Module):
    "解码器由自注意力、源注意力和前馈网络组成（定义如下）"
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)

    def forward(self, x, memory, src_mask, tgt_mask):
        "按照图1（右）的连接顺序。"
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)

我们还修改了解码器堆栈中的自注意力子层，以防止位置关注后续位置。这种掩蔽，加上输出嵌入偏移一个位置的事实，确保位置 i 的预测只能依赖于位置小于 i 的已知输出。

def subsequent_mask(size):
    "掩蔽后续位置。"
    attn_shape = (1, size, size)
    subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
        torch.uint8
    )
    return subsequent_mask == 0

下面的注意力掩码显示了每个目标词（行）允许查看的位置（列）。在训练期间，单词被阻止关注未来的单词。

def example_mask():
    LS_data = pd.concat(
        [
            pd.DataFrame(
                {
                    "Subsequent Mask": subsequent_mask(20)[0][x, y].flatten(),
                    "Window": y,
                    "Masking": x,
                }
            )
            for y in range(20)
            for x in range(20)
        ]
    )
    return (
        alt.Chart(LS_data)
        .mark_rect()
        .properties(height=250, width=250)
        .encode(
            alt.X("Window:O"),
            alt.Y("Masking:O"),
            alt.Color("Subsequent Mask:Q", scale=alt.Scale(scheme="viridis")),
        )
        .interactive()
    )
show_example(example_mask)

注意力

注意力函数可以被描述为将查询和一组键值对映射到输出，其中查询、键、值和输出都是向量。输出是值的加权和，其中分配给每个值的权重由查询与相应键的兼容性函数计算得出。我们将我们特定的注意力称为“缩放点积注意力”。输入由维度 d_k 的查询和键以及维度 d_v 的值组成。我们计算查询与所有键的点积，每个除以 sqrt(d_k)，并应用 softmax 函数以获得值的权重。

在实践中，我们同时在多个查询上计算注意力函数，这些查询被打包成一个矩阵 Q。键和值也分别被打包成矩阵 K 和 V。我们计算输出矩阵为：

[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V ]

def attention(query, key, value, mask=None, dropout=None):
    "计算 '缩放点积注意力'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

两种最常用的注意力函数是加性注意力（https……