WSAI-Math：基于严格验证蒸馏数据训练的 0.8B 数学推理模型

上一篇文章介绍了 WSAI-Math-Distill——3,867 条逐条判分验证、对评测集 8-gram 去污染的数学长思维链蒸馏数据集。现在，基于该数据集训练的模型 WSAI-Math 发布了。

模型基于 Qwen3.5-0.8B，以 DeepSeek deepseek-v4-flash 为教师做长思维链蒸馏冷启动，LoRA SFT 50 个优化步，merge_and_unload 合并导出 BF16 全量权重。这次实验要回答的问题很明确：高纯度蒸馏数据 + 极小训练量，能把 0.8B 基座的数学推理行为改造到什么程度。

模型信息

项	值
基座	Qwen3.5-0.8B（0.85B 参数，混合线性注意力架构，原生上下文 262,144）
参数量	878,544,960（可训练 25,559,040，占 2.91%）
权重精度	bfloat16（LoRA 合并后的全量权重，model.safetensors 1.7GB）
训练方法	LoRA SFT → merge_and_unload 导出
训练量	50 个优化步（有效 batch 16，共见约 800 条样本）
验证损失	val_loss = 0.711（固定 1% 验证集，与训练 loss 同步下降）

基座 Qwen3.5-0.8B 采用 Qwen3-Next 风格的混合架构：Gated DeltaNet 线性注意力与全注意力按 3:1 交错，24 层，GQA 8H/2KV，head_dim 256，词表 248,320。

训练数据

直接取自 WSAI-Math-Distill v1.0：

OpenR1-Math-220k（93,733 题）
  │ 题面清洗 + 来源筛选（amc_aime / olympiads / aops_forum）
  ▼
8-gram 去污染 vs AIME2024 + AIME2025 + HLE 文本全集
  │ 指纹 264,265 个，剔除 3,444 题
  ▼
干净题库（6,245 题）
  │ v4-flash 逐题蒸馏（thinking 模式，max_tokens=16384）
  ▼
严格判分 —— 61% 通过 → 3,867 条验证正确
  │ 固定种子抽样 1,410 条 + 质量过滤（reasoning 800–30,000 字符）
  ▼
LoRA SFT（50 步）→ 合并导出 → WSAI-Math

判分协议在数据集那篇文章里已详细说明，核心是三条：只取最后一个 \boxed{...} 作预测；精确字符串 / math_verify 无预处理 / 纯数字精确相等三判全过才记正确；零规范化，格式不同即判错。训练数据与评测使用同一套口径。

训练配置

• LoRA：r=64，α=128，dropout=0，注入全部 q/k/v/o/gate/up/down 投影；
• 损失掩码：仅对 assistant 段计算（<think> 思维链 + 解答），prompt 全掩码；
• 优化器：AdamW（β=0.9/0.95，eps=1e-8，lr=1e-4，无 weight decay），梯度裁剪 1.0；
• 批次：有效 batch 16 = micro-batch 2 × 梯度累积 8；
• 损失轨迹：train ema20 从 1.10 降至 0.64–0.73 区间，val_loss 0.711，无过拟合。

一个值得单独记录的工程细节：大词表的 logits 分块计算。Qwen3.5 词表 248,320，全量 materialize 一层 logits 的峰值内存约 8GB。实现上将 logits 按 chunk=1024 分块计算，并用 gradient checkpoint 包裹反向重算，峰值内存压到 1GB 以下。这是在消费级硬件上训练大词表小模型时必须处理的问题。

评测结果

协议：thinking 模式，Qwen 官方采样参数（temperature=1.0, top_p=0.95, top_k=20, min_p=0），每题 1 采样，最大生成 5120 token，本地 CPU bf16 推理。

AIME

基准	WSAI-Math	基座 Qwen3.5-0.8B	提升幅度
AIME 2024 (avg@1)	3.3%	0.0%	+3.3pp
AIME 2025 (avg@1)	0.0%	0.0%	—

行为学对比

数字分数之外，行为层面的差异更显著：

观测项	WSAI-Math	基座 Qwen3.5-0.8B
AIME24/25 截断率	43% / 77%	100% / 100%（60/60 全部打转至 5120 上限）
平均输出长度	3,552 / 4,328 token	5,120 / 5,120 token（全部打满）
思维链风格	教师化：简练、复核关键计算、\boxed{} 收尾、干净终止	大量 “Wait…” 原地循环，无法收敛
简单数学（17×23）	146 token 直出 \boxed{391}	数百 token 徘徊后勉强作答

50 个优化步、约 800 条样本、2.91% 可训练参数，基座从「100% 截断、无法收敛」转变为「生成完整思维链并以 \boxed{} 干净终止」。AIME 2024 上 3.3% 对 0% 的提升也发生在这个训练量下。对 0.8B 参数量的模型，AIME 本身难度偏高，这个结果的参考价值主要在行为学对比，而非绝对分数。

使用方法

Ollama

ollama run wangshen/wsai-math

transformers

from transformers import AutoTokenizer, AutoModelForImageTextToText
import torch

tok = AutoTokenizer.from_pretrained("WSAI-Math")
model = AutoModelForImageTextToText.from_pretrained(
    "WSAI-Math", dtype=torch.bfloat16, device_map="cpu")

messages = [{"role": "user", "content": "您的数学题"}]
text = tok.apply_chat_template(messages, tokenize=False,
                               add_generation_prompt=True, enable_thinking=True)
inputs = tok(text, return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=5120,
                     do_sample=True, temperature=1.0, top_p=0.95, top_k=20)

思维链位于 <think>...</think>，最终答案以 \boxed{...} 给出。

致谢与许可

• 基座：Qwen/Qwen3.5-0.8B
• 教师：DeepSeek deepseek-v4-flash
• 训练数据：WSAI-Math-Distill（题源 open-r1/OpenR1-Math-220k，Apache-2.0）

模型采用 Apache-2.0 许可。