PLAN.md

# 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)

```python
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)

**前向扩散（加噪）**：
```python
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
```

**反向扩散（去噪）**：
```python
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)

```python
class DiffusionTransformer:
    """预测噪声的Transformer"""
    - 输入: x_t [batch, len, d_model], t [batch]
    - 输出: predicted_noise [batch, len, d_model]
    
    结构:
    - 语言特定的输入投影
    - 时间步嵌入 (sinusoidal)
    - N层 Transformer blocks
    - 语言特定的输出投影
```

**参数**（CPU优化）：
- `n_layers = 4`
- `n_heads = 4`
- `d_ff = 512`
- 总参数量：约 2M

### 2.5 语言切换器 (switcher.py)

```python
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. 反向传播更新参数
```

**损失函数**：
```python
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: 基础框架（确保可训练）

1. **配置文件** - 定义所有超参数
2. **分词器** - 实现中英文分词
3. **数据集** - 加载tatoeba数据
4. **嵌入层** - 简单的token嵌入
5. **扩散核心** - 前向和反向扩散
6. **简单模型** - 基础噪声预测网络
7. **训练脚本** - 带进度条的训练循环

**验证目标**：能在少量数据上跑通训练，loss下降

### Phase 2: 完整架构

1. **语言切换器** - 实现切换判断
2. **变长处理** - 实现长度嵌入
3. **完整模型** - 整合所有模块
4. **推理脚本** - 可视化扩散过程

**验证目标**：完整训练流程，能输出翻译结果

### Phase 3: 优化加速

1. **DDIM采样** - 减少推理步数
2. **训练加速** - 混合精度、梯度累积
3. **模型调优** - 调整超参数

**验证目标**：提升训练和推理速度，改善翻译质量

## 五、训练策略

### 快速验证模式

```bash
# 使用tatoeba前1000条数据
# batch_size=8, epochs=10
python train.py --quick --samples 1000
```

### 完整训练模式

```bash
# 使用全部数据
# 支持暂停/继续
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环境的优化：

1. **小模型**：参数量控制在2-5M
2. **小批量**：batch_size = 4-16
3. **梯度累积**：模拟更大batch
4. **简单架构**：减少层数和维度
5. **内存优化**：及时释放中间变量

## 八、预期效果

| 指标 | 目标 |
|------|------|
| 训练速度 | 1-2 it/s (CPU) |
| 推理速度 | 1-5秒/句 (DDIM 50步) |
| 翻译质量 | 简单句子可理解 |
| 模型大小 | < 50MB |

---

*计划制定完成，待用户确认后开始实现*