214 lines
6.8 KiB
Markdown
214 lines
6.8 KiB
Markdown
# BitFlow MoE 架构设计文档
|
||
|
||
## 1. 设计目标
|
||
|
||
在 BBC micro:bit（ARM Cortex-M4, 64MHz, 16KB RAM, 256KB Flash）上运行一个具有**理解能力**的微型中文语言模型。
|
||
|
||
核心挑战：
|
||
- **内存限制**：可用 RAM ~108KB，模型权重 + 推理状态必须在这个范围内
|
||
- **计算限制**：64MHz CPU，无 FPU（但有硬件乘法）
|
||
- **存储限制**：256KB Flash，预训练模型通常几百 MB，需极致压缩
|
||
|
||
## 2. 架构选型
|
||
|
||
### 2.1 为什么不用 Transformer？
|
||
|
||
标准 Transformer（即使是 2 层微型版）在 micro:bit 上存在问题：
|
||
|
||
| 组件 | 标准 Attention | GLA (我们的选择) |
|
||
|------|---------------|------------------|
|
||
| 注意力矩阵 | O(n²) 状态矩阵 | O(d) 累积矩阵 |
|
||
| 最大上下文 | ~16 tokens | 48 tokens |
|
||
| 内存峰值 | 32KB+（仅注意力） | <5KB |
|
||
| 推理速度 | 极慢（会卡死） | 实时 |
|
||
|
||
**结论**：标准 Attention 的 O(n²) 状态维护在 micro:bit 上是灾难性的。GLA 将复杂度降到 O(n×d)，状态矩阵从 n×n 压缩到 d×d。
|
||
|
||
### 2.2 为什么不用纯 RNN？
|
||
|
||
RNN 虽然内存 O(d)，但：
|
||
- 无法并行训练（梯度消失/爆炸）
|
||
- 长程依赖能力弱
|
||
- 无法做到像我们这样的"理解"效果
|
||
|
||
GLA 是 RNN 的改进版，保留了 O(d) 内存，但增加了**可并行训练**的能力（通过离散时间注意力）。
|
||
|
||
## 3. 核心架构
|
||
|
||
### 3.1 门控线性注意力 (GLA)
|
||
|
||
GLA 的核心公式：
|
||
|
||
```
|
||
S[t] = g[t] ⊙ S[t-1] + k[t] ⊗ v[t]
|
||
z[t] = g[t] ⊙ z[t-1] + k[t]
|
||
y[t] = Q[t] · S[t] · (Q[t] · z[t])⁻¹
|
||
```
|
||
|
||
其中：
|
||
- **S** 是 d×d 累积状态矩阵（每次更新 O(d²)）
|
||
- **k, q, v** 是通过 Linear 投影得到的向量
|
||
- **g** 是门控（gate），通过 sigmoid(log_decay) 控制衰减
|
||
- **⊙** 是 Hadamard 乘积，**⊗** 是外积
|
||
|
||
**关键优化**：
|
||
1. 不使用矩阵求逆，而是通过标量除法实现
|
||
2. 状态矩阵 S 只存储 d×d=40×40=1600 个浮点数（6.4KB）
|
||
3. 每次更新的复杂度是 O(d²) 而非 O(n²)
|
||
|
||
### 3.2 混合专家 (MoE)
|
||
|
||
```
|
||
输入 → [Token Embedding] + [Position Embedding]
|
||
↓
|
||
[路由器: mean pooling → FC → argmax]
|
||
↓ 选择最优专家
|
||
┌── Expert-0 (事实记忆)
|
||
├── Expert-1 (逻辑推理)
|
||
└── Expert-2 (基础问答)
|
||
↓
|
||
[LayerNorm → GLA Block × 2]
|
||
↓
|
||
[LayerNorm → Output Head]
|
||
↓
|
||
采样输出
|
||
```
|
||
|
||
**为什么 MoE 有效**：
|
||
- 每个专家专精一类任务，不需要一个"全能"模型
|
||
- 路由选择只激活一个专家，**推理计算量不增加**
|
||
- 三个专家共享相同的结构（GLA Block × 2），代码复用率高
|
||
|
||
### 3.3 路由器设计
|
||
|
||
```python
|
||
def router(x):
|
||
# 1. 对所有 token 做 mean pooling → 固定维度 d=40
|
||
x_mean = x.mean(dim=1)
|
||
# 2. 全连接层: 40 → 32 → 3
|
||
h = relu(fc1(x_mean)) # 40 → 32
|
||
logits = fc2(h) # 32 → 3
|
||
return logits # argmax 选专家
|
||
```
|
||
|
||
**训练时**：使用交叉熵损失监督路由（labels 来自数据类别）
|
||
**推理时**：只计算一次路由，选择最优专家
|
||
|
||
## 4. 量化策略
|
||
|
||
### 4.1 int8 量化
|
||
|
||
```
|
||
scale = max(|w|) / 127
|
||
q_int8 = round(w / scale) # [-128, 127]
|
||
w_rec = q_int8 × scale # 反量化
|
||
```
|
||
|
||
- **量化精度**：int8（8-bit 有符号整数）
|
||
- **量化感知训练 (QAT)**：训练后 2000 步（8000-10000 epoch）引入量化噪声，让模型适应 int8 精度
|
||
- **结果**：loss 从 0.0861 仅上升到 0.0863，影响极小
|
||
|
||
### 4.2 内存分析
|
||
|
||
| 组件 | 全精度(float32) | int8 量化 | 压缩比 |
|
||
|------|----------------|-----------|--------|
|
||
| 模型权重 | 236.5KB | 57.7KB | **4.1×** |
|
||
| 推理状态 | ~70KB | ~70KB | 无（运行时需要float） |
|
||
| **总计** | **306.5KB** | **127.7KB** | **2.4×** |
|
||
|
||
## 5. 数据工程
|
||
|
||
### 5.1 三层数据体系
|
||
|
||
```
|
||
基础层 (BASE, 69 条)
|
||
└── 动物行为、颜色判断、位置关系
|
||
↓ 提供基本语义
|
||
事实层 (FACT_QA, 76 条)
|
||
└── 什么是XX、XX吃什么、XX在哪里
|
||
↓ 提供知识推理
|
||
问答层 (CHAT_QA, 258 条)
|
||
└── 问候、礼貌、自我介绍、喜好、情感、能力
|
||
↓ 提供交互能力
|
||
```
|
||
|
||
### 5.2 数据增强
|
||
|
||
- 基础层 × 20 遍
|
||
- 事实层 × 40 遍
|
||
- 问答层 × 80 遍
|
||
|
||
总数据量：290,780 tokens，其中问答类占 74%
|
||
|
||
## 6. 训练流程
|
||
|
||
```
|
||
[10,000 epochs]
|
||
│
|
||
├─ 0 - 8000 epoch: 标准训练
|
||
│ ├─ 学习率: 3e-3 (Cosine Annealing)
|
||
│ ├─ Loss = 预测损失 + 0.5 × 路由损失 + 0.01 × 负载均衡损失
|
||
│ └─ 路由准确率: 34% → 100%
|
||
│
|
||
└─ 8000 - 10000 epoch: QAT 量化感知训练
|
||
├─ 每步对权重做 int8 量化 + 反量化
|
||
├─ 模型适应量化噪声
|
||
└─ loss 微升后稳定
|
||
```
|
||
|
||
## 7. 推理流程
|
||
|
||
```
|
||
[初始化] 加载 int8 权重 → 预解量化小权重 (SMW)
|
||
│
|
||
├─ [种子处理] 逐 token 前向传播，积累运行均值
|
||
│ └─ 路由器根据均值选择专家，激活对应 GLA Block
|
||
│
|
||
├─ [生成循环]
|
||
│ ├─ 计算 token embedding
|
||
│ ├─ 累加运行均值 → 路由 → 选专家
|
||
│ ├─ 专家 GLA Block × 2 层
|
||
│ ├─ LayerNorm → 输出头
|
||
│ ├─ 采样 (temperature=0.3)
|
||
│ └─ 更新频率惩罚 (bcnt)
|
||
│
|
||
└─ [输出] 字符 ID 列表 → 串口发送
|
||
```
|
||
|
||
## 8. micro:bit 内存布局
|
||
|
||
```
|
||
128KB Flash (总容量)
|
||
├── ~20KB: MicroPython 解释器 + 标准库
|
||
├── ~58KB: int8 权重文件
|
||
├── ~7KB: 推理代码 (bitflow_moe.py)
|
||
└── ~43KB: 可用 (代码+数据段)
|
||
|
||
16KB RAM (总容量)
|
||
├── ~108KB: 实际可用 (nRF52833 有 256KB RAM, 但 MicroPython 用了一部分)
|
||
├── 58KB: 解量化后权重 (运行时)
|
||
├── 13KB: GLA 状态矩阵 (2层 × 40×40)
|
||
├── 4KB: 工作缓冲区
|
||
├── 2KB: 预解量化小权重
|
||
└── ~30KB: 余量 (防止栈溢出)
|
||
```
|
||
|
||
## 9. 设计决策总结
|
||
|
||
| 决策 | 选择 | 理由 |
|
||
|------|------|------|
|
||
| 注意力机制 | GLA | O(d) 内存，可并行训练 |
|
||
| MoE 结构 | 3 专家硬路由 | 推理计算量不增加 |
|
||
| 量化 | int8 + QAT | 体积压缩 4×，精度损失 <1% |
|
||
| 训练策略 | 两步训练 + QAT | 先学语义，再适应量化 |
|
||
| 路由方式 | 监督学习 | 保证每个专家都有任务 |
|
||
| 采样温度 | 0.3 | 偏低，保证输出一致性 |
|
||
|
||
## 10. 局限性
|
||
|
||
1. **GLA 的衰减机制**：log_decay 是标量，无法按维度自适应（Mamba 可以做）
|
||
2. **硬路由**：每次只激活一个专家，无法融合多个专家的知识
|
||
3. **词表太小**：149 字符覆盖场景有限
|
||
4. **无上下文理解**：48 token 窗口限制了长对话能力
|
||
5. **未做消融实验**：无法精确量化每个模块的贡献
|