Files
BitFlow-MoE/README.md
T

178 lines
5.8 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# BitFlow MoE — 面向 micro:bit 的混合专家微型中文语言模型
[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
[![micro:bit](https://img.shields.io/badge/target-micro:bit-green)](https://microbit.org/)
[![Chinese NLP](https://img.shields.io/badge/NLP-Chinese-red)](https://www.sohu.com/a/123456789)
[![MoE](https://img.shields.io/badge/architecture-MoE-orange)](https://arxiv.org/abs/2209.01652)
一个完全自研的微型语言模型,可以在 **BBC micro:bit**(ARM Cortex-M4, 256KB Flash, 16KB RAM)上运行。支持中文对话、三种专家路由、int8 量化。
这是目前已知**第一个在 micro:bit 上运行的 MoE 中文语言模型**。
## 项目亮点
- **硬件极限**:在仅 16KB RAM、256KB Flash 的 micro:bit 上运行完整推理
- **架构创新**:门控线性注意力 (GLA) + 混合专家 (MoE) + int8 量化,三管齐下
- **中文支持**:149 字符词表,涵盖 258 条基础问答 + 动物事实 + 逻辑推理
- **完全自研**:从训练、量化到推理全链路自主设计,零外部预训练模型依赖
- **工业级优化**:内存占用 78.3KB(micro:bit 可用 108KB),余量 27.5%
## 架构总览
```
输入字符 → Token Embedding → 位置编码
↓
[路由器: 运行均值 → D→32→3]
↓ 选择最优专家
┌───────────┼───────────┐
↓ ↓ ↓
Expert-0 Expert-1 Expert-2
(事实记忆) (逻辑推理) (基础问答)
↓ ↓ ↓
GLA Block + FFN × 2层
↓
LayerNorm → 输出头
↓
采样 → 下一字符
```
### 核心组件
| 组件 | 技术 | 参数量 |
|------|------|--------|
| Token + Position Embedding | 查找表 | 7,880 |
| 3 × GLA Block (2层) | 门控线性注意力 | 42,720 |
| FFN (隐藏层 80) | GELU 激活 | 6,880 |
| MoE 路由器 | ReLU(32) → Softmax(3) | 1,411 |
| 输出头 | Linear(40→149) | 5,960 |
| **总计** | | **59,131** |
## 性能数据
- **训练**:10,000 epochs,最终 loss = 0.0861
- **路由准确率**:500 epoch 达到 100%(34% → 100%)
- **内存占用**:78.3KB / 108KB 可用(余量 29.7KB / 27.5%)
- **权重体积**:int8 量化后仅 58KB
- **推理速度**:单个字符约 200-500ms(取决于上下文长度)
## 快速开始
### 环境要求
- **PC 训练**:Python 3.8+、PyTorch 1.12+
- **硬件部署**:BBC micro:bit v2(nRF52833)
- **串口工具**:任意串口终端,或本项目 `chat_moe.py`
### 训练模型
```bash
# 1. 安装依赖
pip install torch
# 2. 训练
python train_bitflow_moe.py
# 3. 验证推理
python verify_moe_v4.py
```
### 部署到 micro:bit
```bash
# 1. 部署权重和代码
python deploy_moe.py
# 2. 通过串口对话
pip install pyserial
python chat_moe.py
```
### 使用预训练权重
可以直接下载已训练好的权重文件 `bitflow_moe_weights.bin`(58KB int8),配合 `bitflow_moe.py` 直接部署到 micro:bit。
## 目录结构
```
BitFlow-MoE/
├── train_bitflow_moe.py # 训练脚本(含数据、模型定义)
├── bitflow_moe.py # micro:bit 推理代码(MicroPython)
├── verify_moe_v4.py # PC 端模拟器(验证 micro:bit 推理)
├── chat_moe.py # 串口对话工具(Python)
├── deploy_moe.py # 部署脚本(自动查找 micro:bit)
├── plot_training.py # 训练过程可视化
├── training_analysis.png # 训练过程图(中文)
├── bitflow_moe_weights.bin # 预训练 int8 权重(58KB)
├── bitflow_moe.pt # PyTorch 完整权重(245KB)
├── architecture.md # 架构设计文档
├── quickstart.md # 用户快速上手指南
├── LICENSE # MIT License
└── README.md # 本文件
```
## 三个专家
模型包含三个专用专家,由路由器自动选择:
| 专家 | 功能 | 训练数据量 |
|------|------|-----------|
| Expert-0 (事实记忆) | 事实问答、属性判断 | 7,598 tokens |
| Expert-1 (逻辑推理) | 演绎推理、类比推理 | 7,312 tokens |
| Expert-2 (基础问答) | 问候、礼貌、自我介绍 | 43,248 tokens |
## 训练图表
![训练过程](training_analysis.png)
从上到下:
1. **总损失** — 5.8 → 0.086(下降 67 倍)
2. **预测损失** — 预测能力从 5.2 降到 0.076
3. **路由损失** — 专家选择从随机到完美
4. **路由准确率** — 34% → 100%(仅 500 轮)
5. **负载均衡** — 三个专家使用比例均衡
6. **分阶段对比** — 早期/中期/QAT 量化阶段
## 硬件要求
| 平台 | RAM 可用 | 实际占用 | 余量 |
|------|----------|----------|------|
| micro:bit (nRF52833) | ~108KB | 78.3KB | 29.7KB (27.5%) |
| STM32F4 | ~192KB | 78.3KB | 113.7KB |
| ESP32 | ~360KB | 78.3KB | 281.7KB |
## 已知限制
- 词表 149 字符(覆盖动物、自然、基础对话场景)
- 上下文窗口 48 个 token
- 单次生成速度受限于 micro:bit CPU 频率(64MHz)
- 未进行消融实验对照(单专家 vs MoE、GLA vs Transformer)
## 后续方向
- [ ] RWKV 或 Mamba 替代 GLA(更少的状态维护)
- [ ] 扩展词表到常用汉字(3000-5000 字符)
- [ ] 增加更多专家(情感分析、数学推理)
- [ ] 对比实验(消融研究)
- [ ] 移植到 STM32 / ESP32
## 贡献
欢迎提交 Issue 和 Pull Request!
## 引用
如果你在本研究中使用了 BitFlow MoE,请引用:
```
@misc{bitflow_moe_2026,
author = {shuiyu67},
title = {BitFlow-MoE: 面向 micro:bit 的混合专家微型中文语言模型},
year = {2026},
url = {https://github.com/shuiyu67/BitFlow-MoE}
}
```
## 许可证
[MIT License](LICENSE) — 可以自由使用、修改和分发。