GSG-LLLM — 非自回归 latent 序列 VAE 对话模型 (中文)

把 GSG (Geometric Subdivision Generation) 的"锚点骨架"思想 + 多位置 latent 序列 VAE 应用于中文对话:用一条 latent 序列 z₁..z_L(逐 token 位置)而非单个向量编码上下文,decoder 对整条 z 序列做 cross-attention 后自回归生成回答——彻底消除单 z 瓶颈,使长对话成为可能。

词表与分词复用 frsmash3.7 的 OpenASHVoc(jieba + agent 组合编码,词表 23004)。

核心架构

                ┌─────────── MultiZEncoder ───────────┐
  对话历史 ─token─► Transformer ─► 逐位置 (μᵢ, logvarᵢ) ─► z₁..z_L  (latent 序列)
                └─────────────────────────────────────┘
                                      │ cross-attention (memory = 整条 z 序列)
                ┌─────────── MultiZDecoder ───────────┐
  回答 ◄─自回归─ Transformer ◄─────────────────────────┘

多 z:每个 token 位置一个 latent,信息无损(预训练重建 100%)。
KL=0 确定性编码:推理/训练用 z=μ,无采样噪声——对话任务下 KL 正则会抹平 prompt 细节,故关闭。
两阶段:预训练(文本重建)→ SFT(对话),与自回归 LLM 范式一致。

仓库结构

GSG-LLLM/
├── code/
│   ├── open_ash_voc.py         # OpenASHVoc 分词器
│   ├── exp_minimind_zh.py      # 基础: VAE + anchor + GSG vs AR 验证 (MAX_LEN=48)
│   ├── exp_sft_zh.py           # 长对话单-z SFT (MAX_LEN=256)
│   ├── exp_multiz_zh.py        # 多-z VAE 架构
│   ├── exp_multiz_kl0.py       # 多-z + KL=0 预训练/SFT (产出 vae_mz0_pre.pt)
│   ├── exp_multiz_big.py       # 多-z 大数据 SFT (产出 vae_mz0_big.pt, 最终对话模型)
│   └── infer.py                # 推理/对话入口
├── vocab/
│   └── open_ash_voc_agent.json # 词表 (LFS, 288MB)
├── models/
│   ├── vae_mz0_pre.pt          # 预训练编码器 (KL=0, 100% 重建)
│   └── vae_mz0_big.pt          # 大数据 SFT 对话模型 (推荐推理用)
└── README.md

环境

Python ≥ 3.10,CUDA GPU(推荐 24GB,如 RTX 4090)
依赖:torch≥2.6(+cu124)、jieba、scikit-learn、scipy、numpy、tqdm

# 用 venv
python -m venv .venv && .venv\Scripts\activate
pip install torch --index-url https://download.pytorch.org/whl/cu124
pip install jieba scikit-learn scipy numpy tqdm

# 或用 uv (仓库根目录)
uv sync

Windows 运行需 PYTHONUTF8=1(脚本已内置)。

数据准备(数据集不随仓库分发)

把两份 minimind 风格的 jsonl 放到 data/(或用环境变量 GSG_DATA_DIR 指定):

文件	格式	用途
`pretrain_filtered.jsonl`	`{"text": "...", ...}` 每行一篇	预训练(重建)
`sft_filtered.jsonl`	`{"conversations": [{"role":"user","content":"..."},{"role":"assistant","content":"..."}]}`	SFT(对话)

# 数据来源示例 (三问过滤后的 minimind):
#   https://www.modelscope.cn/datasets/gongjy/minimind_dataset
mkdir data
# 把 pretrain_filtered.jsonl 与 sft_filtered.jsonl 放进 data/

训练步骤

进入 code/ 目录,按顺序执行(每个脚本支持断点续训——已存在对应权重则跳过):

cd code

# ① 多-z 预训练 (KL=0, ~14min on 4090) -> models/vae_mz0_pre.pt
python exp_multiz_kl0.py            # 同时跑一版小数据 SFT (vae_mz0_sft.pt)

# ② 大数据 SFT (~50min) -> models/vae_mz0_big.pt (最终对话模型)
python exp_multiz_big.py

参数可在脚本顶部调整:SFT_N(对话对数)、SFT_EPOCHS、BATCH、MAX_LEN 等。 exp_multiz_big.py 默认用 15 万对话对、8 epoch、MAX_LEN=256。

其它实验脚本(可选)

python exp_minimind_zh.py     # MAX_LEN=48 的 anchor + GSG vs 自回归对照验证
python exp_sft_zh.py          # 单-z 长对话 SFT (对照, 会"答非所问")
python exp_multiz_zh.py       # 多-z 带 KL 的版本 (对照)

推理 / 对话

cd code

# 内置示例问题
python infer.py

# 交互式多轮对话
python infer.py -i

# 指定权重 + 采样温度
python infer.py -m ../models/vae_mz0_big.pt --temp 0.7

结果

版本	改动	对话质量
单 z	prompt 压成 1 个 256 维向量	答非所问(瓶颈)
多 z + KL	逐位置 z, 但保留 KL	重建 33%, 仍跑题(KL 抹平 prompt)
多 z + KL=0	确定性编码	重建 100%, 但小数据(2万对)默认自我介绍
多 z + KL=0 + 大数据	15 万对话对	切题(北京→北京、AI→AI、编程→编程), 结构化回答

最终模型 vae_mz0_big.pt(25.4M 参数,15 万对话对)能识别问题主题并给出切题、结构化的回答。仍存在的局限(小模型常态):回答有重复循环、事实不够精确——继续加大数据/参数可改善。

关于 agent token

OpenASHVoc 的 agent 编码会把高频组合词编为 ts##/te##/rs##/rc##/re## 子 token 序列。解码时 OpenASHVoc.decode 会把它们重组回原词(本仓库 infer.py 已正确处理), 推理输出是干净中文,不会显示这些子 token。

权重下载

# ModelScope SDK
pip install modelscope
python -c "from modelscope import snapshot_download; print(snapshot_download('dfytensor/GSG-LLLM'))"

# 或 git clone (含 LFS 大文件)
git lfs install
git clone https://www.modelscope.cn/dfytensor/GSG-LLLM.git

引用

分词器与三问过滤数据思路来自 frsmash3.7。 GSG 锚点思想源于图像/网格合成中的几何细分生成。