从零开始实现Transformer：深入理解大型语言模型的核心架构

引言

在自然语言处理领域,Transformer架构的出现无疑是一个里程碑式的突破。它不仅推动了BERT、GPT等大型语言模型的诞生,更彻底改变了NLP技术的发展方向。本文将带领读者从零开始实现一个简单的Transformer模型,深入剖析其核心架构,以帮助大家更好地理解当前炙手可热的大型语言模型技术。

Transformer架构概览

Transformer最早由Google团队在2017年的论文《Attention is All You Need》中提出。它抛弃了传统的循环神经网络结构,完全基于注意力机制来捕捉序列中的依赖关系。Transformer主要由编码器和解码器两部分组成,每部分又包含多个相同的层,每层主要由多头注意力机制和前馈神经网络构成。

Transformer架构图

上图展示了Transformer的整体架构。可以看到,输入首先经过词嵌入层转换为向量,然后添加位置编码信息。之后数据流经过多层的编码器和解码器,最后通过一个线性层和softmax层得到输出概率分布。

核心组件详解

1. 词嵌入和位置编码

词嵌入将输入的token转换为固定维度的向量。由于Transformer没有循环或卷积结构,为了让模型感知token的位置信息,还需要添加位置编码。位置编码通常使用正弦和余弦函数生成:

def get_positional_encoding(seq_len, d_model):
    position = np.arange(seq_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
    pos_encoding = np.zeros((seq_len, d_model))
    pos_encoding[:, 0::2] = np.sin(position * div_term)
    pos_encoding[:, 1::2] = np.cos(position * div_term)
    return torch.FloatTensor(pos_encoding)

2. 多头注意力机制

多头注意力机制是Transformer的核心创新。它允许模型同时关注输入的不同位置,捕捉更丰富的特征信息。其计算过程如下:

将输入经过线性变换得到查询(Q)、键(K)和值(V)
计算注意力分数: scores = Q * K.T / sqrt(d_k)
应用softmax得到注意力权重
加权求和得到输出: attention = softmax(scores) * V

多头机制则是将上述过程并行执行多次,然后拼接结果。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
        attn_probs = torch.softmax(attn_scores, dim=-1)
        output = torch.matmul(attn_probs, V)
        return output
        
    def split_heads(self, x):
        batch_size, seq_length, d_model = x.size()
        return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2)
        
    def combine_heads(self, x):
        batch_size, _, seq_length, d_k = x.size()
        return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.num_heads * d_k)
        
    def forward(self, Q, K, V, mask=None):
        Q = self.split_heads(self.W_q(Q))
        K = self.split_heads(self.W_k(K))
        V = self.split_heads(self.W_v(V))
        
        attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
        output = self.W_o(self.combine_heads(attn_output))
        return output

3. 前馈神经网络

在每个编码器和解码器层中,多头注意力后面都跟着一个前馈神经网络。这是一个包含两个线性变换和一个ReLU激活函数的简单全连接网络:

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        
    def forward(self, x):
        return self.linear2(F.relu(self.linear1(x)))

4. 层归一化

为了稳定训练过程,Transformer在每个子层的输出都添加了层归一化:

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super().__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

模型训练

有了上述核心组件,我们就可以组装完整的Transformer模型并进行训练了。以下是一个简化版的训练流程:

# 初始化模型、优化器等
model = Transformer(vocab_size, d_model, num_heads, num_layers, d_ff)
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)

# 训练循环
for epoch in range(num_epochs):
    for batch in dataloader:
        src, tgt = batch;
        
        optimizer.zero_grad()
        output = model(src, tgt[:, :-1])
        loss = criterion(output.contiguous().view(-1, vocab_size), tgt[:, 1:].contiguous().view(-1))
        loss.backward()
        optimizer.step()
        
    print(f"Epoch {epoch+1}, Loss: {loss.item()}")

在实际训练中,我们还需要考虑学习率调度、梯度裁剪等技巧来提升训练效果。

从零训练的挑战与解决方案

尽管理论上我们可以从零开始训练一个Transformer模型,但在实践中这往往面临诸多挑战:

数据量需求大:Transformer模型参数众多,需要海量数据才能充分训练。
计算资源消耗高:训练大型Transformer模型需要强大的GPU集群支持。
训练不稳定:从随机初始化开始训练容易出现梯度消失或爆炸等问题。

针对这些挑战,业界已经发展出一些有效的解决方案:

使用预训练模型:如BERT、GPT等通用预训练模型,可以在此基础上进行微调。
模型压缩:通过知识蒸馏、剪枝等技术减小模型规模,降低资源需求。
优化器改进:使用AdamW、RAdam等优化器提升训练稳定性。
渐进式训练:从小模型开始,逐步增加模型规模和数据量。

实战示例:训练一个简单的语言模型

让我们通过一个具体的例子来展示如何从零训练一个简单的Transformer语言模型。我们将使用一个小型的文本数据集,目标是让模型学会预测下一个单词。

首先,我们需要准备数据:

import requests
import tiktoken

# 下载数据集
url = "https://huggingface.co/datasets/goendalf666/sales-textbook_for_convincing_and_selling/raw/main/sales_textbook.txt"
response = requests.get(url)
text = response.text

# 使用tiktoken进行分词
encoder = tiktoken.get_encoding("gpt2")
tokens = encoder.encode(text)

# 创建训练数据
seq_length = 64
sequences = [tokens[i:i+seq_length+1] for i in range(0, len(tokens) - seq_length, seq_length)]
inputs = torch.tensor([seq[:-1] for seq in sequences])
targets = torch.tensor([seq[1:] for seq in sequences])

接下来,我们定义一个简化版的Transformer模型:

class SimpleTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model, max_len=5000)
        encoder_layers = TransformerEncoderLayer(d_model, nhead)
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers)
        self.fc_out = nn.Linear(d_model, vocab_size)
        
    def forward(self, src):
        embedded = self.embedding(src) * math.sqrt(d_model)
        embedded = self.pos_encoder(embedded)
        output = self.transformer_encoder(embedded)
        return self.fc_out(output)

现在我们可以开始训练过程:

# 初始化模型
vocab_size = encoder.n_vocab
d_model = 256
nhead = 8
num_layers = 4
model = SimpleTransformer(vocab_size, d_model, nhead, num_layers)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练循环
num_epochs = 50
batch_size = 64
for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    for i in range(0, len(inputs), batch_size):
        batch_inputs = inputs[i:i+batch_size]
        batch_targets = targets[i:i+batch_size]
        
        outputs = model(batch_inputs)
        loss = criterion(outputs.view(-1, vocab_size), batch_targets.view(-1))
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    
    avg_loss = total_loss / (len(inputs) // batch_size)
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {avg_loss:.4f}")

# 保存模型
torch.save(model.state_dict(), "simple_transformer.pth")

训练完成后,我们可以使用模型生成一些文本:

def generate_text(model, start_text, max_length=100):
    model.eval()
    tokens = encoder.encode(start_text)
    input_ids = torch.tensor(tokens).unsqueeze(0)
    
    with torch.no_grad():
        for _ in range(max_length):
            outputs = model(input_ids)
            next_token_logits = outputs[0, -1, :]
            next_token = torch.argmax(next_token_logits).item()
            input_ids = torch.cat([input_ids, torch.tensor([[next_token]])], dim=1)
            
            if next_token == encoder.eot_token:
                break
    
    generated_text = encoder.decode(input_ids[0].tolist())
    return generated_text

start_text = "The key to successful sales is"
generated_text = generate_text(model, start_text)
print(generated_text)