K1 論文詳解

背景

大型語言模型在處理 context 裡的「引用鏈」時,有個容易被忽略的限制。引用鏈是指程式第 3 行引用第 2 行的變數,第 2 行再引用第 1 行,模型得一路追到底才答得出來。這類任務本質上是多跳推理:每多一跳,模型就要多做一輪「找到被引用的那一行,再取得它的值」。

這篇論文量測了這種能力的上限。作者測了 13 個預訓練 base model,發現它們在不使用 chain-of-thought 的情況下,只能穩定追蹤約 1.4 到 3.6 行的引用鏈。這個數字遠低於模型的層數:一個幾十層的 transformer,理論上有很大的 forward pass 深度可以拿來做串接運算,實際上卻只用了一小部分。標題「Stop Thinking Too Early」說的就是這件事:模型在還有很多層可用時,就已經停止了這類運算。

另一個直覺是,既然深度不夠,就增加深度。論文也檢驗了這條路。像 Ouro 這種 looped(迴圈式)語言模型,會重複使用同一組層來增加有效深度。作者發現,預訓練階段就加上去的額外 loop 帶來的好處很有限。所以問題不只是層數或迴圈次數不夠,而是模型沒有學會把多出來的深度用在這種串接運算上。

這也呼應了一個常見的困擾:模型的「預設輸出」可能低估了它內部真正具備的運算能力。這篇論文想問的是,能不能用極小的改動,把這部分被閒置的深度喚醒?

方法

核心做法很簡單:凍結模型的所有權重,只在某一個早期層放一個 rank-8 的 LoRA,並用這個引用鏈任務去訓練它。LoRA 的可訓練參數量相對於 8B 模型可以忽略,所以論文才稱它是 "tiny edit"。

論文接著回答這個 LoRA 到底做了什麼,並提出一個「relay(接力)」機制:

• LoRA 在早期層啟動接力,讓每一行程式把自己所屬的「鏈身分」(chain identity)寫進表徵。
• 這份身分資訊會在一小段中間層裡逐層往後傳遞,像接力賽一樣。
• 模型中原本凍結的 attention heads 會讀取這份資訊,而且越往後的層,讀到的鏈上位置越遠。也就是說,每多過幾層,模型就能多追蹤幾行。
• 作者用干預實驗驗證因果關係:移除對「父行」(parent line,也就是被引用的那一行)的 attention,接力就會中斷。

論文還提出一個不需要訓練的診斷方法。在凍結模型上直接量測,可以估計「最後一個還有用的介入層」在哪裡,也就是最晚能在哪一層放 LoRA 而不失效。在四個 held-out 模型中,這個量測有三個落在容許誤差內。這表示介入位置不必全靠試錯,可以從模型本身的性質預測。

除了合成的引用鏈任務,作者也訓練了 task-specific 的 LoRA,在真實的多跳問答 benchmark MuSiQue 上測試。

實驗結果

論文摘要給出的關鍵數字如下:

• 基線:13 個 base model 只能可靠追蹤 1.4 到 3.6 行。預訓練時額外加的 loop 幾乎沒有幫助。
• Qwen3-8B:在 24 行的引用鏈上,exact accuracy 從 15.5% 提升到 99%。這是用 rank-8、單一早期層的 LoRA 做到的,模型本身的權重全部凍結。
• 更長的訓練:訓練時間拉長後,同一類 LoRA 可以讓 Qwen3-8B 做到 50 行。
• Ouro-1.4B:這個 looped 模型在 4 次 loop 後可到 60 行,8 次 loop 後至少 160 行。這說明 loop 能提供額外深度,但前提是模型學會利用它。這和「預訓練 loop 沒什麼用」的說法並不矛盾:有沒有用,取決於模型有沒有被教會把深度用在串接運算上。
• 機制驗證:移除 parent-line attention 會讓接力停止,這是接力假說的因果證據。
• 預測介入層:在凍結模型上量測最後有用的介入層,四個 held-out 模型中有三個命中。
• MuSiQue:task-specific LoRA 也能改善 MuSiQue 的表現。我手上的來源只有摘要層級的描述,沒有提供具體分數與提升幅度,所以這裡不列數字。

需要說明的是,我讀到的 arXiv 頁面與 Hugging Face 頁面幾乎只有摘要內容,沒有表格、消融細節或訓練設定(例如 LoRA 放在第幾層、訓練資料規模、不同鏈長的完整曲線)。因此上面的數字都來自摘要,細節請以論文全文與作者的互動 demo 為準。

意義

這篇論文最有價值的地方,是重新界定了「模型做不到」的意思。15.5% 到 99% 的跳升,不是靠擴大模型或加入新知識換來的,而是靠一個極小的改動,讓模型本來就有的 attention heads 能被串接起來使用。這表示預設輸出會低估模型內部能取用的運算量。評估模型的推理能力時,只看 base model 的零樣本表現,可能得到偏悲觀的結論。

對實務有幾個啟發。第一,如果某類任務卡在「需要的運算深度」而不是「缺乏知識」,那麼用小 LoRA 微調的 ROI 可能很高。第二,relay 機制給了可解釋性研究一個具體的案例:多步推理在 transformer 內部可以被拆解成逐層傳遞的資訊流,而且能用 attention 干預去驗證。第三,「在凍結模型上預測最後有用的介入層」這個工具,讓我們不必暴力搜尋放 LoRA 的位置。

不過有幾點需要保留。訓練任務是合成的引用鏈,和真實世界的推理有差距。MuSiQue 上有改善,但摘要沒有說明幅度,是否能廣泛遷移還要看全文。預測介入層的方法也只在四個模型中命中三個,還不是萬用的。另外,這種 LoRA 是 task-specific 的,每換一類任務可能都要重新訓練。

整體來看,這是一篇把「深度利用不足」這個模糊直覺,變成可量測、可干預、可解釋的研究。對關心 looped transformer、test-time compute 與機制可解釋性的讀者,值得追一下作者放出的程式碼與互動 demo。

背景

Masked diffusion models(MDMs)是近年很受關注的非自迴歸文字生成路線。它從一段全被 mask 的序列開始,每個 denoising step 同時還原多個 token。autoregressive 模型一次只能產生一個 token,MDM 理論上可以靠「一步多個 token」換取速度。

這個速度優勢有一個前提。標準 MDM 的 reverse process 是**對位置因子化(factorized)**的:同一步裡,各個位置的 token 被視為條件獨立,各自從自己的分布取樣。如果一步只還原一個 token,這個假設沒有影響。一步還原很多個 token 時,位置之間的相關性就被丟掉了。

舉個直觀的例子。假設某個位置可能是「紐約」或「倫敦」,相鄰位置可能是「Times Square」或「Big Ben」。因子化的取樣可能分別抽出「紐約」和「Big Ben」,拼出不連貫的組合。步數越少,同一步要同時決定的 token 越多,這種不一致越嚴重。而少步生成正是 diffusion 相對 autoregressive 最有賣點的區間,所以這是 MDM 的核心瓶頸。

已有的做法是引入連續的高斯潛變數,用 variational autoencoder(VAE)的方式訓練,讓潛變數攜帶跨位置的相關資訊。但這類方法容易遇到 posterior collapse:解碼器發現不用潛變數也能把 loss 壓得不錯,就悄悄忽略它。潛變數形同虛設,模型退化回因子化的版本。

方法

E-MoE(Enhanced Mixture-of-Experts)不另外加一個要訓練的連續潛變數,而是把 MoE backbone 內建的專家路由決策當作離散的共享潛變數。

核心想法如下:

1. 離散共享潛變數:MoE 的每一層會決定每個 token 交給哪些 expert 處理。把這些路由選擇合起來,就得到一個離散的「配置」,可以視為整個序列共享的潛變數 z。
2. 混合分布:給定 z,各位置的條件分布仍然是因子化的,所以每個 z 對應一個因子化分布。對 z 做邊際化,整個 reverse process 就變成多個因子化分布的混合(mixture of factorized distributions)。單一因子化分布無法表達位置間的相關性,混合分布可以。典型例子是雙峰分布:「全是紐約」或「全是倫敦」,各自內部一致,混在一起後就有了相關性。
3. 不增加 active parameters:MoE 本來就是稀疏啟動,每個 token 只走 top-k 個 expert。E-MoE 沒有加額外的 latent encoder 或解碼分支,是重新詮釋並利用既有的路由機制。因此與因子化 baseline 相比,每次 forward 的 active parameters 不變。

與連續高斯潛變數相比,這個設計有兩個可能的好處。第一,路由決策本來就決定了計算路徑,模型不能忽略它,所以不容易發生潛變數被默默忽略的 posterior collapse。第二,離散潛變數天然適合對應離散的多模態結構。

需要說明的是,我能取得的資料只有摘要和 arXiv 頁面的概述。路由如何被訓練成有意義的潛變數、訓練目標的具體形式、是否需要額外的 regularization(例如負載平衡損失與潛變數資訊量之間的取捨),這些細節我沒有取得,不在此推測。

實驗結果

論文在三類任務上評估:

• 合成的多模態 benchmark:這類資料有明確的多峰結構,正是因子化模型最容易失敗的情境,適合直接檢驗模型有沒有捕捉到位置間相關性。
• Binarized MNIST:影像生成任務,像素之間高度相關。
• LM1B:One Billion Word Benchmark,是離散擴散語言模型常用的語言建模資料集。

摘要的結論是:在這三類任務上,E-MoE 在少步生成時都優於因子化 baseline。比較對象是「active parameters 相同的因子化 baseline」,所以增益並非來自單純加大模型。

具體數字(例如 LM1B 上的 perplexity 或生成品質指標、各步數下的差距、模型參數量、expert 數量)我沒能從可取得的頁面讀到。Hugging Face 頁面抓取時回傳的內容與這篇論文無關,arXiv 頁面也只有摘要等級的資訊。我不想編造數字,所以這裡不列出。需要精確的 benchmark 對照,請直接看論文全文的表格。

另外要留意,摘要強調的是少步區間的改善。步數增多時,因子化假設的傷害本來就會縮小,E-MoE 的優勢可能也會縮小。這是從問題性質推論的,不是論文的結論。

意義

這篇論文有三點值得注意。

第一,它直接針對 diffusion LM 的主要弱點。 如果 MDM 要在實際部署上贏過 autoregressive,少步、高並行度的生成品質是關鍵。E-MoE 不靠增加步數,而是提高每一步的表達能力,方向與這個目標一致。

第二,它把 MoE 的路由機制重新詮釋成潛變數。 這個視角很有意思。MoE 在業界已經是主流架構,如果路由本身就能擔任「捕捉全域結構」的角色,就不必再引入一套會 posterior collapse 的 VAE 機制。工程上也容易採用,因為只是在既有 MoE 骨架上改變訓練與取樣的方式,不改變推論成本的量級。

第三,它對 posterior collapse 提出了結構性的解法。 過去對這個問題多半靠 KL annealing、free bits 之類的訓練技巧。E-MoE 的思路是讓潛變數成為計算圖中不可繞過的部分。

也有一些尚待釐清的問題:

• 實驗規模(LM1B、binarized MNIST、合成資料)相對小,在更大規模、更長序列,或接近現代 LLM 的設定下是否成立,還需要驗證。
• 路由決策被拿來承載資訊後,和 MoE 原本的負載平衡目標可能互相牽制,摘要沒有討論。
• 離散潛變數的組合空間很大,要用多少 expert、多少層的路由才足以表達位置間相關性,取捨尚不清楚。

整體而言,E-MoE 提供了一個低成本、概念簡潔的方向:不增加 active parameters,就讓 diffusion LM 在少步生成時不再假裝各位置互相獨立。對關注 diffusion 語言模型與 MoE 交會處的讀者,值得看全文,確認實際的增益幅度與適用範圍。

背景

用強化學習(RL)訓練大型語言模型時,目標常常不只一個。例如在工具呼叫任務中,模型要輸出正確的工具與參數,還要符合指定的輸出格式。在數學推理任務中,模型要答對,同時要遵守長度限制。實務上的做法是替每個目標設計一個獨立的 reward function,再把多個 reward 合併成訓練訊號。

這種多獎勵設定有個老問題:各個 reward 的數值尺度和分布差異很大。直接加總時,變異大的 reward 會壓過其他 reward。GDPO 這類方法的做法是 reward-wise normalization:先在每個 rollout group 內,針對每個 reward 各自做標準化,再聚合成 advantage。這樣每個 reward 的「組內相對優劣」都被保留下來,不會被其他 reward 的尺度淹沒。

不過這篇論文指出,這樣做之後學習進度仍然不均。不同目標被優化的速度還是不同,有些行為(例如格式合規)要等很久才學會。原因是 normalization 只處理了「組內」的尺度,沒有處理「批次層級」的訊號量差異。

方法

作者先提出一個診斷量,叫 advantage energy:某個 reward 在一個 batch 內,所有 advantage 的平方和。它可以粗略理解為「這個 reward 在這個 batch 裡貢獻了多少梯度訊號」。

接著是論文的理論觀察。在理想化的 GDPO normalization 下,advantage energy 正比於 active-group density,也就是「該 reward 能提供非零相對 advantage 的 rollout group 所占的比例」。

直覺如下:

• 如果一個 group 裡所有 rollout 在某個 reward 上得分相同(例如全部都符合格式,或全部都不符合),組內標準化後 advantage 全為零,這個 group 對該 reward 沒有貢獻。
• 只有 group 內出現「有的好、有的壞」的差異時,該 reward 才有非零 advantage。標準化後,每個這類 group 貢獻的平方和大致固定。
• 所以整個 batch 的 energy 大致就由「有多少比例的 group 是 active」決定。

結論是,稀疏的 reward(多數 group 裡每個 rollout 得分都一樣)即使經過 reward-wise normalization,在整個 batch 中的總訊號量仍遠小於常常 active 的 reward。這就是標題所說的 residual batch-level signal imbalance。

由此得出的校正方式是 Density-Aware Reward Aggregation(DARA)。energy 隨權重的平方放大,所以要把各 reward 的 energy 拉平,權重就該和密度的平方根成反比。論文因此推導出 inverse-square-root density correction:密度越低的 reward 權重越大。這個推導是我依據摘要推演的,細節以論文為準。

DARA 還有幾個設計特點:

• 權重逐 batch 計算:每個 rollout batch 都重新估計各 reward 的 active density,所以訓練中 reward 活躍度改變(例如格式學會後變得很少有差異)時,權重會自動調整。
• 不改動底層優化目標:DARA 只改變 reward 聚合的方式,不動 policy optimization 的 loss。它可以套在 GDPO 這類流程上,不必另外調整一組 reward 權重超參數。

實驗結果

摘要列出兩類任務:tool calling 與 mathematical reasoning,並以 GDPO 為主要比較對象。我能取得的頁面內容有限,模型大小與完整的 benchmark 數字無法確認,以下只列摘要明確給出的數字。

• 工具呼叫:DARA 達到高 format compliance 所需的訓練步數,最多比 GDPO 少 26%。
• 數學推理:DARA 達到接近飽和的 length compliance,所需步數最多比 GDPO 少 65%。
• 最終表現:兩類任務上,DARA 的最終表現與 GDPO 維持相當(competitive)。

這裡的「最多」(up to)代表這是最佳情況的收斂加速,不是所有設定下的平均值。這些加速主要發生在格式與長度這類「稀疏但明確」的行為目標上,看不出最終準確率有明顯提升。換句話說,DARA 的價值在於學得更快,不是學得更好。論文也附有程式碼(github.com/zhaihaotian/DARA),想看具體實驗設定的人可以直接查。

意義

這項工作的價值有三點。

第一,它給了多獎勵 RL 一個可以量測的診斷工具。過去常憑經驗調 reward 權重,advantage energy 則讓人能在訓練中直接檢查,是否有某個 reward 因為太少 group 有差異,而在梯度裡幾乎沒聲音。

第二,修正方式簡單且有理論依據。權重由 batch 內的 active density 直接算出,不需要額外的網格搜尋,也不需改動 RL 演算法本身,工程整合成本低。對正在用 GDPO 或類似流程訓練 agent 的團隊來說,這是容易試的改動。

第三,它對應到實際訓練成本。RL 訓練很貴,若格式合規、長度控制這類「基本功」能少花 26% 到 65% 的步數,同樣預算就能把更多步數留給真正難的能力。

也有幾點需要保留:

• 結論主要建立在「理想化 GDPO normalization」的假設上,真實訓練中標準差為零、group 大小有限、reward 分布不理想時,近似程度有多好,需要看論文的完整分析。
• 加權依據是 active density,所以它擅長補強「稀疏」的 reward。若稀疏的 reward 本身雜訊很大,放大它的權重也可能放大雜訊。
• 實驗只涵蓋兩類任務,是否能推廣到更多 reward 數量、更大模型或其他領域(如對話安全、程式碼),還有待後續驗證。

整體來看,DARA 是一個針對性強的小改動:它指出了多獎勵 RL 中被忽略的「訊號稀疏度」問題,並用一個簡潔的 1/√density 修正去處理,對重視訓練效率的多目標對齊工作值得關注。