shuiyuPyLab/BitFlow-MoE

0 Fork 0

Files

T

shuiyu67 586917e75c Initial commit: BitFlow MoE - 微型中文 MoE 语言模型 for micro:bit

2026-07-26 02:04:51 +08:00

6.8 KiB

Raw Blame History

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 乘积，⊗ 是外积

关键优化：

不使用矩阵求逆，而是通过标量除法实现
状态矩阵 S 只存储 d×d=40×40=1600 个浮点数（6.4KB）
每次更新的复杂度是 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 路由器设计

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. 局限性

GLA 的衰减机制：log_decay 是标量，无法按维度自适应（Mamba 可以做）
硬路由：每次只激活一个专家，无法融合多个专家的知识
词表太小：149 字符覆盖场景有限
无上下文理解：48 token 窗口限制了长对话能力
未做消融实验：无法精确量化每个模块的贡献

6.8 KiB Raw Blame History Unescape Escape