TM20K：老师读全文，学生读缩写

TM20K 用「老师读全文、学生读缩写」的两阶段蒸馏范式破解超长序列的效率困境：一次性训练的 teacher 吃完整 20K token 但永不上线，学生模型用三种由注意力观测反推出来的 token merge 把 20K 压到平均 1.8K 后上线，再靠 KD 把 teacher 约 85% 的收益捞回来——最终线上 ADSS +1.036%，serving 延迟只涨 5.6%。

超长序列建模的「效果 - 效率」死结

超长用户行为序列建模已被反复证明能同时刻画长期与短期兴趣、显著提升推荐质量，是当前公认的 scaling 方向之一。但字节在自家电商广告系统里把最大序列长度从 5K 拉到 20K 时发现：即使已经上了 FlashAttention 和 M-Falcon serving，代价依然是灾难性的。

现有两条路线，各有硬伤

论文把业界的 trade-off 方案归成两类，并逐一点出问题。

路线 A：序列长度压缩（Sequence Length Compression）

路线 B：轻量注意力架构（Lightweight Attention）

TWIN 的 ETA、STCA 的堆叠 target-to-history cross attention、EST 的 LCA + CSA、HyTRec 的「线性注意力管长期 + softmax 注意力管近期」混合、MTFM 的 GQA + 混合 target attention、LONGER 的 Perceiver 式「底层 CA + 上层 SA」……

路线 C（论文选择的方向）：两阶段框架

除了上述端到端方案，还有一类更灵活的思路：两阶段。

• 知识蒸馏路线：以 ExFM 为代表——构建一个 foundation-model 式的 teacher，同时服务多个 vertical model（学生），把构建成本摊薄。
• Foundation model 路线：通过缓存 user embedding / item embedding 来提升下游任务，以及近期的 instance-as-token 范式。

TM20K 属于第一类，但差异是：别人用 KD 解决「模型容量/特征规模」的 scaling，本文用 KD 专门解决「序列长度」的 scaling——teacher 保留全部 token，student 做高效训练与推理。作者也明确说，foundation model 的整合留作未来工作。

TM20K 在方法谱系中的位置

方法部分的每一个设计都不是拍脑袋来的。论文先做了两组前置研究：（1）架构选型实验决定用 FA 还是 TA；（2）注意力分数统计直接反推出三种 token merge 策略。这一节是理解方法的必要前提。

3.1 FA vs TA：为什么坚持全注意力

先把两种架构写清楚。Full Attention（FA）把 target 拼到序列前面做整体自注意力：

而 Target Attention（TA）把注意力简化成「target 对序列的一次 cross attention」：

实验结论（基线 = 5K + TA）

最关键的证据：seq-seq 注意力不能丢

这一段是在回答一个非常具体的工程问题：注意力矩阵里，除了 target 那一行，其他行能不能直接砍掉？先把矩阵画出来才好理解。

行 = query（谁在提问），列 = key（问谁）。假设历史只有 5 个行为 $h_1 \sim h_5$（$h_5$ 最近）加上 target $t$，矩阵是 $6\times 6$（因果掩码使右上角为空，每行 Softmax 后加起来 = 1）：

论文的做法是把矩阵里的数值分成两堆——堆 A = 最后一行的所有数值（target-seq），堆 B = 前 5 行的所有数值（seq-seq）——各画一条概率密度曲线，看两堆数值的分布形态。结果是两条曲线基本重合。

为什么「形态重合」能算证据？因为如果 seq-seq 是纯粹浪费算力的计算，我们应该看到两种退化形态之一：

而实测是：seq-seq 的分布和 target-seq 一样有结构、一样有区分度。既然 target-seq 那一行公认有用（TA 就只用它），形态相同的 seq-seq 就很难说是噪声。进一步放大到「5 个 query token（1 个 target + 4 个序列 token）× 最近 50 个 key token」的局部热力图后可以看到：那 4 个序列 token 的行同样出现明显亮格子，而且亮的位置各不相同——每个历史 token 都在有选择地关注特定的其他历史 token，这种「有选择」就是在做真实的信息加工。

真正定案的是表，图只是在解释「为什么会掉点」。这个顺序不能反——如果只有图没有表，「seq-seq 有结构」完全可以是「有结构但对 CVR 没用」。

3.2 三组注意力分数观测 → 三种 token merge

这是全文最精彩的部分：三种 merge 策略，一一对应三组真实统计观测，不是凭直觉设计的。这一节的实验设置写得很压缩，下面拆开讲。

第一步：为什么必须先训一个 20K 的 FA 模型

注意力分数是模型学出来的，不是数据自带的属性。随机初始化的模型，注意力矩阵基本是均匀噪声，看不出任何规律。所以必须先有一个训到收敛的、吃完整 20K 序列的 FA 模型——这个模型就是后来的 teacher。

第二步：为什么只取一个局部子矩阵——因为全量存不下

完整注意力矩阵有多大？序列 20000 + 1 个 target，矩阵是 $20001 \times 20001$：

一个样本、一个 head、一层就 1.6GB，而模型有 6 层、多个 head，还要跑一批样本——根本不可能落盘。所以作者做了三刀裁剪：

裁完就是 $A \in \mathbb{R}^{5\times 20000}$ = 10 万个数，轻松落盘。

第三步：这个矩阵具体长什么样

行 = query，列 = key，每行经过 Softmax 所以一行加起来 = 1。关键是每一列都挂着两个元信息标签：

第四步：三组观测 = 用三种方式切同一块数据

这是这节最优雅的地方——一次 dump，三种切法：

注意观测一、二都在同一份数据的列上做文章，只是分组标签不同；只有观测三需要换一层再 dump 一次。

观测一：同 product ID 的 token，注意力分数高度接近 → LITM

对第 $t$ 个 query，计算任意两个 key token 之间的注意力分数绝对差：

结果：同 product ID 的 token 对，$\Delta$ 分布明显更集中于 0。

回看 Eq. 2：$E_{s_i} = \mathrm{SumPooling}(E_{p_i},\, E_{o_i})$ ——token 不只是 PID。$p_i$ 是关键 ID 特征，$o_i$ 是其他特征（类目、交互类型、时间戳等）。所以同 PID 的两个 token，$E_{p_i}$ 那一半相同，$E_{o_i}$ 那一半不同。差异来源按影响排序：

所以完整的结论应该是：同 PID 且位置相近时，注意力差异可忽略；位置一旦拉远，差异就恢复了。$T=3$ 正是这条曲线的拐点。「为什么 $\Delta \ne 0$」这个疑问和这个消融实验，本质是同一件事的两面。

观测二：最近 10% 的 token 占掉约一半注意力 → PATM

把序列按位置索引分桶（位置索引越小 = 越近期），对每个桶 $[i_k, i_{k+1}]$ 累加注意力质量：

结果：最近的 top 10% token 就贡献了接近一半的注意力总权重（Figure 3b 的累积概率曲线在 10% 处已到 ~0.5）。

这个判断不需要猜，从方法本身就能反推。看 PATM 的配置 $R = [[1,1000],[1001,5000],[5001,10000],[10001,20000]]$、$K=[1,2,3,4]$ ——这是固定的绝对区间。而且更关键的是：

那为什么横轴用「%」而不是绝对位置？因为要跨用户可比——用户序列长度差异极大：

如果横轴用绝对位置，两类用户的曲线没法叠在一张图上。用百分比归一化后才能画出统一的累积分布曲线。

观测三：上层注意力更尖锐、底层更均匀 → LPTM

逐层计算注意力概率分布的熵：

熵的直观含义：

结果：上层（upper layer）熵更小、注意力分配更尖锐；底层（bottom layer）熵更大、注意力更均匀。这与 LLM 领域的注意力分析结论一致。注意这一组观测需要换层重新 dump 一次，不像观测一、二那样共用同一份数据。

整个 3.2 节的流程串起来

4.1 整体框架：一个 teacher，多个 student

TM20K 的输入输出链路是这样的：原始电商行为序列（最长 20K，每个 token = 商品 ID + 侧信息）→ CPU 侧数据处理做 LITM + PATM 合并 →查 embedding 表并 sum pooling → 6 层 full attention transformer（层间做 LPTM）→ 序列表示 $x_{\text{seq}}$ 与非序列特征一起送入 RankMixer 做高阶交互 → 输出 CVR 预估。Teacher 和 student 共享这套架构，唯一区别是 teacher 不做任何 token merge。

形式化地，单个 token 的表示由各个特征 embedding 求和得到：

然后走序列建模与特征交互：

4.2.1 LITM：局部同 ID Token 合并

动机来源：观测一——同 product ID 的 token 注意力分数高度接近。

做什么：在局部窗口内把相同关键 ID（product ID）的 token 求和聚合成一个 token。注意关键词是 local-wise（局部）——不是全序列去重，而是只合并「位置间隔不超过阈值 $T$」的同 ID token。

4.2.2 PATM：位置自适应 Token 合并

动机来源：观测二——最近 10% 的 token 占掉约一半注意力。

做什么：把序列切成 $B$ 个不重叠分段，每段用不同的压缩因子：近期段几乎不压，远期段狠狠压。

4.2.3 LPTM：层间金字塔 Token 合并

动机来源：观测三——上层注意力更集中，底层更均匀。

做什么：随着 transformer 层数上升，逐步把序列对折，形成「底层长、上层短」的金字塔。

4.2.4 工程实现：CPU / GPU 的分工艺术

三种 merge 按顺序叠加：LITM → PATM → LPTM。但落地时的关键决策是「谁在哪算」。

4.3 Teacher 模型与知识蒸馏

Teacher 的定位非常明确：不追求效率、独立训练、缓存 logits 供所有 student 使用。

Teacher 的损失

Student 的双塔损失

Student 设计了两个独立的预测头：main tower（只学真实标签）和 distillation tower（同时学标签和 teacher）。

4.4 两个额外的工程设计

论文的基线已经用了一整套优化：User-Level Training（类似 STCA 的 RLB）、FlashAttention、混合精度、Remove-Padding、M-Falcon Serving。在此之上又加了两个。

4.5 与相邻工作的核心区别

5.1 实验设置

5.2 主实验：效果与效率的全景对比

5.3 三种 Token Merge 的消融

从 20K Full（= teacher）出发逐步叠加 merge：

5.4 超参与蒸馏权重

LITM 的间隔阈值 $T$

结论：$T=1 \to 3$ 是纯赚（−39% → −43%，AUC 不变）；$T=3 \to 10$ 才开始付代价（AUC 损失从 0.02% 翻倍到 0.04%）。说明「同一次浏览会话内的重复点击」大约就在 3 个位置的间隔内，超出这个范围就是不同的兴趣事件了。

PATM 的分段配置

结论：PATM 的效果损失对压缩强度更敏感（0.05% ~ 0.09% 的区间，比 LITM 的 0.02% ~ 0.04% 大一倍）。这符合直觉——PATM 是「不管内容、按位置强制合并」，而 LITM 有语义依据（同商品）。所以 PATM 应该调得比 LITM 保守。

蒸馏权重 $\lambda$

其他实现选择的消融（表 6）

5.5 线上 A/B 实验

在字节电商广告线上系统做 A/B，基线是 5K 序列的生产模型。

5.6 作者自陈的局限

这篇论文最值得学的三件事

几个我觉得特别锐利的判断

关于 EffLen 的揭穿。表 1 那一列 EffLen 是全文最"不留情面"的设计。业界论文都爱标榜 MaxLen 支持到 $10^6$，但真正参与计算的只有 $10^2$。这等于在说：你们号称支持百万序列，实际上模型只看到了几百个 token，剩下的都在检索阶段被扔掉了。把 MaxLen 和 EffLen 并列摆出来，是一种非常有力的学术批评。以后读长序列论文，第一件事就该是问 EffLen 是多少。

关于与 EST 结论冲突的处理。面对「EST 说 seq-seq 依赖不重要」这个矛盾结论，作者没有简单说对方错了，而是指出EST 用了检索式压缩，检索本身在强化 target-seq 相关性，从而系统性压低了 seq-seq 的边际价值。这是教科书级别的「结论依赖于实验设定」辨析——同样的实验现象，在不同的数据预处理下会导向相反结论。这提醒我们：任何"某模块不重要"的消融结论，都要先检查上游预处理是否已经替它做了工作。

关于 FA 的 GPU 利用率更高。「计算量大的方案反而更快」是这篇论文最反直觉的发现。现代 GPU 的瓶颈早已从 FLOPs 转向访存带宽和算子规整度。TA 的一个 query × 长 key 是典型的 memory-bound 细碎操作；FA 的大矩阵乘法能打满 Tensor Core。所以「减少计算量」不等于「变快」，选架构时必须看实测利用率而不是理论复杂度。

可以直接借鉴的工程 checklist

• CPU 侧只搬整数，不搬浮点。规则型预处理（去重、分段、分组）放 CPU，但要合并原始特征 ID 而非 embedding，CPU↔GPU 之间只传 ID + 分组元信息。这一条能省掉一个数量级的通信量。
• KD 权重的量级校准比数值本身重要。$\lambda = 50$ 看着离谱，但它的意义只是「让 $\lambda\ell_{kd}$ 和 $\ell_{ce}$ 同量级」。换场景时先算两个 loss 的量级比，而不是照搬 50。且 30~100 都行，鲁棒性很好。
• 蒸馏一定要用独立的塔。不要把 KD loss 加到上线用的那个头上，否则预估值校准会被软标签污染。共享骨干 + 独立头，让知识走梯度而不走输出。
• 超长序列训练必须看长周期曲线。QK Norm 的消融就是血泪教训：前期 +0.00%，后期 −0.45%。任何「早期看起来无用」的稳定性组件，都要跑到后期再判断。
• 盯 P90L 而不只是 AvgL。P90L 决定 padding 张量大小和长尾延迟。TM20K-S 的 P90L 从 20K 降到 2.6K（−87%），这才是延迟能压在 +5.6% 的真正原因。
• 训练后期可以关掉 KD。知识吸收完之后 KD loss 就没用了，关掉可以省掉 teacher logits 的缓存读取和存储成本。
• Stack Sequence 值得单独抄。把 batch 内有效 token 重新均匀分配 + 记录真实计数，替代 pad-to-max，20K 场景省 10GB 显存。这是个与序列建模方法无关的通用优化。

我的疑问与可延伸的方向

• Teacher 上限就一定是天花板吗？Student 拿到 teacher 的 85%，但 teacher 自己也只有 +0.26%。如果 teacher 再往上（比如 50K 序列、更宽的模型），student 能捞回的绝对值会不会更高？论文没有探索 teacher scaling 与 student 收益之间的关系曲线，这可能是下一篇的内容。
• 三种 merge 的顺序能否交换或联合优化？现在是固定 LITM → PATM → LPTM。如果先 PATM 再 LITM，或者把 LITM 也做成层间的，效果如何？论文只给了「叠加」的消融，没给「排序」的消融。
• 可学习的 merge 会不会反而更差？作者把「让 merge 可学习」列为未来工作，但我怀疑这里有个隐藏权衡：规则式 merge 可以在 CPU 上做、可以缓存、可以和训练流水线并行；一旦变成可学习的（需要 embedding 参与决策），就必须搬到 GPU 上，那个 7.5× 吞吐提升可能就没了。这不是纯粹的效果问题，而是效果与可实现性的冲突。
• 为什么 KD 对 ADVV 的提升（+51% 相对）远大于对 AUC 的提升？ADVV +0.515% → +0.780%，而 AUC 只从 +0.15% 到 +0.22%。这暗示 teacher 教给 student 的知识在高价值样本上尤其有效——可能是因为高价值转化往往依赖更长期的兴趣线索，正是 20K 全序列的独有优势。如果这个猜想成立，那么按样本价值加权的蒸馏可能还有额外收益。
• PATM 的分段配置能否随用户自适应？现在 $R$ 和 $K$ 是全局固定的。但重度用户和轻度用户的行为密度差异巨大——重度用户最近 1000 个行为可能只跨 3 天，轻度用户可能跨 3 个月。按时间窗而不是位置索引分段，或者按用户活跃度选择配置，可能比全局固定更合理。