← Back to list
KITE

KITE: KV-Invariant Transformer Expansion for Efficient Agentic LLM Scaling

LLM StepFun
Abstract 7 │ Reading 6 │ Rating —
2026-09-23
Zhiheng Hu, Yixun Wei, Jian Zhou, Yizhuang Zhou, Ji Li, Xing Chen, Yang Li, Bojun Wang, Yibo Zhu, Xiangyu Zhang, Daxin Jiang
StepFun
StepFun 的 KITE 在训练中途把小模型扩成两塔:源模型保留为唯一产 KV 的 Prefiller,复制出只读逐层 KV 的 Decoder 联合续训,使 bulk prefill 保持源模型尺寸;SST 实例(33.8B→67B MoE)在等理论训练 FLOPs 下训练 loss 1.5900 优于 47B/63B 从头训练(1.6006/1.5921)、7 项下游全领先而 held-out NLL 与 63B 持平,75:25 P:D 下参数量代理推理代价低 6.7%/31.6%,但无实测 serving、无生长/KV 复用基线、无消融。
评分原因
摘要评分:通用 LLM 扩参新范式:从小模型 upcycle 扩大,并把新增参数放在不影响 KV 的塔中,使 prefill/KV 成本只依赖小模型,67B MoE 规模验证、对长序列推荐推理成本有迁移价值;但摘要无可核实的工业部署/线上证据,推理节省为估算值,按“有创新但工业验证不足”给 7。
精读评分:KV-invariant 生长范式清晰可组合、等 FLOPs 核算严谨(三者均≈5.168e21);但单配置单种子、仅两个同家族从头训练基线,相对 63B 的 margin 仅训练 loss 0.0021 / held-out NLL 0.0003 且曲线收敛,无生长基线(如 36 层堆叠)、无 YOCO 类 KV 复用基线、无任何消融,KV-invariance 的质量代价与收益来源无法归因;推理节省仅为不含注意力/KV 访存的激活参数代理,无实测,相对 47B 优势依赖 prefill 权重≥65.5%;StepFun 署名但无部署证据。
transformer moe parameter-scaling inference-serving industrial

KITE: KV-Invariant Transformer Expansion for Efficient Agentic LLM Scaling

作者:Zhiheng Hu, Yixun Wei, Jian Zhou, Yizhuang Zhou, Ji Li, Xing Chen, Yang Li, Bojun Wang, Yibo Zhu, Xiangyu Zhang, Daxin Jiang 机构:StepFun arXiv:2609.27294(2026-09-23,cs.LG,13 页)

Figure 1: KITE: training route and observed trade-offs. Both panels use the same 47B transformer trained from scratch as the reference; lower is inward. Active counts exclude embedding/head. The inference proxy is 0.75p + 0.25d, where p and d are bulk-prefill and decode active-body counts normalized to the reference.

一句话概括

KITE 提出一种「扩参范式」:先训一个小模型,训练中途把它扩成两塔结构——原模型保留为只负责产出 KV 的 Prefiller,再复制出一座只读 Prefiller KV、自己不产 KV 的 Decoder,两塔继续联合训练。这样新增的参数既享受「中途生长」的训练折扣,又不进入 prompt 预填充(bulk prefill)路径。其具体实例 SST(Step Scale Transformer) 把 33.8B MoE 扩成 67B MoE,在与 47B / 63B 从头训练 MoE Transformer 理论训练 FLOPs 相同时,训练 loss 更低(1.5900 vs 1.6006 / 1.5921)。在 75:25 的 prefill:decode 权重下,基于激活参数量的解析推理代价代理比 47B 低 6.7%、比 63B 低 31.6%。作者明确声明这些都不是实测的 serving 加速。

1. 研究动机与背景

1.1 训练成本与推理成本要「同时」下降

LLM 的 scaling 性质既是福音也是诅咒:更大的模型更强,但训练和推理的计算量都随之上涨。作者把「以更低的算力代价 scale 模型」作为核心问题,并指出 MoE 之所以成功,一个容易被忽视的原因是它同时降低了训练和推理的计算量:在固定参数量下省算力,或者说在固定算力下长容量。

作者指出,只在一侧省算力的方案往往被另一侧抵消:

  • 渐进式 upcycling / 模型生长(Net2Net、SPARKLING 等):先训小模型,训练中途长成目标模型,大量 token 在更便宜的计算图上跑完,省了训练。但业界普遍认为 upcycled 模型打不过同尺寸从头训练的模型,于是训练省下的算力在推理时被还回去:serve 的是完整的目标模型,拿到的却不是完整目标模型的质量。
  • KV 复用(YOCO 家族、DeepSeek-V4.1-Flash、WeChat 的 KV-Mirror 等):深层复用浅层算出的 K/V,模型变大时 prefill 仍然便宜,省了推理。但删掉产 KV 的层会损失质量,要补回来就得更大的模型或更长的训练,等价于训练成本损失。

因此作者主张把两种成本联合起来看,追求一个让「模型质量、训练成本、推理成本」三者都朝有利方向移动的 scaling 范式。

1.2 Agentic 负载是 prefill 主导的

KITE 押注的负载特征是:agentic 场景中模型被反复调用,每次都要处理累积的上下文和新的工具输出(引用 SWE-agent、ReAct)。即使有 prompt caching,仍然有大量未命中缓存的输入。作者用 OpenRouter 公开流量佐证(Table 1):六个模型中未缓存输入的 token 量都超过输出,在「未缓存输入 + 输出」计费中也占多数。

Table 1:OpenRouter 选定流量的输入需求(2026-08-22 至 09-20 UTC,按 09-21 基准价估算;Input 含缓存 token,Output 含推理 token;Charge share = 未缓存输入计费 /(未缓存输入 + 输出计费),不含 cache-read 计费)

模型 Input/output tokens Uncached input/output tokens Uncached-input charge share
GPT-6 Astra 106.43× 14.64× 74.6%
GPT-6 Astra Pro 63.15× 14.01× 73.8%
Claude Fable 5 63.38× 16.29× 76.7%
Claude Sonnet 5 59.88× 12.12× 70.8%
Gemini 3.7 Flash 37.79× 10.15× 67.4%
Gemini 3.8 Flash 46.14× 10.51× 67.9%

解读:未缓存输入 token 是输出的 10–16 倍。由于输出单价约为输入的 5 倍,折成计费份额后未缓存输入占 67–77%。这是 KITE 整个推理收益论证的「权重」来源,后面的 75:25 prefill:decode 代价比就取自这里。需要注意,这是价格份额,不是 GPU 计算份额;作者在附录 A 里也写明「这些统计只是说明输入重的推理形态,并不测量 GPU 成本」,而且这批流量没有被单独标记为 agentic 用量。

1.3 KITE 的核心思路

在 scale 时把模型划成两个区域:决定 KV 张量的区域和不影响 KV 的区域,只在后者里扩容量。作者给出两层理由:

  1. 新增容量是训练中途才出现的,而不是从第 0 个 token 开始,所以它买到的质量在训练成本上是划算的(生长的训练折扣);
  2. 扩容不影响 KV 的计算,所以推理时的 prefill 代价保持不变(KV 复用的推理折扣)。

作者把省下的算力用来适度扩大模型,就像 MoE 那样,并声称 KITE 与 MoE、hybrid attention 正交、互不干扰。作者也坦承:SST 很可能不是最优的 KITE 设计,只是足以展示这一范式的正面效果。

2. 核心方法:KITE 设计原则与 SST 架构

2.1 设计原则:Prefiller / Decoder 分解

设因果模型由 Prefiller $P$ 和 Decoder $D$ 组成:

$$(M_t, z_t) = P(x_{1:t}; \theta_P), \qquad \ell_t = D(x_t, z_t, M_t; \theta_D) \tag{1}$$

其中 $M_t$ 是可复用的注意力记忆(即 KV),$z_t$ 是位置 $t$ 上可选的 Prefiller 特征,$\ell_t$ 是下一个 token 的 logits。先完成一个小模型训练阶段,然后在中间 checkpoint 上新增或扩大 $D$,同时保留 Prefiller 的计算图和记忆接口,两部分继续训练。

不变量是结构性的:权重和 KV 值都不冻结。约束只有两条:

  1. Decoder 不得产生未来位置需要的记忆,也就是 Decoder 不产 KV;
  2. 为了在生成时省掉历史 Decoder 计算,给定 Prefiller 输出后,Decoder 保留的位置必须与被省略的 Decoder 位置相互独立。

有了这两条,预测容量可以增长,而 bulk prefill($P$ 执行的全 prompt KV 构建)不增加。仍然必须计算的是:首个输出 token 在 prompt 边界位置的 readout,以及 decode 阶段 Decoder 的全部工作。

2.2 Step Scale Transformer(SST)

SST 用两座 Transformer 塔实例化上述原则:Prefiller(上标 $P$)与 Decoder(上标 $D$),实验中各 18 层。Prefiller 先执行并产出逐层 KV;Decoder 每个 block 计算自己的 query,并从对应层的 Prefiller KV 读取:

$$q^D_{l,t} = Q^D_l\big(h^D_{l-1,t}\big), \qquad a^D_{l,t} = \mathrm{Attn}\Big(q^D_{l,t},\ K^P_l[I_l(t)],\ V^P_l[I_l(t)]\Big) \tag{2}$$

  • $Q^D_l$ 包含 query 投影和归一化(Decoder 有自己的 Q 投影和 Q norm,没有 K/V 投影);
  • $I_l(t)$ 是原因果 full / sliding 注意力掩码允许的位置集合,即 Decoder 第 $l$ 层沿用 Prefiller 第 $l$ 层的注意力模式(full 或 sliding window);
  • 两塔保留源模型的 position ID;
  • Decoder 有独立的 block 参数,残差和 FFN/MoE 运算都是 token-local 的。

Figure 2 (Stage 1): 源模型 / Prefiller 单塔训练——Token embedding → L 个 Transformer block → Output norm + LM head。

Figure 2: KITE instantiated with SST. Stage 1 trains the Prefiller; expansion adds a Decoder initialized from the source weights, and Stage 2 continues training both towers jointly. Within Stage 2, the Prefiller executes in full before the Decoder. Each Decoder block uses aligned Prefiller KV. A token-local entry bridge combines the separately normalized embedding and final Prefiller state. The final Decoder state is normalized and passed to the shared output norm and LM head for prediction.

为什么 Decoder 的历史位置可以省略:Decoder 在位置 $t$ 的计算只依赖 (a) 同位置的 token-local 运算和 (b) Prefiller 在 $\le t$ 位置产出的 KV。Decoder 自己的历史 hidden state 不会被未来位置读到,所以「给定 Prefiller 输出,每个 Decoder 位置都独立于更早的 Decoder 状态」。梯度则通过两条路径流回 Prefiller:entry bridge 和被复用的 KV。

2.3 训练与推理的执行方式

  • 训练:扩展后两塔联合训练。Prefiller 处理全序列,Decoder 计算每一个被监督的位置,梯度穿过复用的 KV 流入 Prefiller。注意:训练时 Decoder 要跑满全序列,所以训练阶段的每 token 前向 FLOPs 并不省,SST 甚至比 63B 还高(见 §3.3)。
  • Prefill:Prefiller 处理完整 prompt $x_{1:n}$,产出逐层 KV cache。Decoder 只需计算最后一个 prompt 位置 $n$(attend 到 Prefiller 的 KV),用于预测第一个输出 token。更早的 Decoder 位置都可以省略,因为 Decoder 的其他运算都是 token-local 的。
  • Decode:每个新 token 要穿过两座塔:Prefiller 扩展 KV cache,Decoder 读取这些 KV 并给出下一个 token 的预测。

由此可见 KITE 的推理账是不对称的:bulk prefill 只付源模型尺寸的钱,decode 要付两座塔的钱。

2.4 Decoder 的输入与输出

记共享 embedding/stem 为 $e_t$,Prefiller 最终状态为 $h^P_{L,t}$。SST 的 entry bridge 为

$$h^D_{0,t} = \mathrm{RMS}_\epsilon(e_t) + \mathrm{RMS}_\epsilon\big(h^P_{L,t}\big), \qquad \mathrm{RMS}_\epsilon(v) = \frac{v}{\sqrt{\mathrm{mean}(v^2) + \epsilon}} \tag{3}$$

Decoder 的最终状态产出 logits:

$$\ell_t = W_{\text{out}}\, \mathrm{Norm}_{\text{out}}\big(h^D_{L,t}\big) \tag{4}$$

  • entry bridge 中的 RMS 运算没有可训练参数,$\epsilon = 10^{-5}$,在 FP32 中累加并返回输入 dtype;
  • $\mathrm{Norm}_{\text{out}}$ 是完整的输出归一化:先对 Decoder 最终状态做一次无仿射的 RMS,再接源模型原有的可学习 output norm(附录 D);
  • 两塔共享一套 embedding、output norm 和 output head;embedding 与 head 权重不绑定(untied)。

设计动机:式 (3) 把「原始 token 身份」和「Prefiller 已经算好的上下文表示」各自归一化后相加,保证两路信号在进入 Decoder 时量级相当,避免任何一路主导。式 (4) 让预测完全由 Decoder 给出,Prefiller 的最终状态只通过 entry bridge 间接参与预测,其 LM-head 通路在扩展后被弃用。

3. 关键技术细节:模型构建与训练

3.1 从源模型到扩展模型

SST 把一个 33.819B 参数的源模型扩成 66.959B,源塔保留为 Prefiller(Table 2)。训练路线包含 167.98B 源阶段 token 和 222.55B 续训 token(Figure 3),两个阶段都计入训练预算。

两个对比基线 47B classic 和 63B classic 是同一模型家族 scaling ladder 中 33.8B 源模型之上的下两档,通过加深加宽得到,保留相同的 MoE 设计和注意力模式。47B 作为中间尺度参考,63B 作为与 SST 67B 总参数量接近的参考。

Table 2:模型配置

Setting Source SST 47B classic 63B classic
Layers 18 18 + 18 20 22
Hidden width 2304 2304 2560 2816
MoE layers 16 16 + 16 18 20
Routed / selected experts 512 / 8 512 / 8 512 / 8 512 / 8
Attention pattern SSSF SSSF SSSF SSSF
Total parameters (B) 33.819 66.959 46.727 62.691
Decode active body parameters (B/token) 1.120 2.155 1.477 2.016

注:每座塔都以两层稠密 FFN 开头。SSSF 表示「三层 sliding + 一层 full attention」循环,18 / 22 层时余下的两层为 SF。计数描述的是有效架构。

按这个模式推算(笔者推算,非原文):18 层的 full-attention 层位于第 4/8/12/16/18 层,共 5 层;20 层(47B)也是 5 层;22 层(63B)是 6 层。这一点在后文讨论长上下文 KV 流量时很关键。

Figure 3(训练路线,矢量流程图未提取为位图):Training route to the evaluated SST checkpoint: 13,350 source steps followed by 17,687 continuation steps。内容为:源训练 167.98B token → 在 167.98B 处转换 checkpoint(把权重复制进两座塔)→ Prefiller/Decoder 联合训练 222.55B token,至 390.54B 为最终评估模型,共 31,037 次更新(13,350 源步 + 17,687 续训步)。

3.2 扩展与续训

  • 扩展方式:把源模型的每一层复制进 Prefiller 和 Decoder 的对应层,也就是 Decoder 的初始权重与源模型逐层相同(对齐初始化)。两塔共享 embedding、output norm 和 output head。
  • 优化器:BF16 联合训练,沿用源模型的 Muon 与 Adam 参数组。Prefiller 和共享参数保留优化器状态,Decoder 初始化全新的优化器状态。
  • 学习率:所有参数组沿用源模型的 cosine 调度,把剩余衰减拉长到覆盖续训阶段,Decoder 不单独 warmup。

值得注意的是,由于 Decoder 初始化为源模型的逐层拷贝,且 entry bridge 把 embedding 与 Prefiller 最终状态相加后再送进一个「本来只见过 embedding」的第 1 层,扩展后的模型不是函数保持的。Figure 4 中转换点处 loss 从约 1.745 跳升到约 1.90 就是这个原因,随后快速回落。论文 Related work 也明确说「对齐初始化改变了预测计算图,随后依靠联合适应」。

3.3 累计训练算力

对源阶段与续训阶段的 token 数 $D_s, D_c$ 和每 token 前向 FLOPs $f_s, f_c$:

$$C_{\text{train}} = 3\,(D_s f_s + D_c f_c) \tag{5}$$

其中因子 3 近似前向加反向。两个训练阶段的完整成本都计入。以最终评估的 47B classic checkpoint 定义 $B_{\text{ref}} \approx 5.16836 \times 10^{21}$ FLOPs,SST 与 63B 的最终点都约用了 100% 的预算。

每 token 前向成本(附录 C):47B classic 约 3.887488B FLOPs/token,63B classic 5.148513B,源模型 3.087252B,SST 5.410678B。SST 保留两座塔的注意力和 FFN/MoE 计算,但只有一套 embedding、output head 和 KV 投影。所有前向系数按序列长度 4096、sliding window 512、padded 训练词表 128,896 计算,乘加计为 2 FLOPs,包含注意力投影与矩阵乘、embedding、output head、FFN/MoE、router 和估计器的激活项;不含归一化、RoPE、残差与 bridge 运算、重计算、优化器、loss/aux loss、转换、通信和 checkpoint I/O。

笔者复核:

  • 47B:$3 \times 443.16\text{B} \times 3.887488\text{B} \approx 5.168 \times 10^{21}$;
  • 63B:$3 \times 334.62\text{B} \times 5.148513\text{B} \approx 5.168 \times 10^{21}$;
  • SST:$3 \times (167.98\text{B} \times 3.087252\text{B} + 222.55\text{B} \times 5.410678\text{B}) \approx 3 \times (5.186 + 12.041) \times 10^{20} \approx 5.168 \times 10^{21}$。

三者确实严格等 FLOPs,源阶段占 $3 \times 5.186\times10^{20} / B_{\text{ref}} \approx 0.301$,与 Figure 4 的 0.301 $B_{\text{ref}}$ 一致。这说明训练账算得很仔细。同时也可以看出,SST 续训阶段每 token 前向(5.41B)比 63B(5.15B)还贵 5%,因为训练时 Decoder 要跑满全序列。KITE 在训练侧的全部折扣都来自前 30% 算力跑在便宜的 3.09B/token 源模型上。

3.4 训练配方与续训调度器(附录 D)

  • 原生 BF16 参数,Muon + Adam 两组;Muon 用 6 步 Newton–Schulz,packed attention 按 head 切分、GLU 按投影类型切分;Adam $\epsilon = 10^{-15}$;
  • 源模型峰值 base 学习率约 $1.03302 \times 10^{-3}$;续训 base weight decay 0.15,沿用参数组乘子;
  • token batch:源 / SST / 63B 为 12,582,912,47B 为 14,680,064(47B 的 batch 不同);
  • SST 续训:128 张 NVIDIA H800,expert parallelism 8,序列长度 4,096,global batch 3,072 条序列;续训 17,687 步,接在源训练 13,350 步之后,共 31,037 步。

记已消耗续训 token 为 $u$,续训总时域 $U = 222{,}553{,}964{,}544$,源模型在转换时的归一化 cosine 相位为 $q_0$,SST 的调度为

$$q(u) = q_0 + (1 - q_0)\,\min(u/U, 1) \tag{6}$$

$$\eta(u) = \eta_{\min} + (\eta_{\text{anchor}} - \eta_{\min})\,\frac{1 + \cos(\pi q(u))}{1 + \cos(\pi q_0)} \tag{7}$$

源调度在前 12.582912B token 内 warmup,衰减到 335.9605248B token 为止;转换发生在 167.9818752B token 处,所以 $q_0$ 是在 warmup 之后的衰减区间内度量的,而不是用转换位置除以总时域。学习率常数为 $\eta_{\text{anchor}} \approx 5.96548 \times 10^{-4}$、$\eta_{\min} \approx 1.03302 \times 10^{-4}$。

含义:式 (6) 把余下的 cosine 相位 $[q_0, 1]$ 线性摊到整个续训时域上;式 (7) 按比例缩放,使 $u=0$ 时 $\eta = \eta_{\text{anchor}}$(即转换瞬间源调度的学习率),$u=U$ 时衰减到 $\eta_{\min}$,保证续训 LR 曲线在转换点连续、在终点完整退火。

(原文附录 A 的计费公式与推理代价代理放在 §5.3 介绍;本文公式按出现顺序重新编号。)

4. 实验设置

  • 对比对象:只有两个同家族、从头训练的 MoE Transformer:47B classic(20 层、宽 2560)和 63B classic(22 层、宽 2816)。此外只有作为训练曲线参照的 33.8B 源模型。
  • 控制变量:tokenizer、数据配方、训练协议和评测集完全一致;比较点为等理论累计训练 FLOPs(约 $1.0\,B_{\text{ref}}$)。
  • 质量指标:
  • 训练 LM loss 的 EMA-200(span 200,adjust=False,在绘图抽样前计算并保留最后一个观测;SST 在转换后重启 EMA);
  • 7 个下游任务:OpenBookQA(validation + test 合并)、MMLU(test)、GSM8K、MATH(full test)、HumanEval(few-shot)、MBPP(3-shot)、BBH;
  • held-out ARXIV token NLL:三个模型用同一组 1,024 条冻结的 packed-4K 样本。

  • 推理指标:基于激活 body 参数量的解析代价代理,按不同 prefill:decode 权重扫描;没有任何实测吞吐或延迟。

  • 种子/重复:每个配置单次训练,每个模型只评估一个最终 checkpoint,没有方差估计。

Table 5:详细几何(Source/SST 的维度按每座塔计,均以两层稠密 FFN 开头)

Setting Source / SST 47B classic 63B classic
Dense FFN width 4608 5120 5632
Routed / shared expert width 576 / 576 640 / 640 704 / 704
Head dimension 128 128 128
Full / sliding RoPE dimension 64 / 128 64 / 128 64 / 128
Full / sliding Q heads 24 / 40 24 / 40 32 / 48
KV heads 8 8 8

Table 6:架构参数计数(B,保留三位小数)

Quantity SST 47B classic 63B classic
Total, including embedding/head 66.959 46.727 62.691
Embedding + output head 0.594 0.660 0.725
Total body 66.365 46.068 61.966
Decode active body 2.155 1.477 2.016
Prefiller-only active body 1.120 – –

激活 body 参数不含 embedding/head,计入每个 MoE 层 Top-8 路由专家、共享专家、router 以及一个 token 所需的其余 body 权重。作者强调这是每 token 参数计数,不是实测执行成本。导出词表 128,815 行(不含训练 padding)。在 SST 的有效架构中,只有 Prefiller 含 K/V 投影和 K 归一化。

Table 7:最终 checkpoint 的训练预算

Model Training tokens (B) Relative training FLOPs
47B classic 443.16 1.000
63B classic 334.62 ≈1.000
67B SST 390.54 ≈1.000

5. 主要实验结果

5.1 训练 loss

Figure 4: Training LM loss (EMA-200) against tokens and theoretical cumulative compute. Conversion occurs at 167.98B tokens (0.301Bref).

SST 的 loss 在转换点跳升,随后在续训中持续下降。最终 checkpoint 上:

模型 训练 token 训练 EMA-200 loss vs SST
47B classic 443.16B 1.6006 +0.0106
63B classic 334.62B 1.5921 +0.0021
67B SST 390.54B 1.5900 —

分析:

  • Figure 4(b)(按算力对齐)有三个值得注意的细节。第一,转换时(0.301 $B_{\text{ref}}$)源模型的 loss 本来就略低于同算力的 47B / 63B,这符合 Chinchilla 直觉:训练早期小模型每 FLOP 更高效。SST 实际上是从一个「在该算力点上领先」的位置起跳的。第二,转换后 SST 在约 0.35–0.7 $B_{\text{ref}}$ 区间领先较明显。第三,到 1.0 $B_{\text{ref}}$ 时 SST 与 63B 的曲线已经几乎重合,差距收窄到 0.0021。从曲线走势看,差距在收敛而不是拉开;论文没有给出更大预算下的结果,也没有拟合 scaling law,所以无法判断大预算下 SST 是否仍然领先。
  • Figure 4(a)(按 token 对齐)上,同 token 数时 SST 介于 47B 和 63B 之间,说明 SST 每 token 的学习效率不如 63B,它的优势来自前 168B token 更便宜。
  • 0.0021 的差异来自训练集上的 EMA 训练 loss、单次运行,没有 seed 方差。在 60B 级 MoE 预训练中,这一量级通常处于 run-to-run 噪声和 LR 调度细节(SST 的 cosine 被拉长,63B 在 334.6B 处结束自己的调度)的影响范围内。

5.2 下游任务评测

Table 3:七个下游任务得分(%)与 held-out ARXIV token NLL(越低越好),等理论累计训练 FLOPs,SST 含两个训练阶段。

Task 47B classic 63B classic 67B SST SST − 63B
OpenBookQA 68.50 71.00 75.00 +4.00
MMLU 58.44 59.39 60.54 +1.15
GSM8K 56.56 55.04 58.38 +3.34
MATH 33.06 33.14 34.36 +1.22
HumanEval (few-shot) 35.98 32.93 37.80 +4.87
MBPP (3-shot) 50.40 50.80 51.40 +0.60
BBH 50.36 52.56 52.74 +0.18
ARXIV NLL (↓) 1.4521 1.4424 1.4421 −0.0003

分析:

  • 作者的结论是「等 FLOPs 训练下 SST 全面最好,最低的训练 loss 确实转化为最好的下游表现」。七项全部领先是一个正面信号。
  • 但幅度要分开看:held-out ARXIV NLL 上 SST 与 63B 只差 0.0003,基本打平。这是最干净的泛化指标,它说明 SST 在语言建模质量上与 63B 持平,而不是显著更好。
  • 下游差异中较大的几项(OpenBookQA +4.0、HumanEval +4.87、GSM8K +3.34)来自小测试集(OpenBookQA 约 1000 题,HumanEval 164 题,1 题约 0.61 分)。单 checkpoint、单种子的这类差异方差很大;注意 47B 在 HumanEval 和 GSM8K 上反而高于 63B,说明这些基准在该规模下本身就很噪。MBPP +0.6、BBH +0.18 基本属于噪声范围。
  • 较稳健的读法是:在等训练 FLOPs 下,SST 与 63B 质量相当或略好,明显好于 47B。

5.3 不同 prefill–decode 配比下的推理代价

附录 A 中权重 $w$ 的依据,即 OpenRouter 的未缓存输入计费份额:设 $P$ 为总输入 token,$C$ 为 cache-read token,$U = P - C$ 为未缓存输入,$D$ 为输出(含推理)token;对每个变体 $v$,$a_v$、$o_v$ 为 09-21 的每百万 token 输入/输出基准价,

$$A = 10^{-6}\sum_v U_v a_v, \qquad O = 10^{-6}\sum_v D_v o_v, \qquad s_{\text{uncached}} = \frac{A}{A + O} \tag{8}$$

Standard 与 Batch 分开计费后再求和,cache-read 计费在分子分母中都排除。这些是基准价估算,不是发票;cache-write 溢价、按上下文长度和 provider 定价、历史价格变动与非 token 费用都没有建模。

据此构造推理代价代理:对 prefill 权重 $w$,

$$\text{cost}(w) = w\,p + (1 - w)\,d \tag{9}$$

其中 $p$、$d$ 分别是 bulk-prefill 和 decode 的激活 body 参数量,以 47B classic 为 1 归一化。$w = 0.75$ 即 75:25 的 P:D 代价配比。

Table 4:Table 1 的底层分量(token 单位为 B,计费单位为百万美元)

| Model | U (B) | C (B) | D (B) | A ($M) | O ($M) | |---|---|---|---|---|---| | GPT-6 Astra | 356.199 | 2232.435 | 24.323 | 3.5600 | 1.2135 | | GPT-6 Astra Pro | 45.611 | 159.946 | 3.255 | 0.4554 | 0.1613 | | Claude Fable 5 | 159.115 | 460.103 | 9.770 | 1.5895 | 0.4822 | | Claude Sonnet 5 | 1126.045 | 4434.707 | 92.872 | 2.2514 | 0.9279 | | Gemini 3.7 Flash | 2100.037 | 5722.111 | 207.001 | 1.5666 | 0.7570 | | Gemini 3.8 Flash | 1299.288 | 4405.218 | 123.644 | 0.9714 | 0.4598 |

标准输入/输出价格(每百万 token 美元):Astra、Astra Pro、Fable 5 为 10/50,Sonnet 5 为 2/10,两个 Gemini 为 0.75/3.75;Batch 价格减半。数据覆盖:两个 Astra 变体各 17 个观测日(9/4–9/20),Gemini 3.8 Flash 19 天,其余模型 30 个标准活动日;Claude Fable 5 只有 10 个 Batch 日(8/22–8/31)。

Figure 5: Estimated inference cost across P:D cost mixes, normalized to 47B classic. Shading spans the six-model OpenRouter reference range; open circles mark break-even points and the filled circle marks 75:25.

结果:

  • 以 47B 为 1:SST 的 $p = 1.120/1.477 = 0.758$,$d = 2.155/1.477 = 1.459$;63B 的 $p = d = 2.016/1.477 = 1.365$。
  • 75:25 时 SST $= 0.933$,63B $= 1.365$,即 SST 比 47B 低 6.7%,比 63B 低 31.6%。
  • 盈亏平衡点:与 63B 在 13.4:86.6,与 47B 在 65.5:34.5。prefill 份额越高,SST 越便宜。
  • 以 OpenRouter 的 67.4–76.7% 为参考 prefill 权重时,SST 在整个区间内都低于两个基线:比 47B 低 1.3–7.9%,比 63B 低 27.7–32.5%。作者原话:「These are analytical cost estimates, not measured serving speedups.」

分析:

  • 相对 47B 的优势很薄且依赖假设:SST 的 decode 激活量比 47B 高 46%,只有 prefill 权重超过 65.5% 时才划算。OpenRouter 区间的下沿 67.4% 距盈亏点只有 2 个百分点,对应的优势只有 1.3%。
  • 相对 63B 的 31.6% 优势主要来自 prefill 激活参数量少了 44%(1.120 vs 2.016),decode 其实略贵(2.155 vs 2.016,+6.9%)。
  • 代理忽略了注意力和 KV 访存(笔者分析):式 (9) 只数权重参数。decode 在实际 serving 中通常是访存瓶颈,长上下文 agentic 负载下 KV 读取往往是大头。SST 的 Decoder 每层都要再读一遍对应 Prefiller 层的 KV,所以 decode 时 full-attention 的 KV 读取次数是 5(Prefiller)+ 5(Decoder)= 10 次/step,而 47B 为 5、63B 为 6。长上下文 decode 下,SST 的 KV 带宽开销约为 47B 的 2 倍。prefill 侧 full-attention 的二次方成本 SST 与 47B 相同(都是 5 个 full 层),省下的主要是 FFN/MoE 与投影部分。
  • Prefill 权重本身的依据是价格而非 GPU 成本:75:25 来自 OpenRouter 计费份额,而价格里的输入/输出比(1:5)是 provider 的商业定价,未必等于 prefill 与 decode 的真实算力比。
  • 因此,KV-invariance 所声称的 serving 收益在本文中没有任何实测:没有吞吐、没有 TTFT/TPOT、没有 KV cache 显存对比、没有端到端 agent 任务成本。

6. 消融与分析

论文没有任何消融实验。以下组件都没有被单独验证:

  • entry bridge(式 3)是否必要,或换成只用 embedding / 只用 Prefiller 状态会怎样;
  • 对齐层 KV 复用 vs 反向配对 vs 单一 KV 源复用(Figure 6 列出了这些选项,但明确说「未假设它们有相同的成本或质量」,也没有跑);
  • Decoder 从源权重复制初始化 vs 随机初始化;
  • 转换时点(为什么是 0.301 $B_{\text{ref}}$)和生长倍数;
  • 保留/重置优化器状态、LR 调度拉长方式。

作者在结论中坦率写道:这些观察「不建立 scaling law、不建立每个组件的因果收益、也不建立实测的 serving 加速」。

缺失的关键对照

对照评审关注点逐项核对:

  1. 更强的从头训练基线:只有两档同家族 scaling ladder(47B、63B),没有从头训练的 67B SST 两塔结构。因此无法区分 SST 的优势来自「中途生长」还是「两塔 + KV 复用架构」本身。
  2. 其他生长基线:没有任何生长方法在等算力下对比:没有 Net2Net、SPARKLING、depthwise stacking(直接把 18 层复制成 36 层的普通深模型)、sparse upcycling,也没有 SZP 类的 shrink-perturb。其中最关键的缺失是「把源模型堆叠成 36 层普通 Transformer 再续训」。它与 SST 的训练折扣完全相同,差别只在于新增层是否产出 KV。这个对照才能回答 KV-invariance 的质量代价是多少。SST 只和两个从头训练的模型比,是把「生长带来的训练折扣」和「KV 不变带来的推理折扣」捆在一起报告。
  3. 其他 KV 复用基线:Introduction 把 YOCO / KV-Mirror / DeepSeek-V4.1-Flash 当作主要的推理侧对手,但没有在等算力下对比任何一个。
  4. 消融下降 vs 最强基线的 margin:由于没有消融,这一检查无法执行。可以对照的是,相对最强基线(63B)的 margin 只有训练 loss 0.0021、held-out NLL 0.0003,远小于常见的组件消融幅度,也小到很难与单种子噪声区分。

按档案的既定原则:引入新信号 ≠ 收益来源。SST 引入了「KV-invariant 的新增容量」这一结构,但质量收益在实验中无法归因到 KV-invariance。KV-invariance 本身是一个成本约束(对质量是负约束),质量收益更可能来自生长的训练折扣加上等效深度翻倍(decode 时执行 36 层)。论文展示的是「生长 + KV 不变」这个组合在单一配置下与从头训练打平或小胜,没有展示 KV 不变本身带来任何质量收益,也没有量化它的质量代价。

7. 讨论:KITE 的设计空间

Figure 6: Connectivity choices within KITE. (a) An SST-style backbone with optional hidden-state connections at the Decoder entry, intermediate layers, and output readout. Dashed paths indicate design choices, not connections that must be enabled together. (b) Three illustrative KV assignments: aligned layer-wise reuse, reversed pairing, and reuse of one Prefiller KV source across multiple Decoder layers. The KV-producing computation remains on the Prefiller side; the illustrated choices are not assumed to have equal cost or quality.

作者强调 SST 只是 KITE 的一种实现,而不是唯一的架构处方。其中对齐的逐层 KV 复用、entry bridge、只用 Decoder 做 readout,都是具体的设计选择。核心要求只有一条:新增容量不能扩展产 KV 的计算,也不能引入需要历史 Decoder 状态的依赖(Figure 6 的边界:未来 KV 的产生不得依赖历史 Decoder 状态)。

可变的自由度包括:

  • KV 分配:Prefiller KV 与 Decoder 层之间不必一一对应或保序。可以反序配对(Decoder 第 1 层读 Prefiller 第 L 层),也可以多个 Decoder 层复用同一个 KV 源;
  • Hidden-state 连接:Prefiller 表示可以在 Decoder 入口、中间层进入,也可以参与最终 readout(output fusion);
  • bridge 算子、两个组件的深度和宽度也都可以变化。

这些选择需要保持因果访问和接口兼容,但计算成本和模型质量未必相同。作者表示实验只探索了设计空间的很小一部分,结果只是确立了所选 SST 配置的性能,而没有找到最优连接模式。

与相关工作的关系(原文 §6)

  • Net2Net:函数保持的模型生长;SPARKLING(2602.02472):宽度渐进学习中的信号保持与对称性打破。Dense-to-MoE upcycling(Sparse Upcycling、He et al.):用稠密 checkpoint 初始化稀疏专家模型。SST 则是在训练中途给已经稀疏的源模型加 Decoder;其对齐初始化改变了预测图,随后依靠联合适应。
  • MQA / GQA 在 query head 间共享 KV head 以减少 KV 存储,Cross-Layer Attention 把共享扩展到相邻层。这些方法减少存储的注意力状态;KITE 的目标则是在不扩展 KV 产出路径的前提下扩容量。
  • YOCO:分离产 cache 与消费 cache 的计算,并用于 prefill 早退。SST 把这一执行模式应用到分阶段模型生长上。区别在于 YOCO 的 cross-decoder 共享一份全局 KV 表示,SST 读对齐的逐层 KV,并且通过源到目标的扩展构建。
  • Mixture-of-Recursions:递归 KV 共享,后续递归复用第一次递归的 KV。作者认为这种固定 KV 的循环可以看作 KITE 背后家族的参数绑定特例;SST 使用独立的 Prefiller/Decoder 参数,在训练中扩容量而不扩展 KV 产出路径。

核心贡献总结

  1. 提出 KITE 范式:把「生长」(训练侧折扣)和「KV 复用」(推理侧折扣)组合起来,规则是新增容量只放在不产 KV 的区域,从而让 bulk prefill 保持源模型尺寸。
  2. 给出简洁的实例 SST:两座对齐的 18 层 MoE 塔,Decoder 逐层读 Prefiller KV,通过无参数的 entry bridge(式 3)接入,源权重复制初始化后联合续训。
  3. 在 67B MoE 规模上做了严格等理论 FLOPs 的对比(笔者复核三者均约为 $5.168 \times 10^{21}$):训练 loss 低于 47B / 63B 从头训练基线,7 个下游任务全部领先,held-out NLL 与 63B 持平。
  4. 把 agentic 负载的 prefill 主导性量化为推理代价的权重,并给出盈亏平衡分析(与 47B 在 65.5:34.5,与 63B 在 13.4:86.6)。

与已归档相关工作的对比

SZP SZP: When Is Warmstarting Effective for Scaling Language Models (University of Freiburg, 2026-05-13)

关系:独立并发(本文未引用 SZP,两者都在等 FLOPs 下检验「从小 checkpoint 生长」对从头训练的收益)· 已加载对方精读

  • 共同关注的问题:两者都针对同一个结构性问题:从小模型 checkpoint 长成大模型,能否在等训练算力下打败从头训练?KITE 在 Introduction 里把「upcycled 模型一般被认为不如同尺寸从头训练的模型」作为出发点。SZP 则系统研究了这一现象何时成立,并把原因归结为「对函数保持的执念」和「生长倍数上限 $g_{\text{upper}}$ 不明」。
  • 相近的技术骨架:两者都是「训练小模型 → 某时刻扩展参数 → 继续训练 → 与 IsoFLOP 从头训练曲线比较」。两者也都放弃了精确函数保持:SZP 用 shrink + zero + perturb,SST 用逐层复制加一个会让 loss 从约 1.745 跳到约 1.90 的 entry bridge。SZP 恰好证明了「函数保持不是有效 warmstart 的必要条件」,为 SST 的非函数保持初始化提供了外部背书。
  • 本文的差异与推进:SZP 研究的是同构放大(宽度/深度按原架构放大),不关心推理代价。KITE 的新意在于把新增容量限定在不产 KV 的位置,把生长从「训练技巧」变成「训练 + 推理联合折扣」的架构范式,并在 67B MoE 规模验证,远大于 SZP 的 14M–1.2B。
  • 可比的方法 / 实验差异:SZP 的实践准则是 $g=2$、在不超过 20 tokens/param 的预算内 warmstart 才划算,超过交叉点后从头训练会追上。SST 的生长倍数恰好约为 2(33.8B → 67B),源阶段 168B token 相对 33.8B 总参数约 5 tokens/param,处在 SZP 认为 warmstart 有效的区间。但 SZP 的交叉点预测也提示:Figure 4(b) 中 SST 与 63B 的差距在 1.0 $B_{\text{ref}}$ 附近持续收窄,更大预算下 63B 从头训练可能追平甚至反超,而 KITE 没有做更大预算或 scaling law 拟合。SZP 做了完整的 IsoFLOP 和 scaling law,KITE 只有单一预算点。

Untied-Grow Untied-Grow: How Model Growth, Recursion and Boundary Operators Influence Scaling (NYU / Q Labs, 2026-09-16)

关系:独立并发(本文未引用,两者的生长骨架几乎同构)· 已加载对方精读

  • 共同关注的问题:两者都认为固定深度的模型「从第一步起就在为后期才需要的容量付费」。Untied-Grow 称之为「过早付费」,KITE 称之为「新增容量在中途而非第 0 个 token 出现,所以买到的质量更划算」。两者都主张在训练中途把模型加深,并以等训练算力下的普通 Transformer 为基准。
  • 相近的技术骨架:骨架惊人地相似。(1) 生长方式相同:Untied-Grow 把训练好的 core 复制后解绑,属于 depthwise stacking;SST 把 18 层源模型逐层复制成 Decoder,decode 时执行深度从 18 翻倍到 36。(2) 接入算子几乎同形:Untied-Grow 的边界算子是 $\mathrm{BO}(h,e) = \mathrm{Norm}(h) + \alpha e$,即归一化当前状态并重注入 prelude 输出;SST 的 entry bridge 是 $\mathrm{RMS}(e_t) + \mathrm{RMS}(h^P_{L,t})$,即归一化 Prefiller 最终状态并加回归一化的 embedding。两者都靠「归一化 + 重注入输入」让新增的深度段以满权重写入、并以输入为条件。
  • 本文的差异与推进:Untied-Grow 的新增 core 是普通的 KV 产出层,推理时 prefill 也要跑完整深度。KITE 的关键推进是让复制出来的那一半只读 KV,从而把深度生长的推理成本从 prefill 中剥离。反过来,Untied-Grow 做了 KITE 缺失的东西:逐架构的算力最优配方、$10^{18}$–$10^{20}$ FLOPs 的 scaling ladder、隔离「生长 / 权重共享 / 边界算子」三轴的对照。Untied-Grow 发现仅边界算子就贡献 1.25× compute multiplier、深度生长贡献到 1.55×。这强烈提示:SST 的质量收益很可能主要来自深度生长 + 类边界算子的 bridge,而不是 KV-invariance,这正是 KITE 缺少「堆叠成 36 层普通 Transformer」对照的代价。
  • 可比的方法 / 实验差异:Untied-Grow 是 dense 小模型(最大 7.4B)加 FineWeb 公开数据,单种子,指数差仅 0.003–0.006;KITE 是 67B MoE 私有数据的单点对比。两者都缺多种子。

被剔除的近似候选(未入选):

  • Hidden Decoding Hidden Decoding(WeChat AI):同为 100B+ MoE 规模的「非参数维度扩展」,但它把每个 token 展开成 n 条流、各自带独立 KV,方向与 KITE 相反(放大 KV / 序列维度来换质量),问题和解法都不同构。
  • KSA KSA(Kuaishou):做 KV cache 的序列级压缩,属于纯推理侧的 KV 瘦身,没有生长 / 训练折扣这条腿。
  • Loopie Loopie(IQuest):循环 Transformer 作为算力匹配的 scaling 轴,权重绑定、每次循环都重新产 KV,且不是中途生长。
  • Qwen3.8-Flash-Next Qwen3.8-Flash-Next:同样按「loss + 训练/prefill/decode 代价」三轴评估架构改动,但评估的是注意力 / 残差改动,不涉及生长或 KV 不变扩容。

8. 讨论与局限性

值得借鉴的设计

  • 问题框架清晰:「只在一侧省算力的方案会被另一侧抵消」这一观察很有价值,把 upcycling 的质量折损和 KV 复用的质量折损放进了同一张账本。
  • KV-invariant 扩容是一个可组合的原则:它与 MoE、GQA、hybrid attention 正交,给「推理时 prefill 主导」的负载提供了一条新的扩参路线。对推荐系统的长序列建模同样有启发:用户长历史的 prefill(或者说 KV 构建)只由一个小塔负责,新增的「预测容量」只在候选/当前位置上计算。这与推荐中「用户序列编码与候选打分解耦」的直觉一致。
  • 训练算力核算严谨:两阶段都计入,三者严格等 FLOPs(笔者复核成立),给出了每 token 前向系数和计入/不计入的项目清单。
  • 诚实的自我限定:正文与附录多处声明「不是实测 serving 加速」「不建立 scaling law」「不建立各组件的因果收益」「这批流量没有被标记为 agentic」。

局限与争议

  1. 推理收益完全是估算:6.7% / 31.6% 只是激活参数量的加权和,既不含注意力计算,也不含 KV 访存;而长上下文 agentic 负载恰恰是 KV 访存主导的场景。按笔者推算,SST 在 decode 时 full-attention 的 KV 读取量约为 47B 的 2 倍,这可能抵消甚至反转相对 47B 的那 6.7%。
  2. 相对 47B 的推理优势依赖 prefill 权重 ≥ 65.5%,而这一权重取自价格份额而非 GPU 成本;decode 激活量比 47B 高 46%。在 decode 重的场景(长推理链、低缓存未命中率)SST 反而更贵。
  3. 只有一个配置、一个预算点、单种子:相对最强基线 63B 的 margin 是训练 loss 0.0021、held-out NLL 0.0003,并且曲线在收敛。没有 scaling law,无法判断更大预算下优势是扩大还是消失(SZP 的交叉点理论提示可能消失)。
  4. 缺失所有关键对照:没有 36 层堆叠生长基线(隔离 KV-invariance 的质量代价),没有其他生长方法,没有 YOCO/KV-Mirror 等 KV 复用基线,没有从头训练的 SST 结构,也没有任何消融。
  5. 收益归因不清:按「引入新信号 ≠ 收益来源」原则,KV-invariance 在逻辑上是一个成本约束,质量收益更可能来自生长 + decode 深度翻倍 + 类边界算子的 bridge(Untied-Grow 独立证明了这些都能带来收益)。
  6. 训练侧的折扣有限:续训阶段 SST 每 token 比 63B 还贵 5%,全部训练折扣来自前 30% 算力的便宜源模型阶段。若源阶段更长,折扣更大但生长后可用的续训预算更少;这种权衡没有被研究。
  7. 基线设置的细节差异:47B 的 token batch(14.68M)与其他模型(12.58M)不同;SST 的 cosine 调度被拉长到 390B,而 63B 在 334.6B 结束。这些差异都可能影响 0.002 量级的比较。

工业落地价值

作者单位为 StepFun,续训使用 128 张 H800,模型规模(67B MoE)属于工业预训练量级。但按「工业署名 ≠ 工业证据」原则:论文没有任何线上部署、serving 实测或产品指标;OpenRouter 数据是第三方公开流量(且与 StepFun 自身模型无关),只用来给代理权重提供依据。因此本文属于「工业团队的架构预训练研究」,而不是工业落地报告。

与已有工作的差异

与 YOCO 相比,SST 读对齐的逐层 KV 而非单一全局 KV,并通过生长而非从头构建。与 upcycling/生长相比,新增部分只读 KV。与 MoR 相比,参数不绑定。它的独特之处在于组合,而不是任何单一组件的新颖性。

精读评分

6 / 10。KV-invariant 生长是一个清晰、可组合且有实际意义的范式,训练算力核算严谨,67B MoE 规模的等 FLOPs 对比也有说服力。但实验只有单一配置、两个同家族从头训练基线、单种子。相对 63B 的 margin 仅为训练 loss 0.0021 / held-out NLL 0.0003,且曲线在收敛。论文没有生长基线、KV 复用基线和任何消融,因此 KV-invariance 的质量代价与收益来源都无法归因。推理节省仅是忽略注意力与 KV 访存的参数量代理,没有任何实测,相对 47B 的优势还依赖 prefill 权重 ≥ 65.5%。