Diffutslator 实现计划
基于扩散模型的中英互译系统
一、架构概述
┌─────────────────────────────────────────────────────────────────┐
│ 噪声空间 (共享) │
│ [L×D] │
│ ┌─────────────────────────┐ │
│ │ │ │
│ 中文扩散 ↗ 语言切换判断器 ↖ 英文扩散 │
│ (加噪) [分类器] (加噪) │
│ │ │ │
│ └─────────────────────────┘ │
│ ↓ ↓ │
│ 中文逆扩散 英文逆扩散 │
│ (去噪) (去噪) │
│ ↓ ↓ │
│ ┌────────────┐ ┌────────────┐ │
│ │ 中文解码器 │ │ 英文解码器 │ │
│ └────────────┘ └────────────┘ │
│ ↓ ↓ │
│ 中文输出 英文输出 │
└─────────────────────────────────────────────────────────────────┘
核心设计决策
| 问题 | 决策 | 理由 |
|---|---|---|
| 扩散空间 | 词嵌入连续空间 | 实现成熟、CPU友好、训练稳定 |
| 长度处理 | 变长序列 + 长度嵌入 | 扩散可变长，逆扩散收敛到目标长度 |
| 双向切换 | 可学习分类器 | 让模型自己判断何时切换 |
二、模块设计
2.1 分词器 (tokenizer.py)
中文分词：字符级 + BPE
- 字符级处理中文字符
- BPE处理罕见词和数字
英文分词：BPE
- 使用相同的BPE算法
- 与中文共享词表大小设置
词表：
- 中文词表：8000 tokens
- 英文词表：8000 tokens
- 特殊token：
<pad>,<sos>,<eos>,<mask>,<unk>
2.2 嵌入层 (embedding.py)
class LanguageEmbedding:
"""语言特定的嵌入层"""
- token_embedding: [vocab_size, d_model]
- position_embedding: [max_len, d_model]
- length_embedding: [max_len, d_model] # 长度编码
参数：
d_model = 256（CPU环境下适中）max_len = 128（最大序列长度）
2.3 扩散核心 (diffusion.py)
前向扩散（加噪）：
def forward_diffusion(x_0, t):
"""
x_0: 初始嵌入 [batch, len, d_model]
t: 时间步 [batch]
返回: x_t, noise
"""
# 线性噪声调度
alpha_t = 1 - t / T # 简化调度
noise = randn_like(x_0)
x_t = sqrt(alpha_t) * x_0 + sqrt(1 - alpha_t) * noise
return x_t, noise
反向扩散（去噪）：
def reverse_diffusion(x_t, t, model):
"""
x_t: 当前噪声状态
t: 当前时间步
model: 噪声预测网络
"""
predicted_noise = model(x_t, t)
x_t_minus_1 = denoise_step(x_t, predicted_noise, t)
return x_t_minus_1
时间调度：
- 训练时：T = 1000 步
- 推理时：DDIM加速，可降到 10-50 步
2.4 噪声预测网络 (model.py)
class DiffusionTransformer:
"""预测噪声的Transformer"""
- 输入: x_t [batch, len, d_model], t [batch]
- 输出: predicted_noise [batch, len, d_model]
结构:
- 语言特定的输入投影
- 时间步嵌入 (sinusoidal)
- N层 Transformer blocks
- 语言特定的输出投影
参数（CPU优化）：
n_layers = 4n_heads = 4d_ff = 512- 总参数量：约 2M
2.5 语言切换器 (switcher.py)
class LanguageSwitcher:
"""判断当前噪声更接近哪种语言"""
- 输入: x_t [batch, len, d_model]
- 输出: 语言概率 [batch, 2] # [中文, 英文]
结构:
- 全局平均池化
- 2层MLP
- Softmax输出
2.6 训练流程 (train.py)
训练步骤:
1. 加载中英平行句对 (zh, en)
2. 分别嵌入到连续空间
3. 随机采样时间步 t
4. 对中文嵌入做前向扩散到 t 步 → zh_t
5. 对英文嵌入做前向扩散到 t 步 → en_t
6. 训练噪声预测网络预测噪声
7. 训练切换器判断语言
8. 反向传播更新参数
损失函数：
L_total = L_noise_zh + L_noise_en + λ * L_switcher
L_noise: 噪声预测MSE损失
L_switcher: 语言分类交叉熵损失
2.7 推理流程 (inference.py)
中文→英文翻译：
1. 中文输入 → 中文嵌入
2. 完整前向扩散到纯噪声 (T步)
3. 迭代反向扩散:
for t in [T, T-1, ..., 1]:
- 切换器判断当前语言
- 若判断为中文→用中文去噪
- 若判断为英文→切换到英文去噪
- 输出当前步骤状态（可视化）
4. 最终噪声状态 → 英文解码 → 英文输出
英文→中文翻译：对称过程
三、文件结构
diffutslator/
├── TASK.md # 任务描述
├── PLAN.md # 本文件
├── config.py # 超参数配置
├── tokenizer.py # 分词器
├── embedding.py # 嵌入层
├── model.py # 扩散模型
├── diffusion.py # 扩散过程
├── switcher.py # 语言切换器
├── dataset.py # 数据集加载
├── train.py # 训练脚本
├── inference.py # 推理脚本
├── main.py # 主入口
├── utils.py # 工具函数
└── checkpoints/ # 模型检查点
四、实现步骤
Phase 1: 基础框架（确保可训练）
- 配置文件 - 定义所有超参数
- 分词器 - 实现中英文分词
- 数据集 - 加载tatoeba数据
- 嵌入层 - 简单的token嵌入
- 扩散核心 - 前向和反向扩散
- 简单模型 - 基础噪声预测网络
- 训练脚本 - 带进度条的训练循环
验证目标：能在少量数据上跑通训练，loss下降
Phase 2: 完整架构
- 语言切换器 - 实现切换判断
- 变长处理 - 实现长度嵌入
- 完整模型 - 整合所有模块
- 推理脚本 - 可视化扩散过程
验证目标：完整训练流程，能输出翻译结果
Phase 3: 优化加速
- DDIM采样 - 减少推理步数
- 训练加速 - 混合精度、梯度累积
- 模型调优 - 调整超参数
验证目标：提升训练和推理速度，改善翻译质量
五、训练策略
快速验证模式
# 使用tatoeba前1000条数据
# batch_size=8, epochs=10
python train.py --quick --samples 1000
完整训练模式
# 使用全部数据
# 支持暂停/继续
python train.py --full
# Ctrl+C 暂停，自动保存检查点
# python train.py --resume 继续训练
训练输出
Epoch 1/10: 100%|████████| 125/125 [02:30<00:00, loss=0.452]
预计剩余: 22:30 | 速度: 0.5 it/s
最新检查点: checkpoints/model_epoch1.pt
按 Ctrl+C 停止训练（自动保存）
六、推理展示
$ python inference.py --zh "你好世界"
翻译模式: 中文 → 英文
输入: 你好世界
扩散过程:
Step 1000: [噪声状态 - 切换器: 中文 95%]
Step 900: [噪声状态 - 切换器: 中文 78%]
Step 800: [噪声状态 - 切换器: 中文 52%]
Step 700: [噪声状态 - 切换器: 英文 61%] ← 语言切换!
Step 600: [噪声状态 - 切换器: 英文 89%]
...
Step 50: [接近完整句子 - 切换器: 英文 99%]
Step 1: [完整句子]
输出: Hello world
七、环境适配
针对CPU环境的优化：
- 小模型：参数量控制在2-5M
- 小批量：batch_size = 4-16
- 梯度累积：模拟更大batch
- 简单架构：减少层数和维度
- 内存优化：及时释放中间变量
八、预期效果
| 指标 | 目标 |
|---|---|
| 训练速度 | 1-2 it/s (CPU) |
| 推理速度 | 1-5秒/句 (DDIM 50步) |
| 翻译质量 | 简单句子可理解 |
| 模型大小 | < 50MB |
计划制定完成，待用户确认后开始实现