forthezero/vAIbe_diffutslator

Diffutslator

基于扩散模型的中英互译系统。使用非自回归并行生成，通过DDIM加速推理

原理

扩散翻译的核心思想

传统翻译模型（如Transformer）是自回归的，逐token生成。扩散模型则是非自回归的，并行生成所有token：

自回归:  [SOS] → [token1] → [token2] → [token3] → [EOS]
                  ↓           ↓           ↓
扩散:    噪声 ──同时去噪──→ 完整句子（一步生成所有token）

双向翻译架构

┌─────────────────────────────────────────────────────────────────────┐
│                           噪声空间 (共享)                            │
│                             [L × D]                                 │
│                                                                     │
│    中文嵌入 ──前向扩散(q_sample)──→ 噪声 ←──前向扩散── 英文嵌入      │
│                    ↓                       ↓                        │
│              中文去噪网络            英文去噪网络                    │
│                    ↓                       ↓                        │
│              中文逆扩散              英文逆扩散                      │
│                    ↓                       ↓                        │
│                中文输出               英文输出                       │
└─────────────────────────────────────────────────────────────────────┘

翻译流程

以 中译英 为例：

1. 编码: 中文句子 → 中文token → 中文嵌入向量
2. 前向扩散: 中文嵌入添加噪声到指定时间步（或到纯噪声）
3. 逆扩散去噪:
   • 前半段：用中文去噪网络（保持源语言特征）
   • 后半段：切换到英文去噪网络（转向目标语言）
4. 解码: 最终嵌入 → 英文token → 英文句子

为什么扩散能做翻译？

扩散过程将数据逐步加噪变成纯噪声，逆扩散则从噪声恢复数据。关键洞察：

• 两种语言嵌入经过充分加噪后，在噪声空间中变得"不可区分"
• 从这个共享噪声空间出发，用不同语言的去噪路径，可以恢复到不同语言
• 类比：把中文和英文都"打散"成同样的积木，再用英文的说明书拼回去

安装

依赖

pip install torch tqdm

硬件要求

• CPU训练可用（本项目针对CPU优化）
• 内存：至少4GB
• 推荐：GPU可大幅加速

快速开始

训练

# 快速验证模式（1000条数据，5轮）
python train.py --quick

# 完整训练
python train.py

# 从检查点续训
python train.py --resume checkpoints/epoch_1.pt

训练中按 Ctrl+C 可安全中断，自动保存 checkpoints/interrupted.pt。

推理

# 中译英
python inference.py --text "你好世界" --zh

# 英译中
python inference.py --text "Hello world" --en

# 交互模式
python inference.py --interactive

详细使用

训练命令

# 基本训练
python train.py

# 快速验证（小数据集，少轮次）
python train.py --quick

# 从检查点续训
python train.py --resume checkpoints/best.pt

# 使用更多数据
python train.py --max-samples 10000

# 指定轮次和批量
python train.py --epochs 20 --batch-size 32

推理命令

# 基本推理（中译英）
python inference.py --text "今天天气很好" --zh

# 英译中
python inference.py --text "The weather is nice today" --en

# 使用DDPM（更慢但可能更准）
python inference.py --text "你好" --zh --ddpm

# 交互模式
python inference.py --interactive

# 指定检查点
python inference.py --text "你好" --zh --checkpoint checkpoints/best.pt

# 静默模式（不显示扩散过程）
python inference.py --text "你好" --zh --quiet

配置参数

模型配置 (ModelConfig)

参数	默认值	说明
d_model	256	嵌入维度，影响模型容量
n_heads	4	多头注意力头数
n_layers	4	Transformer编码器层数
d_ff	512	前馈网络隐藏层维度
max_len	128	最大序列长度
dropout	0.1	Dropout比率
vocab_size_zh	8000	中文词表大小
vocab_size_en	8000	英文词表大小

扩散配置 (DiffusionConfig)

参数	默认值	说明
timesteps	1000	训练时的扩散总步数
ddim_steps	50	DDIM推理采样步数
beta_start	0.0001	噪声调度起始值
beta_end	0.02	噪声调度结束值

训练配置 (TrainingConfig)

参数	默认值	说明
batch_size	64	批量大小
learning_rate	1e-4	学习率
weight_decay	0.01	权重衰减
warmup_steps	500	学习率预热步数
epochs	10	训练轮次
save_every	1	每N轮保存检查点

数据配置 (DataConfig)

参数	默认值	说明
max_samples	None	最大样本数（None=全部）
min_len	2	最小句子长度
max_len	128	最大句子长度

架构说明

分词器 (tokenizer.py)

使用BPE（Byte Pair Encoding）算法：

• 中文: 字符级为主，BPE处理罕见词和数字
• 英文: 标准BPE子词分割
• 词表大小：各8000 tokens
• 特殊token: <pad>, <sos>, <eos>, <unk>, <mask>

# 示例
tokenizer_zh.encode("你好世界")  # [123, 456, 789]
tokenizer_en.encode("hello world")  # [234, 567]

嵌入层 (embedding.py)

class LanguageEmbedding:
    token_embedding    # [vocab_size, d_model]
    position_embedding # [max_len, d_model]
    length_embedding   # [max_len, d_model]

将离散token转换为连续向量，加入位置信息。

噪声预测网络 (model.py)

class DiffusionTransformer:
    """基于Transformer的噪声预测网络"""
    
    # 输入: x_t [batch, len, d_model], t [batch], lang [str]
    # 输出: predicted_noise [batch, len, d_model]
    
    # 结构:
    # 1. 时间步嵌入 (sinusoidal)
    # 2. 语言特定输入投影
    # 3. N层 Transformer blocks
    # 4. 语言特定输出投影

扩散过程 (diffusion.py)

# 前向扩散（加噪）
x_t, noise = diffusion.q_sample(x_0, t)  # x_0 → x_t

# 反向扩散（去噪）
x_t_minus_1 = diffusion.p_sample(x_t, t, predicted_noise)

使用线性噪声调度，支持DDIM加速采样。

语言切换器 (switcher.py)

class LanguageSwitcher:
    """判断当前噪声状态更接近哪种语言"""
    
    # 输入: x_t [batch, len, d_model]
    # 输出: lang_prob [batch, 2]  # [中文概率, 英文概率]

在推理时判断何时切换去噪路径。

文件结构

diffutslator/
├── config.py       # 超参数配置
├── tokenizer.py    # BPE分词器
├── embedding.py    # 嵌入层
├── model.py        # 噪声预测网络 (Transformer)
├── diffusion.py    # 扩散过程 + DDIM采样
├── switcher.py     # 语言切换分类器
├── dataset.py      # 数据加载（流式）
├── train.py        # 训练脚本
├── inference.py    # 推理脚本
├── main.py         # 主入口
├── utils.py        # 工具函数
├── .cache/         # 分词器缓存
│   ├── tokenizer_zh.json
│   └── tokenizer_en.json
└── checkpoints/    # 模型检查点
    ├── best.pt
    ├── epoch_1.pt
    └── interrupted.pt

数据集

• _dataset/cveto/
• _dataset/tatoeba.tsv

---

上面是AI生成的，我到这补充一下

生成这个项目的模型是GLM-5，用iflow cli，在我的电脑上训练了九个半小时，用了2.8w条数据，权重在checkpoints下