Next Concept Prediction in Discrete Latent Space Leads to Stronger Language Models

1 · 动机:为什么要在 token 之上再造一层"概念"

1.1 历史脉络:NTP 已经吃饱了它能吃的红利

从 GPT-2 / GPT-3 到 Llama-3 / Qwen3,LLM 走的是同一条路: 纯 Next Token Prediction,词表是 BPE 子词,loss 是 cross-entropy。这条路把 scaling law 推到了万亿参数,但作者引用 LeCun 2022 的论点:低层级 (low-level) 的预测目标一旦逼近 token 信息熵的下限,继续 scale model 只能改善常数项,不能改变曲率。

而别的领域早就跨越了"原始信号"这一关:

• 视觉:VAR / VQ-GAN / DALL·E 都把 pixel-level prediction 换成了 latent token prediction;JEPA 直接预测 representation,不预测 input。
• RL/Robotics:Decision Transformer / Diffusion Policy 把 trajectory 当 sequence 建模,而不是 step-wise mapping。
• 多模态:LCM (Large Concept Model, Meta 2024) 干脆把句子 embedding 当作预测单位。
• 音频:wav2vec / HuBERT 把 raw waveform 换成 discrete hidden units。

语言侧也有零碎的尝试 (Hourglass / ContextLM / BLT / H-Net / MegaByte / MTP),但要么是结构上压缩 (改 tokenizer / 改 chunk),要么仍然死守 next-token loss。"专门给语言模型设计一个高层级的自监督预训练目标"还是空白。这就是 NCP 想填的位置。

1.2 别的方案为什么不够 — 一张比较表

1.3 为什么这事不平凡:三个工程/理论上的暗坑

把 "high-level prediction" 这个口号落到 LLM pretraining 上,有三件事拦路:

(a) 概念词表从哪来? 视觉里 VQ-GAN 可以拿一个固定 codebook (因为图像有 reconstruction 监督),语言里没有 reconstruction loss;如果 codebook 仅靠 NCP loss 训,容易 codebook collapse (作者在附录 E.3 报告过:vanilla VQ + 浅插入位置就崩)。需要一个能 jointly 优化、且抗崩塌的量化方案。

(b) 离散监督怎么传梯度? 直接 hard sample 一个 codebook entry 不可导。作者用预测 logits 对 codebook 做 weighted sum,然后 MSE 回归到真值 concept hidden,绕开了 cross-entropy + Gumbel-softmax 的复杂度。

(c) 信息泄漏。 概念是 k 个 token 的池化,即位置 t 的 concept 包含了未来 k−1 个 token 的信息。如果直接把 ground-truth concept 加回 token hidden 去预测下一 token,模型就偷看了未来。NCP 的精髓是:Decoder 必须用预测出来的 concept (而非 ground-truth),并且 broadcast 时再后移 k−1 位。这是论文最容易被忽视、但最容易实现错的工程点。

2 · 背景速查

2.1 关键术语

2.2 VQ-VAE 流程速复习

1. Encoder 输出连续向量 z ∈ ℝ^d
2. 查询 codebook {e_1, ..., e_K},取最近邻 e_q = argmin_k ||z − e_k||²
3. Decoder 用 e_q 重建输入
4. Loss: reconstruction + ||sg(z) − e_q||² (codebook loss) + β||z − sg(e_q)||² (commitment loss)

本文复用了 (sg, commitment) 这套结构,但去掉了 reconstruction loss——concept 不需要重建任何东西,它只需要被下游的 Decoder "用"上即可。

2.3 Product Quantization 直觉

朴素 VQ 想表达 N^S 种概念,需要一张 size = N^S 的 codebook,既爆显存也爆参数。Product Quantization (Jegou 2010 经典 ANN 加速方法) 把 d 维向量沿特征维切成 S 段,每段独立查 N 项的小表,总组合数 N^S 但参数量只有 S × N × (d/S) = N × d。论文里 S = head 数 (~12),N = 64,概念词表实际大小约 64^12 ≈ 7.2 × 10^22,远大于 BPE 词表组合 V^k = 50000^4 ≈ 6.25 × 10^18。

3 · ConceptLM 架构总览

三段式: Token-level Encoder → Concept-level Module → Token-level Decoder。Encoder 和 Decoder 共享原 backbone (Pythia / GPT-2 / Llama) 的层,在某个深度插入两层 concept transformer。

3.1 NTP vs NCP 一图对比

4 · VQ 概念词表的构造与 Product Quantization

每个 concept c ∈ ℝ^d (d 是模型 hidden size) 沿特征维切成 S 段:c = [c¹, c², ..., c^S],每段维度 d/S。对第 s 段,有自己的 codebook E^s = {e^s_1, ..., e^s_N}。量化方式:

4.1 为什么要 product 而不是 monolithic

如果单 codebook,要表达 N^S 种概念,codebook size 必须 = N^S,参数量 ~ N^S × d,完全无法承受。Product 之后:

• 参数量 = S × N × (d/S) = N × d (与 S 无关!)
• 表达力 = N^S (与 S 指数相关)
• 查找代价 = S 次小 KNN,可在 GPU 上并行

论文里 Pythia 的 d=768, head=12,所以 S=12,N=64,概念词表 = 64^12 ≈ 7.2 × 10^22。对比 BPE: V^k = 50000^4 ≈ 6.25 × 10^18,概念词表大了 4 个数量级。这意味着概念粒度比 4-gram token 更细,有充足的"语义槽位"。

4.2 防 codebook 坍塌:SimVQ trick

vanilla VQ 容易让大部分 input 都 nearest 到少数几个 e。作者沿用 SimVQ (Zhu 2025): 把 codebook entry 通过两层 MLP + ReLU变换之后再用,即真实参与 KNN 的不是 e^s_n 本身而是 MLP(e^s_n)。原理:非线性变换让 codebook 的有效流形更弯曲,不易塌到一个子空间。

实证 (附录 E.3): 即便用 S=12, N=1024 (1.2 万个 entry) 的大词表,codebook utilization 仍达近 100%。如果不用 ReLU MLP,而且 concept 层插入位置太浅,会观察到 collapse。

4.3 Worked example: 一个 4-token span 的量化

设 backbone 是 Pythia-160M, d=768, S=12, N=64, k=4。输入序列 "The integral of x squared" 经 BPE 分成 6 个 token,我们看前 4 个 token 组成的 concept c0:

x0..x3 = ["The", "Ġintegral", "Ġof", "Ġx"]
h0..h3 ∈ ℝ^768   (Encoder 第 1 层之后的 hidden)
c0 = mean(h0, h1, h2, h3) ∈ ℝ^768

切 12 段:c0¹ ∈ ℝ^64, c0² ∈ ℝ^64, ..., c0^{12} ∈ ℝ^64

对每段并行查 codebook (size 64):
  c0¹ → argmin_n ||c0¹ − MLP(e¹_n)||²  =  e¹_{37}
  c0² → e²_{12}
  ...
  c0^{12} → e^{12}_{5}

离散概念 ID = (37, 12, ..., 5)  ← 一个 12-tuple,概念词表里的一个"词"
量化向量 d0 = concat(e¹_{37}, e²_{12}, ..., e^{12}_{5}) ∈ ℝ^768

这个 12-tuple 就是 c0 在概念词表里的"BPE id"。下一个 chunk c1 (x4..x7 = "squared dx equals one third") 同样得到一个 12-tuple,这个 12-tuple 才是 NCP loss 想"猜对"的目标。

5 · Concept 预测头:weighted sum 而非 hard sample

预测下一个 concept 时,Concept Module 输出 hidden h_c,再过 S 个独立的 prediction head,每个 head 是一个 linear (d_concept → N),输出 N 维 logits → softmax → 权重 w^s ∈ ℝ^N。

关键设计:不做 argmax / Gumbel-softmax 取出某一个 codebook entry,而是直接做加权求和:

然后用 MSE(ĉ, c) 监督——而不是用交叉熵预测 codebook id。这是论文里反复出现但不显眼的设计选择。

5.1 为什么不用交叉熵 + hard target

三个考虑:

1. 可微分性: weighted sum 处处可导,梯度直接流回 head。argmax / Gumbel 需要 reparameterization,引入额外噪声。
2. codebook 共同优化: 因为 ĉ 是 codebook entry 的线性组合,梯度也流回 e^s_n,从而 codebook 能被 LNCP 间接训(虽然主要还是 LVQ 在训)。
3. 避免 over-commitment: 如果只能选一个 entry,模型在概念边界处 (语义模糊段) 会震荡;weighted sum 容许"34% 是 e_3 + 41% 是 e_7 + ..."的软表达。

6 · 三个 loss:NTP + NCP + VQ 各自做什么

6.1 三者的分工

总 loss: L = L_NTP + L_NCP + L_VQ (默认权重 1:1:1,β 是 commitment 项内部权重,文中没显式给具体值,但是 VQ-VAE 经典设置 0.25)。

6.2 Ablation 关键数字 (Pythia-160M, 30B tokens, Table 4)

读法: 单加 VQ 或单加 NCP 都不稳赢,有时 PPL 反而更差;只有三者齐全才同时拿到 Acc + PPL 双赢。MTP 在这个尺度 (160M) 反而比 baseline 还差——印证了 Gloeckle 2024 自己说的"3B 以下 MTP 没用"。

6.3 为什么 NTP 不能去掉

论文给的解释更微妙: NCP 学的是稀疏的、每 k 步一次的语义信号,Decoder 仍需要 token-级稠密监督才能维持基本语言能力 (语法 / 局部连贯)。如果只用 NCP 训,Decoder 拿到的 broadcast 概念向量只是"语义提示",但没有任何信号告诉它"该把这个语义铺成哪个具体 token"。NTP 是 dense supervision,负责"概念 → token 表面形式"的映射。

7 · 防信息泄漏:k−1 位移 trick

这是论文里最容易踩坑的一处。预测下一个 concept c_{t/k} 需要看完整的前 k 个 token 的 hidden h_{(t/k-1)·k .. t-1}——等价于看到了过去一个完整的 chunk。但 broadcast 之后,decoder 在位置 t 拿到的 ĉ 是对应哪个 chunk 呢?

8 · 公式拆解 + 一个 token 的全程追踪

8.1 总目标函数

8.2 物理直觉

• L_VQ 的两项: 第一项 sg(c) − d 拉 codebook 去贴近输入分布 (codebook 学习);第二项 c − sg(d) 给 commitment loss,让 Encoder 也愿意把表征"靠拢"到现有的 codebook 上,防止 c 漂走。sg 把两条梯度路径解耦,这也是为什么 LVQ "不影响主表征"。
• L_NCP 选 MSE 而非 CE: 见 §5.1。MSE 在 ℝ^d 上是稠密信号;CE 在 N^S 离散空间太稀疏。
• L_NTP 仍是 CE: 在 token 词表 V (50k) 上 CE 已足够稠密,且与"概念已经被提示进 hidden"形成互补。

8.3 一个 token 的全程追踪

设 Pythia-160M, T=2048, d=768, k=4, S=12, N=64。输入是数学语句:

"The integral of x squared dx equals one third x cubed plus C"
tokens (12): [The, Ġintegral, Ġof, Ġx, Ġsquared, Ġdx, Ġequals, Ġone, Ġthird, Ġx, Ġcubed, ...]

追踪 t=4 这个 token Ġsquared 的预测过程:

位置→概念 index 的映射有点 tricky,公式里写 ⌊t/(k−1)⌋ 是 paper 自己的写法,实践中需要小心 off-by-one。

8.4 一个数值 sensitivity 表 (chunk size k)

9 · 实验关键结果与 scaling 行为

9.1 Scaling 主图 (论文 Fig 1) 的核心三个数

• 同 perplexity 减 37% 参数: GPT-2 1.5B 的 PPL ≈ ConceptLM 1.0B 的 PPL
• 同 perplexity 减 24% 训练 token: 同尺寸下,只跑 76% 数据就 match 全量 token-level baseline
• 同 perplexity 减 19% FLOPs: 综合参数 + 数据的训练成本节省

9.2 Pythia 全尺寸表 (300B token, 选自 Table 1, 平均 PPL ↓)

注意 160M 这一档跳点最大 (+33.5 PPL),作者归因于 S=head 数 + N=64 的固定配置在这个尺寸最贴合。这暗示 NCP 的 hyperparams 不是 size-agnostic 的,大模型可能需要 N 更大 / S 更大。

9.3 Continual Pretraining: 在 Llama-3.1-8B 上加 NCP (Table 3, 9.6B token)

Continual pretrain 的提升只有 +0.4 (vs PM) / +0.1 (vs Trained),不算惊艳——但作者强调:这是从已有 NTP 模型扩出来,新加的 concept 层和原 backbone 对齐难,而且只跑了 9.6B token (相对 Llama-3.1 总训练量微不足道)。LVQ 还在持续下降,说明没收敛。

9.4 长文实验 (Fig 4)

训练长度从 512 → 2048,ConceptLM 比 PM 基线的 loss gap 单调扩大。原因:concept 层运行在 T/k 长度上,attention 二次项缩 16 倍,把"省下来的算力"投到了 long-range dependency 上。这是 NCP 的副产品。

9.5 计算复杂度 (Table 6, Llama-3.1-8B 实测)

实测: PM-8B 训练 80h,ConceptLM-8B 71h (↓11.3%);推理 99min vs 86min (↓13.1%)。理论 FLOPs 只省 4.69%,但实测省更多——作者归因于IO/KV cache: concept 层处理的序列短了 4 倍,KV cache 直接缩 4 倍,缓解了 memory bandwidth bottleneck。

10 · 与同类工作对比

定位总结: ConceptLM 占据的"格子"是 [语言域 + pretraining objective + discrete latent + 与 NTP 联合训练],四个维度中前人的工作至少缺一个。最接近的是 ContextLM (同实验室前作),区别就在 显式 VQ + NCP loss vs implicit concept 层——这个差距在 Pythia 实验上稳定值 +1.7 ~ +2.5 Acc。

11 · 局限 / 个人 take / 待验证问题

11.1 局限

• 规模不够大。 主实验最大 1.5B,continual 8B 只跑 9.6B token。MTP 在 3B 以下没收益是已知现象,NCP 在 8B+ 时和 NTP 的 gap 是否扩大或收窄,本文没给。
• k=4 是固定的。 没有动态 chunk (BLT 风格),也没考虑数学/代码这种 token 边界刚性的领域 chunk 应不应该跨语法边界。
• S, N 也是固定的。 论文承认在不同 size 上不是 size-optimal (160M 跳得最猛,1.5B 几乎没跳)。
• 生成效率 unclear。 推理时仍要先 forward Encoder 算 c,再过 concept 层预测 ĉ,再 broadcast,再过 decoder——比纯 NTP 多了概念层 forward。Table 6 报告的"推理快"是因为 concept 层替换了 2 个 token 层,不是叠加。如果想要"加而不减",收益要重新算。
• NCP 的 MSE loss 没有概率解释。 不像 CE 可以解读为 KL,MSE 在概念词表上的几何意义不直观,调参 (vs L_NTP) 缺少理论指南。
• 没有 SFT / RLHF 阶段实验。 概念层在 instruction-following 时是否依然提供增益,完全空白。

11.2 个人 take

• 这篇的"立场"我很喜欢: 把 high-level prediction 这个口号真正落到了 pretraining objective 上,而不是 inference 时的潜推理 trick。这是 Coconut / Token Assorted 这类工作差的那一步。
• 但 8B 上只 +0.4 的提升非常 marginal。如果未来在 70B 上不能继续放大,这个故事就只是"小模型 trick"。
• VQ + product quantization 已经在视觉 (VAR / MAGVIT) 验证可以 scale,语言侧的 codebook collapse 风险被作者用 SimVQ + 插入深度 (3/4 处) 解决了,工程上看起来稳。
• 跟 DeepSeek-V3 MTP 的关系值得深思: MTP 在 token 空间组合爆炸,NCP 用 VQ 把"组合"压成离散 codebook,本质上是给 MTP 找一个紧凑的预测目标空间。这个思路可以反过来用 — 把 MTP head 的预测目标换成 concept code 而不是 token tuple。

11.3 待验证问题 (个人想做的实验)

1. 在 Llama-3 70B 上 continual pretrain NCP (而不是 8B),收益是否扩大?如果不扩大,说明 NCP 是"小模型补丁"。
2. k=4 是否对所有 domain 通用?数学代码用 k=2,叙事文用 k=8,可能各有甜点。
3. 把 NCP 的 MSE loss 换成对 codebook id 的 cross-entropy,效果 vs 现在的 weighted-sum + MSE 谁好?这是设计选择的核心 ablation。
4. 推理时是否可以用 concept ID 做 speculative decoding 风格的"概念前瞻"?比如先采样概念,再投机性地展开 k 个 token,target 模型只验证概念边界。
5. Concept 层插入位置(论文报告 3/4 处最好) 在 8B+ 上是否仍然在 3/4 处? 还是和 backbone 深度相关?
6. Codebook (在多语言 / 多 domain 训练后) 学到的是什么? 一个 concept ID 对应的所有 4-token span 在语义上是否聚类? 这是 interpretability 的核心问题,论文完全没碰。

记忆点

精读笔记 v1 · 2026-05-07 · 配套论文 PDF: /data/szhang967/papers/paper-notes/models/NCP_2602.08984.pdf