之前虽然了解过 Transformer 架构，但是没有自己实现过。

最近阅读 transformers 库中 Llama 模型结构，于是想试着亲手实现一个简单的 Transformer。

在实现过程中加深了理解，同时发现之前阅读 Llama 中一些错误的地方，因此做一个记录。

笔者小白，如果实现过程中存在错误，请不吝指出。

Embedding

Embedding 可以将高维的离散文本数据映射到低维的连续向量空间。这不仅减小了输入数据的维度，也有助于减少数据的稀疏性，提高模型的性能和效率。

同时，词嵌入可以捕捉单词之间的语义关系，相似的单词在嵌入空间中会更接近。

使用 Pytorch 可以很方便定义出 Embedding 模型：

class Embedder(nn.Module):
    def __init__(self, vocab_size: int, d_model: int) -> None:
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.embed(x)

Positional Encoding

Transformer 中没有类似 RNN 的循环机制，需要通过位置编码记录单词的位置和顺序。

其计算位置编码的公式如下：

PE_{(pos,2i)}=sin(\frac{pos}{10000^{2i/d_{model}}})

PE_{(pos,2i+1)}=cos(\frac{pos}{10000^{2i/d_{model}}})

其中 $pos$ 是位置，而 $i$ 是维度。

Pytorch 实现位置编码器代码如下：

class PositionalEncoder(nn.Module):
    def __init__(
        self, d_model: int = 512, max_seq_len: int = 2048, base: int = 10000
    ) -> None:
        super().__init__()
        self.d_model = d_model

        inv_freq_half = 1.0 / (
            base ** (torch.arange(0, d_model, 2, dtype=torch.float) / d_model)
        )
        inv_freq = torch.arange(0, d_model, dtype=inv_freq_half.dtype)
        inv_freq[..., 0::2] = inv_freq_half
        inv_freq[..., 1::2] = inv_freq_half

        pos = torch.arange(max_seq_len, dtype=inv_freq.dtype)

        pe = torch.einsum("i, j -> ij", pos, inv_freq)
        pe[..., 0::2] = pe[..., 0::2].sin()
        pe[..., 1::2] = pe[..., 1::2].cos()

        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 使 embedding 相对大一些
        x = x * math.sqrt(self.d_model)
        seq_len = x.shape[1]
        pe = self.pe[:seq_len].to(dtype=x.dtype)
        return x + pe

在 PyTorch 中，nn.Module 类中的 register_buffer() 方法用于将一个张量（或缓冲区）注册为模型的一部分。

注册的缓冲区不会参与模型的梯度计算，但会在模型的保存和加载时保持状态。

register_buffer() 的主要作用是在模型中保留一些不需要梯度更新的状态。

在前向传播中加入位置编码前扩大 embedding 的值，是为了保证原始语言信息不会因为加入位置信息而丢失。

Mask

Mask 在 Transformer 中有很重要的作用：

在 Encoder 和 Decoder 中，Mask 会遮住用于 Padding 的位置。
在 Decoder 中，Mask 会遮住预测剩余位置，防止 Dcoder 提前得到信息。

Multi-Headed Attention

多头注意力是 Transformer 中的核心模块，它们网络结构如下：

在多头注意力中，会将 embedding 分割为 $h$ 个头，每个头的维度为 $d_{model} / h$ 。

In this work we employ $h = 8$ parallel attention layers, or heads. For each of these we use $d_k$ = $d_v$ = $d_{model}/h$ = 64.

多头注意力公式如下：

MultiHead(Q,K,V)=Concat(head_1,...,head_n)W^O

head_i=Attention(QW_i^Q,KW_i^K,VW_i^V)

多头注意力代码如下：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, heads: int = 8, dropout: int = 0.1) -> None:
        super().__init__()

        self.d_model = d_model
        self.heads = heads
        self.d_k = self.d_model // self.heads

        if self.heads * self.d_k != self.d_model:
            raise ValueError(
                f"d_model must be divisible by heads (got `d_model`: {self.d_model}"
                f" and `heads`: {self.heads})."
            )

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(dropout)
        self.o_proj = nn.Linear(d_model, d_model)

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ):
        bsz = q.shape[0]

        # translate [bsz, seq_len, d_model] to [bsz, seq_len, heads, d_k]
        q = self.q_proj(q).view(bsz, -1, self.heads, self.d_k)
        k = self.k_proj(k).view(bsz, -1, self.heads, self.d_k)
        v = self.v_proj(v).view(bsz, -1, self.heads, self.d_k)

        # translate [bsz, seq_len, heads, d_k] to [bsz, heads, seq_len, d_k]
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # calculate attention
        scores = attention(q, k, v, self.d_k, mask, self.dropout)

        # cat multi-heads
        concat = scores.transpose(1, 2).contiguous().view(bsz, -1, self.d_model)
        output = self.o_proj(concat)
        return output

注意力计算公式为：

Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V

其中计算注意力的代码如下：

def attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    d_k: int,
    mask: Optional[torch.Tensor] = None,
    dropout: Optional[nn.Dropout] = None,
) -> torch.Tensor:
    # calculate the scores
    # q: [bsz, heads, seq_len, d_k]
    # k: [bsz, heads, d_k, seq_len]
    scores = torch.matmul(q, k.transpose(-1, -2)) / torch.sqrt(d_k)

    if mask is not None:
        # tanslate [bsz, seq_len, seq_len] to [bsz, 1, seq_len, seq_len]
        mask = mask.unsqueeze(1)
        scores = scores.masked_fill(mask == 0, -1e9)
    scores = F.softmax(scores, dim=-1)

    if dropout is not None:
        scores = dropout(scores)

    output = torch.matmul(scores, v)
    return output

The Feed-Forward Network

Feed-Forward 由两个线性变换和一个激活函数构成。

This consists of two linear transformations with a ReLU activation in between.

其公式如下：

FFN=max(0,xW_1+b_1)W_2+b_2

该网络中输入输出维度为 512，中间线性层维度为 2048。

The dimensionality of input and output is $d_{model} = 512$ , and the inner-layer has dimensionality $d_{ff} = 2048$ .

实现代码如下：

class FeedForward(nn.Module):
    def __init__(
        self, d_model: int = 512, d_ff: int = 2048, dropout: float = 0.1
    ) -> None:
        super().__init__()

        self.linear_1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear_2 = nn.Linear(d_ff, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.dropout(F.relu(self.linear_1(x)))
        x = self.linear_2(x)
        return x

Norm

正则化可以防止数据在不同网络中流动时范围差距过大，保证模型稳定性。

实现代码如下：

class Norm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-6) -> None:
        super().__init__()

        self.dim = d_model

        self.alpha = nn.Parameter(torch.ones(self.dim))
        self.bias = nn.Parameter(torch.zeros(self.dim))

        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = (
            self.alpha
            * (x - x.mean(dim=-1, keepdim=True))
            / (x.std(dim=-1, keepdim=True) + self.eps)
            + self.bias
        )
        return norm

Assemble

Transformer 由多个 EncoderLayer 和 DecoderLayer 组合在一起，首先实现 EncoderLayer。

[注意] 在每个子层输出和下一个子层输入以及正则化前，有一层 dropout。

We apply dropout [33] to the output of each sub-layer, before it is added to the sub-layer input and normalized. In addition, we apply dropout to the sums of the embeddings and the positional encodings in both the encoder and decoder stacks. For the base model, we use a rate of Pdrop = 0.1

EncoderLayer 实现代码如下：

class EncoderLayer(nn.Module):
    def __init__(
        self, d_model: int = 512, heads: int = 8, d_ff: int = 2048, dropout: float = 0.1
    ) -> None:
        super().__init__()

        self.attn = MultiHeadAttention(d_model, heads, dropout)
        self.dropout_1 = nn.Dropout(dropout)
        self.norm_1 = Norm(d_model)

        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.dropout_2 = nn.Dropout(dropout)
        self.norm_2 = Norm(d_model)

    def forward(
        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        x = x + self.dropout_1(self.attn(x, x, x, mask))
        x = self.norm_1(x)

        x = x + self.dropout_2(self.ffn(x))
        x = self.norm_2(x)
        return x

DecoderLayer 实现代码如下：

class DecoderLayer(nn.Module):
    def __init__(
        self, d_model: int = 512, heads: int = 8, d_ff: int = 2048, dropout: float = 0.1
    ) -> None:
        super().__init__()

        self.attn_1 = MultiHeadAttention(d_model, heads, dropout)
        self.dropout_1 = nn.Dropout(dropout)
        self.norm_1 = Norm(d_model)

        self.attn_2 = MultiHeadAttention(d_model, heads, dropout)
        self.dropout_2 = nn.Dropout(dropout)
        self.norm_2 = Norm(d_model)

        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.dropout_3 = nn.Dropout(dropout)
        self.norm_3 = Norm(d_model)

    def forward(
        self,
        x: torch.Tensor,
        enc_output: torch.Tensor,
        src_mask: torch.Tensor,
        tgt_mask: torch.Tensor,
    ) -> torch.Tensor:
        x = x + self.dropout_1(self.attn_1(x, x, x, tgt_mask))
        x = self.norm_1(x)

        x = x + self.dropout_2(self.attn_2(x, enc_output, enc_output, src_mask))
        x = self.norm_2(x)

        x = x + self.dropout_3(self.ffn(x))
        x = self.norm_3(x)
        return x

Encoder 和 Decoder 分别由 N 个 EncoderLayer 和 DecoderLayer 组成。

代码实现如下：

class Encoder(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        N: int = 6,
        d_model: int = 512,
        max_seq_len: int = 2048,
        heads: int = 8,
        d_ff: int = 2048,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()

        self.N = N
        self.embed = Embedder(vocab_size, d_model)
        self.pe = PositionalEncoder(d_model, max_seq_len)
        self.layers = nn.ModuleList(
            [EncoderLayer(d_model, heads, d_ff, dropout) for _ in range(N)]
        )

    def forward(self, src: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        x = self.embed(src)
        x = self.pe(x)
        for layer in self.layers:
            x = layer(x, mask)
        return x


class Decoder(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        N: int = 6,
        d_model: int = 512,
        max_seq_len: int = 2048,
        heads: int = 8,
        d_ff: int = 2048,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()

        self.N = N
        self.embed = Embedder(vocab_size, d_model)
        self.pe = PositionalEncoder(d_model, max_seq_len)
        self.layers = nn.ModuleList(
            [DecoderLayer(d_model, heads, d_ff, dropout) for _ in range(N)]
        )

    def forward(
        self,
        tgt: torch.Tensor,
        enc_output: torch.Tensor,
        src_mask: torch.Tensor,
        tgt_mask: torch.Tensor,
    ) -> torch.Tensor:
        x = self.embed(tgt)
        x = self.pe(x)
        for layer in self.layers:
            x = layer(x, enc_output, src_mask, tgt_mask)
        return x

最后组装成 Transformer！

class Transformer(nn.Module):
    def __init__(
        self,
        src_vocab: int,
        tgt_vocab: int,
        N: int = 6,
        d_model: int = 512,
        max_seq_len: int = 2048,
        heads: int = 8,
        d_ff: int = 2048,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()

        self.encoder = Encoder(src_vocab, N, d_model, max_seq_len, heads, d_ff, dropout)
        self.decoder = Decoder(tgt_vocab, N, d_model, max_seq_len, heads, d_ff, dropout)
        self.out = nn.Linear(d_model, tgt_vocab)

    def forward(
        self,
        src: torch.Tensor,
        tgt: torch.Tensor,
        src_mask: torch.Tensor,
        tgt_mask: torch.Tensor,
    ) -> torch.Tensor:
        enc_output = self.encoder(src, src_mask)
        dec_output = self.decoder(tgt, enc_output, src_mask, tgt_mask)
        output = F.softmax(self.out(dec_output), dim=-1)
        return output

Test

测试一下代码能不能运行，按照如下配置测试：

from transformer_scratch import Transformer
import torch

bsz = 4
max_seq_len = 1024
src_vocab = 128
tgt_vocab = 64
N = 3
d_ff = 512

model = Transformer(src_vocab, tgt_vocab, N=N, max_seq_len=max_seq_len, d_ff=d_ff)

src = torch.randint(low=0, high=src_vocab, size=(bsz, max_seq_len))
tgt = torch.randint(low=0, high=tgt_vocab, size=(bsz, max_seq_len))
src_mask = torch.ones(size=(bsz, max_seq_len, max_seq_len))
tgt_mask = torch.ones(size=(bsz, max_seq_len, max_seq_len))

res = model(src, tgt, src_mask, tgt_mask)
print(f"Output data shape is: {res.shape}")

输出：Output data shape is: torch.Size([4, 1024, 64])

Reference

在编写过程中参考下面的博客，感谢大佬分享自己的经验。

How to code The Transformer in Pytorch | by Samuel Lynn-Evans | Towards Data Science

Pytorch 从零实现 Transformer