# 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. **未做消融实验**：无法精确量化每个模块的贡献