\u200E
最新动态 一手掌握

Dense、MoE之外第三条Scaling路径:交大提出JTok模块,省1/3算力

时间: 2026-03-02 23:05作者: 约翰·明奇



本文第一作者杨烨彬为上海交通大学一年级硕士生,本科毕业于上海交通大学,研究方向为基座大模型架构与scaling law。通讯作者是上海交通大学人工智能学院严骏驰教授,IAPR Fellow、ACM MM大会程序主席、ICML理事。

大模型的发展似乎遵循着一条铁律:根据 scaling law 堆砌参数和数据,模型性能就会遵循负幂律持续增长。但这条路正变得越来越昂贵,因为传统的 scaling 方式始终无法摆脱一个致命的魔咒 ——参数与计算量的深度绑定。在传统的 Dense 模型中,扩展逻辑简单粗暴:加宽、加深 。随之而来的硬伤是:参数规模一旦暴涨,计算量和显存需求就会跟着线性飙升 。在高质量文本数据早就告急的今天,单纯堆稠密参数不仅边际收益骤减,甚至会出现性能倒退的尴尬局面 ——scaling 之路,眼看就要走到头了!

为了松绑参数与计算量,MoE 曾被寄予厚望 。它靠着稀疏激活的专家子网络,在一定程度上实现了模型容量与计算量的解耦 。然而,近期的研究表明,这并非没有代价的免费午餐 :稀疏模型通常具有更低的样本效率 ;随着稀疏度增大,路由负载均衡变得更加困难 ,且巨大的显存开销和通信压力导致其推理吞吐量往往远低于同等激活参数量的 dense 模型 。

LLM 是否还存在新的扩展方向,能带我们走出这个困境?

近日,上海交通大学与交大校友创办的小红书 Hi Lab 联合团队提出了一种全新的 scaling 维度:token-indexed parameters。该方法不依赖主干参数扩展,也不稀疏计算路径,而是通过为每个 token 引入调制向量,以查表 + 逐元素调制的方式提升模型容量。

作者们提出的 JTok / JTok-M 模块,可作为插件形式挂载在 Transformer 每一层,形成静态(JTok)或上下文感知的动态(JTok-M)调制路径,构建 token 与主干路径的交互桥梁。更关键的是,这种机制几乎不增加算力和显存开销,却带来显著且稳定的性能提升。

    在从 650M 到 61B 的模型规模中,JTok-M 显著降低 loss,并在多个下游任务中大幅提分:MMLU +4.1,ARC +8.3,CEval +8.9。达到相同性能,JTok-M 直接砍掉了三分之一的算力需求; 更令人兴奋的是,其 scaling 效果呈现出清晰可预测的幂律规律,可独立扩展、量化收益,完全符合构建新一代 LLM scaling law 的所有要素。

这一创新构建出 Dense、MoE 之外的第三条 scaling 路径,为未来大模型的发展打开了新的方向:

不必加算力、不依赖更多数据,仅靠结构设计与查表调制,就能持续扩大模型容量,重塑性能 - 计算的效率前沿。



    论文标题:JTok: On Token Embedding as another Axis of Scaling Law via Joint Token Self-modulation 论文链接:https://www.arxiv.org/abs/2602.00800

轻量插件式改造,容量飙升却几乎不增计算量

算法设计的核心思想是用 token-id 直接查表取得调制向量,再以元素级乘加方式注入主干,实现模型容量提升,而 FLOPs 几乎不变。



JTok 静态调制:不改主干,外挂轻量插件就搞定

与传统扩展方式不同,JTok 不是加深加宽网络结构,而是为每个 token 引入一个专属调制向量,并在 Transformer 各层通过逐元素乘法对 MLP 残差进行调制,实现无侵入式容量注入。具体来说:所有调制向量保存在一个嵌入表中,每个 token 在每层通过其 ID 查表获得调制向量,调制向量经过归一化后,与当前层 MLP 残差进行逐元素相乘,再写入残差路径。

整个过程无需修改主干结构,仅通过轻量插件式外挂就完成了有效参数的注入。

更关键的是,不会显著增加 FLOPs,也不引入额外通信瓶颈,几乎不影响模型原有的推理 / 训练吞吐。

JTok-M 动态调制:让 Token 懂上下文,适配不同语义

JTok 虽好,但有两个小局限:一是参数量扩展不够灵活;二是同一个 Token 在不同上下文里,语义本就千差万别,总用同一个调制向量,不够贴合实际场景。

虽然静态调制已带来显著收益,但 JTok-M 进一步突破了一个关键限制:

同一个 token 在不同上下文下语义千差万别,调制向量也应因境而异。

JTok-M 为此引入了两个核心机制:

    调制向量池:每个 token 不再只有一个向量,而是拥有一组候选向量,构成语义子空间; 上下文路由器:根据 token 当前上下文的隐状态动态选择 top-K 个向量并加权融合,形成最终调制向量。

这种机制实现了语义敏感 + 稀疏激活 + 插件扩展的三重优势,让模型在几乎不增加核心计算的情况下获得上下文感知能力。

为了保证各向量充分参与训练,JTok-M 还引入了类似 MoE 的路由负载均衡损失,使得调制空间的利用率维持高效、稳定。

三分之一算力节省的工程落地之道

JTok / JTok-M 虽引入了大量新参数,但在系统设计上采取了查表式插件 + 旁路异步调度的范式,使得计算与访存压力都得到了有效隔离与隐藏:

    查表可异步与主干计算重叠,访存开销被调度隐藏; Token 频率长尾分布可利用,同一 token 多次查表合并访问,大幅减少内存压力; 训练阶段支持嵌入并行,推理阶段支持 CPU offload,仅传输需要的向量片段;

在这一系列优化的加持下,即便把 JTok-M 扩展到相当可观的容量,训练吞吐损失也不到 7%;推理阶段的吞吐损失控制在 7.3% 以内,而 GPU 侧几乎不需要额外的显存占用 —— 兼顾了性能提升和工程落地的实用性。

重新定义 Scaling Law:让参数扩展走出计算陷阱

Scaling Law 是理解大模型发展的指南针:

模型的性能表现往往呈现出一种近乎物理规律般的幂律关系 —— 参数越多、数据越多、计算量越大,损失就越低,准确率就越高。

但这个规律,也设下了代价等式:要继续提升性能,就必须同步增加计算量。

JTok-M 的出现,打破了这个性能 - 算力绑定逻辑,建立起了对传统 scaling law 的一次横向扩展。

在论文中,为了量化 JTok-M 的扩展潜力,作者核心解答了两个关键问题:

1. 当主干模型规模扩大时,JTok-M 带来的收益能否保持稳定?

2. 当扩展 JTok-M 本身参数规模时,它表现出怎样的 scaling 行为?

结论让人眼前一亮:

结论 1:适配超大主干,收益不打折,算力省 35%

第一个核心问题:当主干模型从几千万参数涨到几十亿规模时,JTok-M 的性能提升会不会失效?这直接决定了它能不能落地到真正的大模型中。

作者没有只靠实验碰运气,而是先从理论层面把逻辑说透:将 token-indexed 参数融入经典的 scaling law 框架,还提出了 有效参数 的关键假设:

传统模型的性能由主干激活参数(N_c)和训练数据量决定,而 JTok-M 新增的参数(N_n)会和主干参数形成一个扩展比例 η=Nₙ/N_c。再考虑到 JTok-M 的稀疏特性(embedding 本身的稀疏性,Top-K 激活的稀疏性),这些新参数会有一个有效折扣因子 γ,最终组合成有效参数 N_eff。



把有效参数代入原有的 scaling law 公式后,神奇的规律出现了:JTok-M 没有改变模型对算力、数据的依赖逻辑,只是让整个 性能 - 算力 帕累托前沿曲线整体往下平移了一截 —— 这意味着,不管是小模型还是超大模型,要达到同样的性能,JTok-M 需要的算力都会少一大截,而且这个收益是稳定的,和主干规模无关。



后续的严格实验也完全验证了这个猜想:在对数坐标下,JTok-M 的性能 - 算力前沿拟合线和原生模型基线几乎完全平行,只是整体下移了一个固定间隔。计算得出:要达到和原生 MoE 同等的模型性能,JTok-M 能直接节省 35% 的训练算力,这个节省比例在不同模型规模和训练预算下都成立,验证了理论中的 scale-invariance 。对企业来说,这就是实打实的降本增效,同样的预算能训练出更强的模型,或者用更少的成本达到目标效果。



结论 2:自身参数越堆越有用,遵循清晰的幂律缩放

第二个核心问题:JTok-M 自己的参数扩容后,性能会不会很快饱和?能不能成为一个独立的缩放维度?

作者做了一组精准的控制变量实验:固定主干模型结构和训练数据量,只通过调整参数扩展率 η 来增加 JTok-M 的自身容量。结果呈现出非常漂亮的规律性:在对数坐标下,模型的验证损失随着 JTok-M 参数的增加近乎线性下降 —— 每翻倍一次 JTok-M 的参数,测试损失就稳定降低约 0.0118,而且完全没有出现性能饱和的迹象.



这背后的意义非常关键:JTok-M 的 token-indexed 参数,本身就是一个和主干参数、数据量、专家稀疏度完全正交的新缩放维度。它不用依赖主干模型加参,也不用额外堆砌高质量数据,只要单纯扩大 JTok-M 自己的参数规模,就能稳定提升性能,而且效果可以精准预判。

就像当初稠密模型的 scaling law 一样,开发者现在可以拿着这个规律做精准规划:想让模型损失降多少,需要给 JTok-M 加多少参数,一目了然,不用再靠盲目堆参试错。

这正是 scaling law 的核心魅力 —— 它不是一次偶然的性能提升,而是为大模型的发展提供了一条稳定、可持续、低代价的新路径,让后续的模型迭代有章可循。

总结而言,传统 Scaling Law 实际上只提供了两个方向:参数规模(N)与数据规模(D)。MoE 试图用稀疏激活释放计算,但仍困在 N 与 C 的线性关系中。

JTok-M 的意义在于:

    引入了一种新的扩展形式:token-indexed capacity; 构建了从理论 → 实现 → 验证 的完整闭环; 把 Scaling Law 从二维推向三维:参数、数据、token-indexed 结构共塑未来 LLM 的成长路线。

JTok-M 不只是一个插件,更是一种对大模型如何继续扩展的重新定义。

下游泛化能力:知识、推理、数学能力全面飙升

光说缩放律和算力节省还不够,JTok/JTok-M 能不能落地,关键看它在实际下游任务中好不好用 —— 能不能实实在在提升模型的知识储备、推理能力,甚至是代码和数学解题水平?

作者做了超全面的验证:覆盖知识、推理、代码、数学四大类共 14 个子任务,在 dense 和 MoE 两种基座、多种参数规模的模型上测试,全程保持 backbone、训练数据、训练流程完全一致,只外挂 JTok/JTok-M 模块,结果堪称全任务提分,而且模型越大、任务越难,收益越明显!

Dense 基座:小模型也能大提升

在 1.5B 参数的 dense 模型上,加了 JTok 模块后,14 项任务的平均准确率直接暴涨 4.32 个百分点 —— 相当于在原有基础上提升了近 20% 的相对收益,性价比拉满!

其中重点任务的提升尤为亮眼:

    通用知识推理(MMLU):+4.6 个点,知识储备更扎实; 科学常识推理(ARC-C):+5.8 个点,对复杂常识的理解更透彻。



MoE 基座:JTok-M 发力,推理能力突飞猛进

在 MoE 基座上,JTok 已经能稳定带来平均 2.5 个点的提升,而升级后的 JTok-M 直接把收益拉满,尤其是在推理、数学这类高难度任务上,表现堪称惊艳:

    1.5B 总参数量(激活参数量 250M):平均准确率 + 3.91; 3.2B 总参数量(激活参数量 500M):平均准确率 + 5.59,其中科学常识推理(ARC-C)+7.25,数学解题(GSM8K)+6.31。

这说明 JTok-M 不仅帮模型记住了更多知识,更实实在在增强了模型的逻辑推理和复杂问题解决能力 —— 这正是大模型落地的核心需求之一。

17B 大模型验证:越大越能打,样本效率还超高

下游任务的收益可不是小模型专属,作者在 17B 总参数量(激活参数量 2B)的超大 MoE 模型上做了验证,最终总参数量达到 61B,还全程跟踪了 MMLU、ARC-C、CEval 等 6 个核心基准的训练曲线:



曲线趋势一眼就能看出优势:

    样本效率超高:训练刚进行到几 B token 时,加了 JTok-M 的模型就已经全面反超原生 MoE,不用等训练后期,就能看到明显提升; 大模型收益更猛:训练结束时,MMLU 提升约 4 个点,而 ARC-C(科学推理)、CEval/CMMLU(中文知识推理)等难度更高的任务,直接暴涨 8-9 个点!

核心结论:全场景适配,越难越能打

JTok/JTok-M 的下游泛化能力,藏着两个关键亮点:

    适配性极强:不管是 dense 还是 MoE 基座,从小模型到 17B 超大模型,都能稳定提分,不用改动原有训练流程,插件式升级就能见效; 针对性提分:不仅能提升知识记忆类任务的表现,更能强化推理、数学等复杂任务的能力,而且任务越难、模型规模越大,提升效果越显著 —— 这正好戳中了大模型工业化落地的核心痛点,让模型在实际应用中更能打。

相关工作对比

最近,Scaling Embedding 成为了突破大模型参数效率瓶颈的一个核心大方向。在这一探索路径上,业界涌现了多篇极具启发性的工作,包括 DeepSeek 的 Engram、Meta 的 STEM,以及我们最近推出的 JTok-M。 虽然这三篇工作都致力于通过大规模 Embedding 来低成本地扩展模型容量,但它们在设计哲学、切入点以及核心洞察上各有侧重,共同描绘了这一领域的广阔前景:



PS:JTok,技术上代表 Joint Token。但在众多校友心里,Joint 寓意着 “交通”,读音也与 “交通” 撞轨;它更藏着一个交大人心照不宣的梗 ——“JT(交通)OK”。以此命名,并沿袭 Jaccount 的命名文化,向交大 130 周年校庆致敬!