一个基于稀疏混合专家模型（Sparse Mixture of Experts, Sparse MoE） 的 Transformer 语言模型
flyfish
用稀疏混合专家（Sparse MoE）替代传统Transformer的全连接层
py
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 out
class 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 out
class 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)
# 投影并应用dropout
out = self.dropout(self.proj(out))
return out
class 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_k
self.topkroute_linear = nn.Linear(n_embed, num_experts) # 用于计算路由logits
self.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, indices
class 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_k
self.capacity_factor = capacity_factor # 控制专家容量的因子
self.num_experts = num_experts
def 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)
# 展平批量和序列维度，以便独立处理每个token
flat_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_k
expert_capacity = int((tokens_per_batch / self.num_experts) * self.capacity_factor)
updates = torch.zeros_like(flat_x)
# 每个专家处理分配给它的token
for i, expert in enumerate(self.experts):
# 找到分配给当前专家的token
expert_mask = (indices == i).any(dim=-1)
flat_mask = expert_mask.view(-1)
selected_indices = torch.nonzero(flat_mask).squeeze(-1)
# 应用容量限制，防止单个专家处理过多token
if selected_indices.numel() > expert_capacity:
selected_indices = selected_indices[:expert_capacity]
if selected_indices.numel() > 0:
# 专家处理选中的token
expert_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_output
class 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_head
self.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 x
class 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) # 语言模型头，预测下一个token
def 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 = None
else:
B, T, C = logits.shape
logits = logits.view(B*T, C) # 展平以便计算交叉熵
targets = targets.view(B*T)
loss = F.cross_entropy(logits, targets)
return logits, loss
def generate(self, idx, max_new_tokens):
"""生成新的token序列"""
# idx是当前上下文的索引数组，形状为 (B, T)
for _ in range(max_new_tokens):
# 将idx裁剪到最后block_size个token
idx_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 idx
def 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"），用于训练后查看生成结果。 |
模块协作
- 数据处理模块将原始对话文本→数值序列→批量数据，喂给模型；
- 模型核心模块接收批量数据，通过"注意力捕捉上下文+专家分工计算"，输出下一个字符的预测；
- 训练评估模块根据预测结果计算损失，更新模型参数，并监控训练效果；
- 训练完成后，模型核心模块 的
generate方法结合解码工具，生成新的戏剧风格文本。