1. 问题与动机
推荐大模型扩容遇到 embedding collapse:学到的表示集中在低秩子空间,限制表达多样性、浪费扩出的容量。RankMixer 虽比 DCNv2、xDeepFM、AutoInt 等传统模型好,但提升不足。作者用 effective rank(erank = ||X||_F² / ||X||_2²,超越代数秩刻画表示容量)逐层度量:Criteo 上 RankMixer 末层 erank 仅略高于原始 embedding(提升「modest」未能阻止坍缩);Avazu 上 P-FFN 的收缩占主导、collapse 重现。关键现象被概括为「阻尼振荡轨迹(damped oscillatory trajectory)」——token mixing 抬高 erank、P-FFN 压低 erank,逐层交替收缩却始终未能消除坍缩。实验在 Criteo 与 Avazu 两个各 3300 万+ 训练样本的工业级 CTR 基准上进行,坍缩问题在这种规模下尤为现实。
2. 方法详解
RankMixer 三件套:tokenization 把异构 embedding 投到统一 D 维 token 矩阵;token mixing 是无参的 block-transpose(成对置换 token);P-FFN 是 per-token GELU FFN。理论诊断:block-transpose 的秩扩张有界 erank(M) ≤ 2(k+μ)(k、μ 为原矩阵与转置矩阵的有效秩);带正齐次激活的标准 FFN 以概率收缩 erank(Theorem 2.1/2.2)。RankElastor 以「Expand More, Shrink Less」为目标提出两组件:
(1) 参数化全 token mixing:把无参 block-transpose 换成可学习稠密变换 vec(M⊤)=LN((W+I)vec(X⊤)),其中 W∈ℝ^(TD×TD) 为可学习混合矩阵,(W+I) 的残差形式保证初始接近恒等,把 block 粒度从 d>1 降到 d=1(细粒度),实现比 Kronecker 约束的 block mixing 更丰富的坐标交互;Theorem 3.1 证明 d=1 情形严格扩张表达力。
(2) GLU 改进的 P-FFN:把 GELU 换成门控形式 Z_t=(GELU(M_tW₁)⊙(M_tW₂))W₃+M_tW_r(⊙ 为逐元素乘),遵循 GLU 原则用乘性特征交互防止秩收缩;Theorem 3.2 证明乘性项引入 degree-2 交互,代数秩至少抬升 k(k+1)/2,并保证 effective rank 增量 δ>0。
复杂度:token mixing 从 RankMixer 的 O(TD)、零参数升到 O(T²D²) 计算与参数(主要开销来源);P-FFN 加常数因子(扩展比 3r+1 vs r);净结果是训练时间增加 10-15%,显存影响可忽略。
3. 核心技术贡献
(1) 首次系统刻画 RankMixer 的 collapse:把经验 effective-rank 动力学与理论(Theorem 2.1-2.2)结合,解释为何 RankMixer 的提升不彻底。(2) RankElastor 架构:参数化全混合 + GLU P-FFN,给出可证明的 collapse 缓解与「谱稳健(spectrum-robust)表示」理论保证(Theorem 3.1-3.2)。(3) 工业规模验证:在 CTR 与行为序列建模上都取得一致增益、表示多样性提升与更优的参数 scaling 行为,并开源代码。
4. 实验结果分析
CTR (Table 2):Criteo AUC RankElastor 0.81482 vs RankMixer 0.81375、DCNv2 0.81365、MLP 0.81307;Avazu 0.79323 vs RankMixer 0.79270、DCNv2 0.79258。约 0.001 AUC 增益,在工业级基准属「统计显著」(CTR 领域 0.001 AUC 常对应可观线上收益)。collapse 缓解 (Fig 6):RankElastor 跨层维持持续更高的平均 erank,Avazu 上保持稳定的 erank 增长,而 RankMixer 呈阻尼振荡且扩张有限;分布偏移(Fig 5)显示参数化全混合相对 block-transpose 带来明显更大的 erank 增益。scaling (Fig 7-8):宽度、深度、深+宽联合扩展三种方式下,RankElastor 都比 RankMixer 有更大 AUC 提升与更低 LogLoss,且联合扩展收益大于单维扩展,与 LLM 的 scaling 规律一致——这点对「推荐 Scaling Law」尤为关键,说明先解决 collapse 才能让加参数真正兑现。消融 (Table 3):去掉 full mixing→Criteo 0.81413、换 ReLU FFN→0.81326、换 GELU FFN→0.81349,两组件协同——只把 GLU 加到原始 RankMixer 仅 0.81393(对比 0.81375),说明需先增强 token mixing 才有协同效应。泛化 (Table 4):行为序列建模 KuaiVideo gAUC 0.6731 vs 0.6691、AUC 0.7514 vs 0.7482;TaobaoAd gAUC 0.5778 vs 0.5763、AUC 0.6522 vs 0.6508,证明方法不局限于 CTR 估计。效率 (Fig 4):10-15% 运行时增加、GPU 显存相当,作者称对大规模推荐「practically negligible」。
5. 局限与未来方向
作者明确:(1) 有效秩动力学与理论分析专门针对 token-transformation 类(RankMixer 式)架构,结论未必推广到根本不同的设计范式;(2) 理论框架适用面较窄,RankMixer 类仍有大量探索空间;(3) 分析假设了特定初始化与 incoherence 条件,现实偏离未被完全刻画。此外论文未给在线 A/B,也未讨论稠密 (TD×TD) 混合矩阵在 serving 规模下的推理延迟、量化/蒸馏等服务优化策略,以及 T×D 远大于实验设置时的可扩展性。
6. 工业实践意义
对做排序模型放大的团队最有用的一点:扩参数前先量化表示的 effective rank,若已坍缩则单纯增容会被低秩瓶颈吃掉——erank 是判断 dense scaling 是否真正生效的可量化诊断指标,可作为放大实验的前置体检项。RankElastor 的两点改造(可学习全混合替代无参 block-transpose、GLU 替代 GELU FFN)可直接嫁接到 RankMixer 类统一 token 排序模型,10-15% 训练开销换来更好的 scaling 与跨场景泛化(CTR 与行为序列双验证)。但 O(T²D²) 的 token mixing 是 serving 隐患,落地需评估推理延迟、可能配合量化/蒸馏,建议增量上线并密切监控服务侧性能。代码已开源(Zenodo DOI 10.5281/zenodo.20252036)。对社区搜索/推荐排序而言,若已在用统一 token 排序骨干,这是低成本提升表示表达力与放大效率的可借鉴改造。