当前位置：首页 > news >正文

一个基于稀疏混合专家模型（Sparse Mixture of Experts, Sparse MoE）的 Transformer 语言模型

news 2025/10/13 6:52:43

一个基于稀疏混合专家模型（Sparse Mixture of Experts, Sparse MoE）的 Transformer 语言模型

flyfish

用稀疏混合专家（Sparse MoE）替代传统Transformer的全连接层

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.nn import init# 超参数设置
batch_size = 16  # 并行处理的独立序列数量
block_size = 32  # 预测的最大上下文长度
max_iters = 5000  # 最大训练迭代次数
eval_interval = 100  # 评估间隔
learning_rate = 1e-3  # 学习率
device = 'cuda' if torch.cuda.is_available() else 'cpu'  # 设备选择
eval_iters = 400  # 评估迭代次数
n_embed = 128  # 嵌入维度
n_head = 8  # 注意力头数量
n_layer = 8  # 网络层数
dropout = 0.1  # dropout比率
num_experts = 8  # 专家数量，可根据总参数量调整
top_k = 2  # 激活的专家数量torch.manual_seed(1337)  # 设置随机种子，保证结果可复现# 加载并预处理数据
with open('input.txt', 'r', encoding='utf-8') as f:text = f.read()# 获取文本中所有唯一字符
chars = sorted(list(set(text)))
vocab_size = len(chars)
# 创建字符与整数的映射
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
# 编码器：将字符串转换为整数列表
encode = lambda s: [stoi[c] for c in s]
# 解码器：将整数列表转换为字符串
decode = lambda l: ''.join([itos[i] for i in l])# 划分训练集和测试集
data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))  # 前90%作为训练集，其余作为验证集
train_data = data[:n]
val_data = data[n:]# 数据加载函数
def get_batch(split):"""生成输入x和目标y的小批量数据"""data = train_data if split == 'train' else val_data# 随机选择批量数据的起始索引ix = torch.randint(len(data) - block_size, (batch_size,))x = torch.stack([data[i:i+block_size] for i in ix])y = torch.stack([data[i+1:i+block_size+1] for i in ix])x, y = x.to(device), y.to(device)return x, y@torch.no_grad()  # 禁用梯度计算，节省内存并加速
def estimate_loss(model):"""评估模型在训练集和验证集上的损失"""out = {}model.eval()  # 切换到评估模式for split in ['train', 'val']:losses = torch.zeros(eval_iters)for k in range(eval_iters):X, Y = get_batch(split)logits, loss = model(X, Y)losses[k] = loss.item()out[split] = losses.mean()model.train()  # 切换回训练模式return outclass Head(nn.Module):"""自注意力的单个头"""def __init__(self, head_size):super().__init__()self.key = nn.Linear(n_embed, head_size, bias=False)self.query = nn.Linear(n_embed, head_size, bias=False)self.value = nn.Linear(n_embed, head_size, bias=False)# 注册下三角矩阵作为缓冲区（不会被视为模型参数）self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))self.dropout = nn.Dropout(dropout)def forward(self, x):B, T, C = x.shape  # B: 批量大小, T: 时间步长, C: 通道数(嵌入维度)k = self.key(x)    # (B, T, C)q = self.query(x)  # (B, T, C)# 计算注意力分数wei = q @ k.transpose(-2, -1) * C**-0.5  # (B, T, T)，缩放点积注意力# 遮罩未来位置的信息（因果注意力）wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))wei = F.softmax(wei, dim=-1)  # 归一化wei = self.dropout(wei)       # 应用dropout# 对值进行加权聚合v = self.value(x)  # (B, T, C)out = wei @ v      # (B, T, C)return outclass MultiHeadAttention(nn.Module):"""并行的多个自注意力头"""def __init__(self, num_heads, head_size):super().__init__()self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])self.proj = nn.Linear(n_embed, n_embed)  # 投影层，将多头输出映射回嵌入维度self.dropout = nn.Dropout(dropout)def forward(self, x):# 拼接所有头的输出out = torch.cat([h(x) for h in self.heads], dim=-1)# 投影并应用dropoutout = self.dropout(self.proj(out))return outclass Expert(nn.Module):"""专家模块：简单的MLP，包含线性层和非线性激活"""def __init__(self, n_embed):super().__init__()self.net = nn.Sequential(nn.Linear(n_embed, 4 * n_embed),  # 扩展维度nn.ReLU(),                        # 非线性激活nn.Linear(4 * n_embed, n_embed),  # 投影回原维度nn.Dropout(dropout),              # 防止过拟合)def forward(self, x):return self.net(x)class NoisyTopkRouter(nn.Module):"""带噪声的Top-K路由模块，决定每个输入应该由哪些专家处理"""def __init__(self, n_embed, num_experts, top_k):super().__init__()self.top_k = top_kself.topkroute_linear = nn.Linear(n_embed, num_experts)  # 用于计算路由logitsself.noise_linear = nn.Linear(n_embed, num_experts)      # 用于计算噪声强度def forward(self, mh_output):# mh_output是多头注意力块的输出张量logits = self.topkroute_linear(mh_output)  # 路由logits# 计算噪声logits并生成噪声noise_logits = self.noise_linear(mh_output)noise = torch.randn_like(logits) * F.softplus(noise_logits)  # 缩放的高斯噪声# 给logits添加噪声noisy_logits = logits + noise# 选择top-k的专家top_k_logits, indices = noisy_logits.topk(self.top_k, dim=-1)# 创建稀疏logits（只保留top-k的logits，其余为负无穷）zeros = torch.full_like(noisy_logits, float('-inf'))sparse_logits = zeros.scatter(-1, indices, top_k_logits)# 计算路由权重router_output = F.softmax(sparse_logits, dim=-1)return router_output, indicesclass SparseMoE(nn.Module):"""稀疏混合专家模块：根据路由结果选择专家处理输入"""def __init__(self, n_embed, num_experts, top_k, capacity_factor=1.0):super().__init__()self.router = NoisyTopkRouter(n_embed, num_experts, top_k)self.experts = nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])self.top_k = top_kself.capacity_factor = capacity_factor  # 控制专家容量的因子self.num_experts = num_expertsdef forward(self, x):# 假设x的形状为 [batch_size, seq_len, n_embed]batch_size, seq_len, _ = x.shape# 获取路由权重和选中的专家索引gating_output, indices = self.router(x)final_output = torch.zeros_like(x)# 展平批量和序列维度，以便独立处理每个tokenflat_x = x.view(-1, x.size(-1))flat_gating_output = gating_output.view(-1, gating_output.size(-1))# 计算每个专家的容量tokens_per_batch = batch_size * seq_len * self.top_kexpert_capacity = int((tokens_per_batch / self.num_experts) * self.capacity_factor)updates = torch.zeros_like(flat_x)# 每个专家处理分配给它的tokenfor i, expert in enumerate(self.experts):# 找到分配给当前专家的tokenexpert_mask = (indices == i).any(dim=-1)flat_mask = expert_mask.view(-1)selected_indices = torch.nonzero(flat_mask).squeeze(-1)# 应用容量限制，防止单个专家处理过多tokenif selected_indices.numel() > expert_capacity:selected_indices = selected_indices[:expert_capacity]if selected_indices.numel() > 0:# 专家处理选中的tokenexpert_input = flat_x[selected_indices]expert_output = expert(expert_input)# 应用路由权重gating_scores = flat_gating_output[selected_indices, i].unsqueeze(1)weighted_output = expert_output * gating_scores# 累加结果updates.index_add_(0, selected_indices, weighted_output)# 重塑回原始形状final_output += updates.view(batch_size, seq_len, -1)return final_outputclass Block(nn.Module):"""混合专家Transformer块：先通信（多头自注意力）后计算（稀疏混合专家）"""def __init__(self, n_embed, n_head, num_experts, top_k):# n_embed: 嵌入维度, n_head: 注意力头数量super().__init__()head_size = n_embed // n_headself.sa = MultiHeadAttention(n_head, head_size)  # 自注意力模块self.smoe = SparseMoE(n_embed, num_experts, top_k)  # 稀疏混合专家模块self.ln1 = nn.LayerNorm(n_embed)  # 层归一化self.ln2 = nn.LayerNorm(n_embed)def forward(self, x):# 残差连接 + 自注意力x = x + self.sa(self.ln1(x))# 残差连接 + 混合专家x = x + self.smoe(self.ln2(x))return xclass SparseMoELanguageModel(nn.Module):"""稀疏混合专家语言模型"""def __init__(self):super().__init__()# 词嵌入表：每个token直接从查找表中读取嵌入self.token_embedding_table = nn.Embedding(vocab_size, n_embed)# 位置嵌入表：为每个位置提供嵌入self.position_embedding_table = nn.Embedding(block_size, n_embed)# 多个Transformer块self.blocks = nn.Sequential(*[Block(n_embed, n_head=n_head, num_experts=num_experts, top_k=top_k) for _ in range(n_layer)])self.ln_f = nn.LayerNorm(n_embed)  # 最终层归一化self.lm_head = nn.Linear(n_embed, vocab_size)  # 语言模型头，预测下一个tokendef forward(self, idx, targets=None):B, T = idx.shape  # B: 批量大小, T: 时间步长# 获取词嵌入和位置嵌入tok_emb = self.token_embedding_table(idx)  # (B, T, C)pos_emb = self.position_embedding_table(torch.arange(T, device=device))  # (T, C)x = tok_emb + pos_emb  # 合并嵌入 (B, T, C)# 通过Transformer块x = self.blocks(x)  # (B, T, C)x = self.ln_f(x)    # 最终层归一化logits = self.lm_head(x)  # 计算logits (B, T, vocab_size)# 计算损失（如果提供了目标）if targets is None:loss = Noneelse:B, T, C = logits.shapelogits = logits.view(B*T, C)  # 展平以便计算交叉熵targets = targets.view(B*T)loss = F.cross_entropy(logits, targets)return logits, lossdef generate(self, idx, max_new_tokens):"""生成新的token序列"""# idx是当前上下文的索引数组，形状为 (B, T)for _ in range(max_new_tokens):# 将idx裁剪到最后block_size个tokenidx_cond = idx[:, -block_size:]# 获取预测logits, loss = self(idx_cond)# 只关注最后一个时间步logits = logits[:, -1, :]  # 变为 (B, C)# 应用softmax获取概率probs = F.softmax(logits, dim=-1)  # (B, C)# 从分布中采样idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)# 将采样的索引追加到序列中idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)return idxdef kaiming_init_weights(m):"""使用Kaiming初始化权重"""if isinstance(m, nn.Linear):init.kaiming_normal_(m.weight)def main():# 初始化模型并应用权重初始化model = SparseMoELanguageModel()model.apply(kaiming_init_weights)model = model.to(device)# 打印模型参数数量print(f"{sum(p.numel() for p in model.parameters()) / 1e6:.2f} M parameters")# 优化器optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)# 训练循环for iter in range(max_iters):# 定期评估损失if iter % eval_interval == 0 or iter == max_iters - 1:losses = estimate_loss(model)print(f"迭代 {iter}: 训练损失 {losses['train']:.4f}, 验证损失 {losses['val']:.4f}")# 获取批量数据xb, yb = get_batch('train')# 计算损失并优化logits, loss = model(xb, yb)optimizer.zero_grad(set_to_none=True)  # 清空梯度loss.backward()                        # 反向传播optimizer.step()                       # 更新参数# 训练完成后生成一些文本context = torch.zeros((1, 1), dtype=torch.long, device=device)print(decode(model.generate(context, max_new_tokens=500)[0].tolist()))if __name__ == "__main__":main()

一个基于稀疏混合专家模型（Sparse Mixture of Experts, Sparse MoE） 的Transformer语言模型

功能

数据处理
首先加载文本数据（input.txt），构建字符级别的词汇表（将文本中的每个独特字符映射到整数ID），然后将文本转换为整数序列。接着划分训练集和验证集，并用get_batch函数生成模型训练所需的批量数据（输入序列x和对应的目标序列y，其中y是x向右偏移一位的序列，用于训练“预测下一个字符”的能力）。
模型结构
模型基于Transformer架构，引入了“稀疏混合专家”机制，核心组件包括：
- 自注意力模块：包括单个注意力头（Head）和多头注意力（MultiHeadAttention），用于捕捉文本序列中的上下文依赖关系（例如“明天”和“天气”的关联）。
- 专家模块（Expert）：每个专家是一个简单的神经网络（MLP），负责处理不同类型的特征（可以理解为“分工合作”，比如有的专家擅长处理语法，有的擅长处理语义）。
- 路由模块（NoisyTopkRouter）：决定每个输入token由哪些专家处理（通过带噪声的Top-K选择，避免路由决策过于固化）。
- 稀疏混合专家（SparseMoE）：整合路由结果，仅让被选中的专家处理对应token，平衡模型能力和计算效率（相比普通Transformer，能用更少的计算量实现更大的模型容量）。
- Transformer块（Block）：将“多头注意力”和“稀疏混合专家”结合，加上残差连接和层归一化，构成模型的基本单元，重复多层（n_layer）提升模型能力。
训练与评估
- 用estimate_loss函数定期评估模型在训练集和验证集上的损失（交叉熵损失，衡量“预测下一个字符”的准确率）。
- 通过AdamW优化器训练模型，迭代更新参数，最小化预测损失。
文本生成
模型训练完成后，通过generate函数生成新文本：从初始上下文（例如空字符）出发，不断预测下一个字符并拼接，最终生成连贯的文本序列。

用“稀疏混合专家”机制增强Transformer，在保持计算效率的同时提升模型容量，最终实现基于字符的文本生成（例如输入一段文字，让模型续写下去）。

模块

数据处理→模型训练→文本生成

一、数据处理模块：为模型准备“可读懂”的文本数据

负责将原始文本（如你提供的戏剧对话）转换为模型能处理的数值格式，是训练的基础。包含3个关键组件：

组件名称（代码对应）	核心功能
数据加载与编码（`encode`）	1. 读取`input.txt`中的原始文本； 2. 将文本转换为整数序列（字符→ID映射）。
词汇表构建（`stoi/itos`）	1. 提取文本中所有唯一字符（如字母、空格、`:`、`?`等）； 2. 建立“字符→ID”（`stoi`）和“ID→字符”（`itos`）的双向映射，用于编码和解码。
批量数据生成（`get_batch`）	1. 从训练/验证集中随机抽取数据； 2. 生成输入序列`x`（上下文）和目标序列`y`（`x`右移1位，用于训练“预测下一个字符”）； 3. 将数据移动到GPU/CPU设备。

二、模块：实现“稀疏混合专家Transformer”的核心逻辑

这是代码的核心，负责定义模型的结构，核心是“用稀疏专家分工处理文本，兼顾效果与效率”。包含6个关键组件：

组件名称（代码类/函数）	核心功能
自注意力头（`Head`）	单个自注意力单元，负责捕捉文本序列的上下文依赖（如“First Citizen”与后续冒号、台词的关联）；通过“掩码”确保只关注“过去/当前字符”，不泄露未来信息（因果注意力）。
多头注意力（`MultiHeadAttention`）	将多个`Head`并行运行，再通过投影层整合结果；可同时捕捉不同维度的上下文（如有的头关注语法，有的关注语义），提升注意力能力。
专家模块（`Expert`）	单个“专家”，本质是简单的神经网络（MLP）；负责“分工处理”不同类型的文本特征（如处理戏剧对话中“角色名格式”“重复回应”等特定模式）。
噪声Top-K路由（`NoisyTopkRouter`）	模型的“调度中心”： 1. 根据输入文本的特征，计算每个“专家”的适配分数； 2. 加入噪声避免调度固化，选择`top_k`个最适配的专家； 3. 输出专家权重（用于加权输出）和选中的专家索引。
稀疏混合专家（`SparseMoE`）	整合“路由”和“专家”： 1. 按路由结果，将每个文本token分配给`top_k`个专家； 2. 限制单个专家处理的token数量（容量控制），避免过载； 3. 加权汇总专家输出，得到最终结果（仅激活部分专家，节省计算）。
Transformer块（`Block`）	模型的“基本 building block”：将“多头注意力（通信，捕捉上下文）”和“稀疏混合专家（计算，处理特征）”结合，再加上残差连接和层归一化，避免训练梯度消失，提升模型稳定性。
语言模型封装（`SparseMoELanguageModel`）	整合所有组件的“完整模型”： 1. 词嵌入+位置嵌入（将字符ID转换为带位置信息的向量）； 2. 堆叠多个`Block`（默认8层）； 3. 最终线性层（`lm_head`）输出“下一个字符”的预测概率； 4. 内置`generate`方法，实现文本生成。

三、训练与评估模块：让模型“学会”文本规律

负责模型的初始化、训练迭代和效果监控，确保模型能从数据中学习（如戏剧对话的角色格式、重复回应模式）。包含3个关键组件：

组件名称（代码对应）	核心功能
权重初始化（`kaiming_init_weights`）	用Kaiming正态分布初始化模型中线性层的权重，避免训练初期梯度过小或过大，加速模型收敛。
损失估计（`estimate_loss`）	1. 禁用梯度计算（节省资源），切换模型到“评估模式”； 2. 批量计算训练集/验证集的平均交叉熵损失（衡量“预测下一个字符”的准确率）； 3. 监控模型是否过拟合（如验证损失持续上升，说明模型学错了“噪声”）。
主训练流程（`main`）	1. 初始化模型、优化器（AdamW，用于更新参数）； 2. 循环迭代训练（默认5000次）： - 取批量数据→计算损失→反向传播更新参数； - 定期（每100次迭代）评估损失并打印； 3. 训练结束后调用`generate`，生成符合数据集风格的文本（如戏剧对话）。

四、工具/辅助模块：支撑核心流程运行

包含2个基础辅助组件，确保代码可复现和功能完整：

组件名称（代码对应）	核心功能
超参数配置（开头变量）	定义模型和训练的关键参数（如批量大小`batch_size`、专家数量`num_experts`、学习率`learning_rate`等），方便调整和实验。
文本解码（`decode`）	将模型输出的“整数ID序列”转换回原始文本（如`[1,2,3]`→“Fir”），用于训练后查看生成结果。