研究：分解式输出层让小模型在同等算力下用上 131k 词表

BLCNYY · x · 2026-11-16

一项独立研究（Jakub Michałek，代码与日志已开源）比较了小语言模型（6 层、d=384，约 62–108M 参数，FineWeb-Edu 训练）在同等训练 FLOPs 下的三种输出层方案。

核心发现

在 d=384 + 131k 词表设置下，普通 flat softmax 头要吃掉 79.5% 的训练算力。
两级头（8,192 高频词独立 logits + 4,096 个 k-means 聚类）和 adaptive softmax 两种分解式输出层都优于 flat 头：131k 词表相对最佳 50k 模型可做到 -0.0095 / -0.0092 BPB，而 flat 头下 131k 反而差 +0.0495。
论文的反直觉发现：adaptive softmax 常规的窄尾设计反而有害——把尾部分数按全宽度计算后，BPB 改善 -0.0107 / -0.0093，并与两级头打平（-0.0010 / +0.0004，预先约定的 ±0.003 平局规则）。
两级头在 GPU 上训练吞吐达 flat 头的 1.98×（50k 词表）到 4.29×（131k 词表）。

注意限制：样本仅 1–2 个种子，131k BPE 在训练语料上训练而 50k 用 GPT-2 tokenizer，词表大小与领域同时变化，结论需谨慎外推。

「研究」频道最新