samaidev-pnet-2b — PonderNet × MiniCPM5-2B 循环思考推理
在 MiniCPM5-2B（Llama 架构）上实现 PonderNet 式循环思考：一个可学习的 halting 头（2049 参数）决定每个位置"还想几步"，思考块循环执行 1..K 次， 按停机分布 w 加权混合各步表示后继续解码。目标是让模型学会语义级 **"难题多想、简单题快停"**。
本仓是四轮实验 + 第五轮三方向的独立代码仓（权重与数据集在
tchbcb/MiniCPM5-2B-cpu 的
pondernet/ 子目录）。
核心结果一览
| 轮次 | 方案 | hard−easy 步数差 | 全局 CE | 结论 |
|---|---|---|---|---|
| R1 | 模板化短答案 CE | ~0 | — | CE 收益 38%，无步数分化 |
| R2 | 真实多步中文 CoT (1200) | 反向 | — | 收益 12%，hard 反而最少想 |
| R3 | 失败自举 (AdaptThink 式) | +0.012 | 0.376 | 方向转正，幅度不足 |
| R4 | 步数硬监督 + 冻结两阶段 | +0.498 ✓ | 0.6114 | 突破 0.3 终判线 |
| R5 | ①加权/目标消融 ②逐步 CE 回收 ③答案级 RL | 部分完成 | — | 见 docs/round5_results.md |
R5 已落地的消融结论：**分桶加权 (A1, +0.297, CE 0.590) 严格优于全局加压 (A2, +0.731 但 CE 崩至 1.279)**；easy 组权重压 4× 仅 −0.076 步 → 支持"线性头可分性上限"假说，容量扩展 (A4 MLP) 与表征适配 (B 逐步 CE) 是剩余候选解。
R4 逐样本语义验证：Spearman(步数, probe_ce)=0.362；probe_ce 四分位步数 严格单调 (2.756/2.913/3.133/3.199)；启发式桶反向对照 → 头学到的是 模型实测难度而非表面标签。
安装
pip install -r requirements.txt
pip uninstall -y torchao # peft 0.20 LoRA dispatch 兼容性 (T4 实测必需)
快速开始
from pondernet import PonderLlamaForCausalLM, PonderLlamaConfig
from safetensors.torch import load_file
model = PonderLlamaForCausalLM.from_ponder(
"tchbcb/MiniCPM5-2B-cpu",
ponder_kwargs={"ponder_signal": "learned", "max_ponder_steps": 6})
head = load_file("out_r4/ponder_head.safetensors")
model.ponder_head.load_state_dict({k.replace("ponder_head.", ""): v
for k, v in head.items()})
out = model.generate(**tokenizer("一只鸡每天下6个蛋...", return_tensors="pt").to(model.device))
代码结构
pondernet/
modeling_ponder.py # PonderLlamaForCausalLM: 思考块循环 + halting 头
# R5 新增: 逐步 CE / forced-steps / MLP 头 (config 开关)
train_ponder.py # 训练: head / head+lora / head+block 三模式
# R5 新增: --sup-weight-map --sup-target-preset --head-mlp --step-ce-weight
probe_difficulty.py # K=1 逐样本 CE 探针 → easy/medium/hard 实测难度标签
eval_ponder.py # 分桶步数表 + 全局 CE + 逐样本落盘 (--dump)
train_ponder_rl.py # 方向③: REINFORCE with baseline, 奖励 = 答案对 − λ·步数
make_gsm8k_zh_data.py # GSM8K 中文多步 CoT 数据构造
configs/
round4_repro.sh # R4 复现 (判据 +0.498)
round5_A_ablation.sh # 方向①: 加权/目标/容量消融
round5_B_step_ce.sh # 方向②: 逐步 CE 回收
round5_C_rl.sh # 方向③: 答案级 RL
dev/
r5_sweep.py # 交错流水线总控 (训练↔评估, 幂等断点)
r5_bootstrap.py # 新 T4 一键恢复 (环境+HF 拉取+内嵌 sweep)
r5_smoke.py # 三路径冒烟
docs/
round4_results.md # R4 完整实验记录
round5_results.md # R5 三方向结果
data_sample/ # 数据样例 (完整 720 train + 180 eval 在 HF pondernet/data_zh_round2 等)
方法要点
1. PonderNet 机制 (modeling_ponder.py)
每个位置独立 halting：思考块第 k 次执行后，halting 头输出 λ_k，
累计概率 cum ≥ 1−ε 即停；最终表示 h = Σ_k w_k·h_k + remaining·h_N。
训练损失 = CE(mix) + β·KL(w‖Geometric(p_g)) [+ γ·逐步 CE] [+ 监督项]。
2. 步数硬监督 (R4 突破的关键)
β·KL 只给"方向"且拉力不足（R3 分化仅 +0.012）。改为直接监督 w 分布：
loss_sup = CE(w*‖w) = −Σ w*_n·log w_n
easy → [1,0,..] medium → [.55,.30,.15,0..] hard → 均匀 1/K
配合冻结 LoRA 两阶段（--train-mode head --init-adapter），probe 标签
永不过时，2049 参数的 head 即可完成分化。
3. 第五轮三方向
- ① 消融：easy 目标 [1,0..] 已到 one-hot 极限仍停在 2.62 步 →
分桶加权 (
--sup-weight-map easy:8) / 全局加压 (--sup-weight 10) / 目标压尖 (--sup-target-preset sharp) / MLP 头 (--head-mlp 256, 打破线性可分性上限) - ② 逐步 CE 回收：R4 的 CE 只算在混合表示上，后几步表示从未被单独要求
"可解码" → OOD 摊薄 mix，CE 0.376→0.611。
--step-ce-weight 1.0对每个 思考步单独算 CE 按 w 加权，配 LoRA 低 LR (5e-5) 让思考块适配多步执行 - ③ 答案级 RL：REINFORCE with baseline，
r = 1[答案对] − λ·(n−1)/(K−1)，ponder_forced_steps=n强制恰好 n 步 rollout，只更新 halting 头
复现路线
# 0) 难度标签: K=1 探针 (实测难度是正确标签源)
python -m pondernet.probe_difficulty --model $MODEL --adapter <LORA> \
--data train.jsonl --out probe_results
# 1) R4: 步数硬监督 (head-only) → configs/round4_repro.sh
# 2) R5-A: 加权/目标/容量消融 → configs/round5_A_ablation.sh
# 3) R5-B: 逐步 CE 回收 (思考块适配) → configs/round5_B_step_ce.sh
# 4) R5-C: 答案级 RL → configs/round5_C_rl.sh
工程备忘（T4/CPU 实测踩坑）
- T4 (sm75) 无原生 bf16，emulated bf16 慢 138× → 一律 fp16 + GradScaler
pip uninstall torchao（peft 0.20 dispatch 崩溃）- halting 头恒 fp32（fp16 下 KL 的 log(0)·0 = NaN 事故）
- 共享 T4 避让：显存 <9GB 先等待；长任务 nohup + 幂等断点
- pkill -f 会自匹配杀死自己的 shell → 只按显式 PID 杀
引用
若本仓对您有帮助，请引用：
@misc{samaidev-pnet-2b,
title = {PonderNet Adaptive Computation on MiniCPM5-2B: Semantic-Level
Hard-Easy Step Differentiation via Hard Distribution Supervision},
author = {tchbcb},
year = {2026},
url = {https://huggingface.co/tchbcb/samaidev-pnet-2b}
}