Dense、MoE之外第三條Scaling路徑：交大提出JTok模塊，省1/3算力

本文第一作者楊燁彬?yàn)樯虾＝煌ù髮W(xué)一年級(jí)碩士生，本科畢業(yè)于上海交通大學(xué)，研究方向?yàn)榛竽Ｐ图軜?gòu)與scaling law。通訊作者是上海交通大學(xué)人工智能學(xué)院嚴(yán)駿馳教授，IAPR Fellow、ACM MM大會(huì)程序主席、ICML理事。

大模型的發(fā)展似乎遵循著一條鐵律：根據(jù) scaling law 堆砌參數(shù)和數(shù)據(jù)，模型性能就會(huì)遵循負(fù)冪律持續(xù)增長(zhǎng)。但這條路正變得越來(lái)越昂貴，因?yàn)閭鹘y(tǒng)的 scaling 方式始終無(wú)法擺脫一個(gè)致命的魔咒 ——參數(shù)與計(jì)算量的深度綁定。在傳統(tǒng)的 Dense 模型中，擴(kuò)展邏輯簡(jiǎn)單粗暴：加寬、加深 。隨之而來(lái)的硬傷是：參數(shù)規(guī)模一旦暴漲，計(jì)算量和顯存需求就會(huì)跟著線性飆升 。在高質(zhì)量文本數(shù)據(jù)早就告急的今天，單純堆稠密參數(shù)不僅邊際收益驟減，甚至?xí)霈F(xiàn)性能倒退的尷尬局面 ——scaling 之路，眼看就要走到頭了！

為了松綁參數(shù)與計(jì)算量，MoE 曾被寄予厚望 。它靠著稀疏激活的專(zhuān)家子網(wǎng)絡(luò)，在一定程度上實(shí)現(xiàn)了模型容量與計(jì)算量的解耦 。然而，近期的研究表明，這并非沒(méi)有代價(jià)的免費(fèi)午餐 ：稀疏模型通常具有更低的樣本效率 ；隨著稀疏度增大，路由負(fù)載均衡變得更加困難 ，且巨大的顯存開(kāi)銷(xiāo)和通信壓力導(dǎo)致其推理吞吐量往往遠(yuǎn)低于同等激活參數(shù)量的 dense 模型 。

LLM 是否還存在新的擴(kuò)展方向，能帶我們走出這個(gè)困境？

近日，上海交通大學(xué)與交大校友創(chuàng)辦的小紅書(shū) Hi Lab 聯(lián)合團(tuán)隊(duì)提出了一種全新的 scaling 維度：token-indexed parameters。該方法不依賴(lài)主干參數(shù)擴(kuò)展，也不稀疏計(jì)算路徑，而是通過(guò)為每個(gè) token 引入調(diào)制向量，以查表 + 逐元素調(diào)制的方式提升模型容量。

作者們提出的 JTok / JTok-M 模塊，可作為插件形式掛載在 Transformer 每一層，形成靜態(tài)（JTok）或上下文感知的動(dòng)態(tài)（JTok-M）調(diào)制路徑，構(gòu)建 token 與主干路徑的交互橋梁。更關(guān)鍵的是，這種機(jī)制幾乎不增加算力和顯存開(kāi)銷(xiāo)，卻帶來(lái)顯著且穩(wěn)定的性能提升。

• 在從 650M 到 61B 的模型規(guī)模中，JTok-M 顯著降低 loss，并在多個(gè)下游任務(wù)中大幅提分：MMLU +4.1，ARC +8.3，CEval +8.9。達(dá)到相同性能，JTok-M 直接砍掉了三分之一的算力需求；
• 更令人興奮的是，其 scaling 效果呈現(xiàn)出清晰可預(yù)測(cè)的冪律規(guī)律，可獨(dú)立擴(kuò)展、量化收益，完全符合構(gòu)建新一代 LLM scaling law 的所有要素。

這一創(chuàng)新構(gòu)建出 Dense、MoE 之外的第三條 scaling 路徑，為未來(lái)大模型的發(fā)展打開(kāi)了新的方向：

不必加算力、不依賴(lài)更多數(shù)據(jù)，僅靠結(jié)構(gòu)設(shè)計(jì)與查表調(diào)制，就能持續(xù)擴(kuò)大模型容量，重塑性能 - 計(jì)算的效率前沿。

• 論文標(biāo)題：JTok: On Token Embedding as another Axis of Scaling Law via Joint Token Self-modulation
• 論文鏈接：https://www.arxiv.org/abs/2602.00800

輕量插件式改造，容量飆升卻幾乎不增計(jì)算量

算法設(shè)計(jì)的核心思想是用 token-id 直接查表取得調(diào)制向量，再以元素級(jí)乘加方式注入主干，實(shí)現(xiàn)模型容量提升，而 FLOPs 幾乎不變。

JTok 靜態(tài)調(diào)制：不改主干，外掛輕量插件就搞定

與傳統(tǒng)擴(kuò)展方式不同，JTok 不是加深加寬網(wǎng)絡(luò)結(jié)構(gòu)，而是為每個(gè) token 引入一個(gè)專(zhuān)屬調(diào)制向量，并在 Transformer 各層通過(guò)逐元素乘法對(duì) MLP 殘差進(jìn)行調(diào)制，實(shí)現(xiàn)無(wú)侵入式容量注入。具體來(lái)說(shuō)：所有調(diào)制向量保存在一個(gè)嵌入表中，每個(gè) token 在每層通過(guò)其 ID 查表獲得調(diào)制向量，調(diào)制向量經(jīng)過(guò)歸一化后，與當(dāng)前層 MLP 殘差進(jìn)行逐元素相乘，再寫(xiě)入殘差路徑。

整個(gè)過(guò)程無(wú)需修改主干結(jié)構(gòu)，僅通過(guò)輕量插件式外掛就完成了有效參數(shù)的注入。

更關(guān)鍵的是，不會(huì)顯著增加 FLOPs，也不引入額外通信瓶頸，幾乎不影響模型原有的推理 / 訓(xùn)練吞吐。

JTok-M 動(dòng)態(tài)調(diào)制：讓 Token 懂上下文，適配不同語(yǔ)義

JTok 雖好，但有兩個(gè)小局限：一是參數(shù)量擴(kuò)展不夠靈活；二是同一個(gè) Token 在不同上下文里，語(yǔ)義本就千差萬(wàn)別，總用同一個(gè)調(diào)制向量，不夠貼合實(shí)際場(chǎng)景。

雖然靜態(tài)調(diào)制已帶來(lái)顯著收益，但 JTok-M 進(jìn)一步突破了一個(gè)關(guān)鍵限制：

同一個(gè) token 在不同上下文下語(yǔ)義千差萬(wàn)別，調(diào)制向量也應(yīng)因境而異。

JTok-M 為此引入了兩個(gè)核心機(jī)制：

• 調(diào)制向量池：每個(gè) token 不再只有一個(gè)向量，而是擁有一組候選向量，構(gòu)成語(yǔ)義子空間；
• 上下文路由器：根據(jù) token 當(dāng)前上下文的隱狀態(tài)動(dòng)態(tài)選擇 top-K 個(gè)向量并加權(quán)融合，形成最終調(diào)制向量。

這種機(jī)制實(shí)現(xiàn)了語(yǔ)義敏感 + 稀疏激活 + 插件擴(kuò)展的三重優(yōu)勢(shì)，讓模型在幾乎不增加核心計(jì)算的情況下獲得上下文感知能力。

為了保證各向量充分參與訓(xùn)練，JTok-M 還引入了類(lèi)似 MoE 的路由負(fù)載均衡損失，使得調(diào)制空間的利用率維持高效、穩(wěn)定。

三分之一算力節(jié)省的工程落地之道

JTok / JTok-M 雖引入了大量新參數(shù)，但在系統(tǒng)設(shè)計(jì)上采取了查表式插件 + 旁路異步調(diào)度的范式，使得計(jì)算與訪存壓力都得到了有效隔離與隱藏：

• 查表可異步與主干計(jì)算重疊，訪存開(kāi)銷(xiāo)被調(diào)度隱藏；
• Token 頻率長(zhǎng)尾分布可利用，同一 token 多次查表合并訪問(wèn)，大幅減少內(nèi)存壓力；
• 訓(xùn)練階段支持嵌入并行，推理階段支持 CPU offload，僅傳輸需要的向量片段；

在這一系列優(yōu)化的加持下，即便把 JTok-M 擴(kuò)展到相當(dāng)可觀的容量，訓(xùn)練吞吐?lián)p失也不到 7%；推理階段的吞吐?lián)p失控制在 7.3% 以內(nèi)，而 GPU 側(cè)幾乎不需要額外的顯存占用 —— 兼顧了性能提升和工程落地的實(shí)用性。

重新定義 Scaling Law：讓參數(shù)擴(kuò)展走出計(jì)算陷阱

Scaling Law 是理解大模型發(fā)展的指南針：

模型的性能表現(xiàn)往往呈現(xiàn)出一種近乎物理規(guī)律般的冪律關(guān)系 —— 參數(shù)越多、數(shù)據(jù)越多、計(jì)算量越大，損失就越低，準(zhǔn)確率就越高。

但這個(gè)規(guī)律，也設(shè)下了代價(jià)等式：要繼續(xù)提升性能，就必須同步增加計(jì)算量。

JTok-M 的出現(xiàn)，打破了這個(gè)性能 - 算力綁定邏輯，建立起了對(duì)傳統(tǒng) scaling law 的一次橫向擴(kuò)展。

在論文中，為了量化 JTok-M 的擴(kuò)展?jié)摿Γ髡吆诵慕獯鹆藘蓚€(gè)關(guān)鍵問(wèn)題：

1. 當(dāng)主干模型規(guī)模擴(kuò)大時(shí)，JTok-M 帶來(lái)的收益能否保持穩(wěn)定？

2. 當(dāng)擴(kuò)展 JTok-M 本身參數(shù)規(guī)模時(shí)，它表現(xiàn)出怎樣的 scaling 行為？

結(jié)論讓人眼前一亮：

結(jié)論 1：適配超大主干，收益不打折，算力省 35%

第一個(gè)核心問(wèn)題：當(dāng)主干模型從幾千萬(wàn)參數(shù)漲到幾十億規(guī)模時(shí)，JTok-M 的性能提升會(huì)不會(huì)失效？這直接決定了它能不能落地到真正的大模型中。

作者沒(méi)有只靠實(shí)驗(yàn)碰運(yùn)氣，而是先從理論層面把邏輯說(shuō)透：將 token-indexed 參數(shù)融入經(jīng)典的 scaling law 框架，還提出了 有效參數(shù) 的關(guān)鍵假設(shè):

傳統(tǒng)模型的性能由主干激活參數(shù)（N_c）和訓(xùn)練數(shù)據(jù)量決定，而 JTok-M 新增的參數(shù)（N_n）會(huì)和主干參數(shù)形成一個(gè)擴(kuò)展比例 η=N?/N_c。再考慮到 JTok-M 的稀疏特性（embedding 本身的稀疏性，Top-K 激活的稀疏性），這些新參數(shù)會(huì)有一個(gè)有效折扣因子 γ，最終組合成有效參數(shù) N_eff。

把有效參數(shù)代入原有的 scaling law 公式后，神奇的規(guī)律出現(xiàn)了：JTok-M 沒(méi)有改變模型對(duì)算力、數(shù)據(jù)的依賴(lài)邏輯，只是讓整個(gè) 性能 - 算力 帕累托前沿曲線整體往下平移了一截 —— 這意味著，不管是小模型還是超大模型，要達(dá)到同樣的性能，JTok-M 需要的算力都會(huì)少一大截，而且這個(gè)收益是穩(wěn)定的，和主干規(guī)模無(wú)關(guān)。

后續(xù)的嚴(yán)格實(shí)驗(yàn)也完全驗(yàn)證了這個(gè)猜想：在對(duì)數(shù)坐標(biāo)下，JTok-M 的性能 - 算力前沿?cái)M合線和原生模型基線幾乎完全平行，只是整體下移了一個(gè)固定間隔。計(jì)算得出：要達(dá)到和原生 MoE 同等的模型性能，JTok-M 能直接節(jié)省 35% 的訓(xùn)練算力，這個(gè)節(jié)省比例在不同模型規(guī)模和訓(xùn)練預(yù)算下都成立，驗(yàn)證了理論中的 scale-invariance 。對(duì)企業(yè)來(lái)說(shuō)，這就是實(shí)打?qū)嵉慕当驹鲂В瑯拥念A(yù)算能訓(xùn)練出更強(qiáng)的模型，或者用更少的成本達(dá)到目標(biāo)效果。

結(jié)論 2：自身參數(shù)越堆越有用，遵循清晰的冪律縮放

第二個(gè)核心問(wèn)題：JTok-M 自己的參數(shù)擴(kuò)容后，性能會(huì)不會(huì)很快飽和？能不能成為一個(gè)獨(dú)立的縮放維度？

作者做了一組精準(zhǔn)的控制變量實(shí)驗(yàn)：固定主干模型結(jié)構(gòu)和訓(xùn)練數(shù)據(jù)量，只通過(guò)調(diào)整參數(shù)擴(kuò)展率 η 來(lái)增加 JTok-M 的自身容量。結(jié)果呈現(xiàn)出非常漂亮的規(guī)律性：在對(duì)數(shù)坐標(biāo)下，模型的驗(yàn)證損失隨著 JTok-M 參數(shù)的增加近乎線性下降 —— 每翻倍一次 JTok-M 的參數(shù)，測(cè)試損失就穩(wěn)定降低約 0.0118，而且完全沒(méi)有出現(xiàn)性能飽和的跡象.

這背后的意義非常關(guān)鍵：JTok-M 的 token-indexed 參數(shù)，本身就是一個(gè)和主干參數(shù)、數(shù)據(jù)量、專(zhuān)家稀疏度完全正交的新縮放維度。它不用依賴(lài)主干模型加參，也不用額外堆砌高質(zhì)量數(shù)據(jù)，只要單純擴(kuò)大 JTok-M 自己的參數(shù)規(guī)模，就能穩(wěn)定提升性能，而且效果可以精準(zhǔn)預(yù)判。

就像當(dāng)初稠密模型的 scaling law 一樣，開(kāi)發(fā)者現(xiàn)在可以拿著這個(gè)規(guī)律做精準(zhǔn)規(guī)劃：想讓模型損失降多少，需要給 JTok-M 加多少參數(shù)，一目了然，不用再靠盲目堆參試錯(cuò)。

這正是 scaling law 的核心魅力 —— 它不是一次偶然的性能提升，而是為大模型的發(fā)展提供了一條穩(wěn)定、可持續(xù)、低代價(jià)的新路徑，讓后續(xù)的模型迭代有章可循。

總結(jié)而言，傳統(tǒng) Scaling Law 實(shí)際上只提供了兩個(gè)方向：參數(shù)規(guī)模（N）與數(shù)據(jù)規(guī)模（D）。MoE 試圖用稀疏激活釋放計(jì)算，但仍困在 N 與 C 的線性關(guān)系中。

JTok-M 的意義在于：

• 引入了一種新的擴(kuò)展形式：token-indexed capacity；
• 構(gòu)建了從理論 → 實(shí)現(xiàn) → 驗(yàn)證 的完整閉環(huán)；
• 把 Scaling Law 從二維推向三維：參數(shù)、數(shù)據(jù)、token-indexed 結(jié)構(gòu)共塑未來(lái) LLM 的成長(zhǎng)路線。

JTok-M 不只是一個(gè)插件，更是一種對(duì)大模型如何繼續(xù)擴(kuò)展的重新定義。

下游泛化能力：知識(shí)、推理、數(shù)學(xué)能力全面飆升

光說(shuō)縮放律和算力節(jié)省還不夠，JTok/JTok-M 能不能落地，關(guān)鍵看它在實(shí)際下游任務(wù)中好不好用 —— 能不能實(shí)實(shí)在在提升模型的知識(shí)儲(chǔ)備、推理能力，甚至是代碼和數(shù)學(xué)解題水平？

作者做了超全面的驗(yàn)證：覆蓋知識(shí)、推理、代碼、數(shù)學(xué)四大類(lèi)共 14 個(gè)子任務(wù)，在 dense 和 MoE 兩種基座、多種參數(shù)規(guī)模的模型上測(cè)試，全程保持 backbone、訓(xùn)練數(shù)據(jù)、訓(xùn)練流程完全一致，只外掛 JTok/JTok-M 模塊，結(jié)果堪稱(chēng)全任務(wù)提分，而且模型越大、任務(wù)越難，收益越明顯！

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

在 1.5B 參數(shù)的 dense 模型上，加了 JTok 模塊后，14 項(xiàng)任務(wù)的平均準(zhǔn)確率直接暴漲 4.32 個(gè)百分點(diǎn) —— 相當(dāng)于在原有基礎(chǔ)上提升了近 20% 的相對(duì)收益，性價(jià)比拉滿！

其中重點(diǎn)任務(wù)的提升尤為亮眼：

• 通用知識(shí)推理（MMLU）：+4.6 個(gè)點(diǎn)，知識(shí)儲(chǔ)備更扎實(shí)；
• 科學(xué)常識(shí)推理（ARC-C）：+5.8 個(gè)點(diǎn)，對(duì)復(fù)雜常識(shí)的理解更透徹。

MoE 基座：JTok-M 發(fā)力，推理能力突飛猛進(jìn)

在 MoE 基座上，JTok 已經(jīng)能穩(wěn)定帶來(lái)平均 2.5 個(gè)點(diǎn)的提升，而升級(jí)后的 JTok-M 直接把收益拉滿，尤其是在推理、數(shù)學(xué)這類(lèi)高難度任務(wù)上，表現(xiàn)堪稱(chēng)驚艷：

• 1.5B 總參數(shù)量（激活參數(shù)量 250M）：平均準(zhǔn)確率 + 3.91；
• 3.2B 總參數(shù)量（激活參數(shù)量 500M）：平均準(zhǔn)確率 + 5.59，其中科學(xué)常識(shí)推理（ARC-C）+7.25，數(shù)學(xué)解題（GSM8K）+6.31。

這說(shuō)明 JTok-M 不僅幫模型記住了更多知識(shí)，更實(shí)實(shí)在在增強(qiáng)了模型的邏輯推理和復(fù)雜問(wèn)題解決能力 —— 這正是大模型落地的核心需求之一。

17B 大模型驗(yàn)證：越大越能打，樣本效率還超高

下游任務(wù)的收益可不是小模型專(zhuān)屬，作者在 17B 總參數(shù)量（激活參數(shù)量 2B）的超大 MoE 模型上做了驗(yàn)證，最終總參數(shù)量達(dá)到 61B，還全程跟蹤了 MMLU、ARC-C、CEval 等 6 個(gè)核心基準(zhǔn)的訓(xùn)練曲線：

曲線趨勢(shì)一眼就能看出優(yōu)勢(shì)：

• 樣本效率超高：訓(xùn)練剛進(jìn)行到幾 B token 時(shí)，加了 JTok-M 的模型就已經(jīng)全面反超原生 MoE，不用等訓(xùn)練后期，就能看到明顯提升；
• 大模型收益更猛：訓(xùn)練結(jié)束時(shí)，MMLU 提升約 4 個(gè)點(diǎn)，而 ARC-C（科學(xué)推理）、CEval/CMMLU（中文知識(shí)推理）等難度更高的任務(wù)，直接暴漲 8-9 個(gè)點(diǎn)！

核心結(jié)論：全場(chǎng)景適配，越難越能打

JTok/JTok-M 的下游泛化能力，藏著兩個(gè)關(guān)鍵亮點(diǎn)：

• 適配性極強(qiáng)：不管是 dense 還是 MoE 基座，從小模型到 17B 超大模型，都能穩(wěn)定提分，不用改動(dòng)原有訓(xùn)練流程，插件式升級(jí)就能見(jiàn)效；
• 針對(duì)性提分：不僅能提升知識(shí)記憶類(lèi)任務(wù)的表現(xiàn)，更能強(qiáng)化推理、數(shù)學(xué)等復(fù)雜任務(wù)的能力，而且任務(wù)越難、模型規(guī)模越大，提升效果越顯著 —— 這正好戳中了大模型工業(yè)化落地的核心痛點(diǎn)，讓模型在實(shí)際應(yīng)用中更能打。

相關(guān)工作對(duì)比

最近，Scaling Embedding 成為了突破大模型參數(shù)效率瓶頸的一個(gè)核心大方向。在這一探索路徑上，業(yè)界涌現(xiàn)了多篇極具啟發(fā)性的工作，包括 DeepSeek 的 Engram、Meta 的 STEM，以及我們最近推出的 JTok-M。 雖然這三篇工作都致力于通過(guò)大規(guī)模 Embedding 來(lái)低成本地?cái)U(kuò)展模型容量，但它們?cè)谠O(shè)計(jì)哲學(xué)、切入點(diǎn)以及核心洞察上各有側(cè)重，共同描繪了這一領(lǐng)域的廣闊前景：

PS：JTok，技術(shù)上代表 Joint Token。但在眾多校友心里，Joint 寓意著 “交通”，讀音也與 “交通” 撞軌；它更藏著一個(gè)交大人心照不宣的梗 ——“JT（交通）OK”。以此命名，并沿襲 Jaccount 的命名文化，向交大 130 周年校慶致敬！