K1 論文詳解

背景

大型語言模型愈來愈常被拿來輔助高風險決策,例如醫療、法律和金融。這時有一個問題變得很難迴避:單一一個答案到底能信幾分?麻煩在於,語言模型就算答錯,也能用流暢的文字和篤定的語氣包裝結論,使用者很難從表面判斷該不該採信。

學界處理這個問題的方向叫不確定性量化(uncertainty quantification, UQ),目標是替每一個預測估計可靠度。但現有方法有幾個痛點:

1. 成本高:不少方法需要對同一個問題重複生成多次,例如 Semantic Entropy 要取樣 10 次,TokUR 要跑 5 個 pass。也有方法需要另外訓練分類器或探針(probe),這類 supervised 方法常常要用有正確性標註的資料。
2. 只給一個純量:一個分數無法告訴你不確定性出現在推理的哪一步,也看不出它在推理過程中如何演變。
3. 長度混淆(length confounding):近期研究發現,生成長度本身就和不確定性估計值、答案是否正確高度相關。推理模型常常在碰到難題時寫得更長,所以一個估計器看起來很準,可能只是因為它間接抓到了「輸出很長」。

這篇論文想同時回應這三點。作者的切入點是機制可解釋性(mechanistic interpretability):既然模型內部狀態可以和人類可理解的概念對應,為什麼不直接在 hidden state 裡找出「懷疑」與「確定」的方向,把它當成讀取不確定性的鏡頭?

方法

論文提出 U-Space,一個低維度的子空間,用來讓模型「正在演變的不確定性」變得可量測、可解讀。建構流程大致如下:

第一步:選語意錨點(semantic anchors)。 作者替「懷疑」和「確定」兩端各挑一組詞,並依不確定性的性質分成四個類別:ambiguity(模糊)、incompleteness(資訊不全)、conflicting evidence(證據衝突)以及 general uncertainty(一般性不確定)。錨點詞限定為單一 token 的詞彙。

第二步:把詞彙方向映射回 residual space。 這些錨點原本活在詞彙(unembedding)空間裡。作者透過平均 Jacobian 把它們「拉回」residual stream,得到對應的內部方向。

第三步:對比並正交化。 每個類別都用「不確定那端減去確定那端」得到一條對比方向,再把四條方向正交化成一組基底,這就是 U-Space。

第四步:U-Lens 讀出。 U-Lens 把每個 token 的 hidden state 投影到這組基底上,得到 token 層級的不確定性地圖。這張圖可以直接檢視,例如看推理中哪個 token 開始出現「資訊不全」的訊號;也可以彙整成一個純量信心分數。

純量分數由兩個訊號相乘而成:

• 可言語化的不確定性:以「cone alignment」衡量 hidden state 與各類別方向的正向組合有多一致;
• 分布層面的不確定性:整段回答的平均 predictive entropy。

最終分數為 S(r) = A_cone(h_t_eot) × H̄(r)。實作上,讀取深度固定在網路約三分之二處,沒有針對任務調校。

整個流程不需要正確性標籤、不需要重複生成、也不需要訓練,只需要預先算好(或估計)Jacobian。

實驗結果

模型與資料集。 評估三個模型:Gemma 4(31B)、Qwen3.5(27B)和 Magistral 1.1(Small)。benchmark 共四個,各抽樣 1,000 題:MMLU-Pro(多領域推理)、Omni-MATH(奧林匹亞等級數學)、SuperGPQA(研究所程度專業知識)以及 TriviaQA(事實知識)。

比較對象。 單次生成方法有 MSP、predictive entropy、Self-Certainty 和 DeepConf;多次生成方法有 P(True)、TokUR(5 passes)和 Semantic Entropy(10 samples);supervised 方法則是 Feature-Gaps。

長度控制下的表現。 作者把輸出依長度切成十等分(decile bins)再評估,避免模型靠「寫得多不多」取巧。U-Lens 的 AUROC 如下:

• Gemma 4:67.0%,最強 baseline(Entropy)為 64.8%,領先 2.2 個百分點;
• Qwen3.5:66.2%,對 64.3%,領先 1.9 點;
• Magistral 1.1:70.8%,對 66.1%,領先 4.7 點。

論文表示,在 AUROC、AUPRC、AURC 三個指標、三個模型共九組比較中,U-Lens 全部領先。

長度偏誤的實證。 只用「生成長度」當作信心訊號,未校正時 AUROC 高達 68.6%,看起來頗有預測力,長度控制後掉到 52.4%,幾乎等於隨機。U-Lens 只下降 1.8 個百分點,顯示它抓到的訊號並不只是輸出長短。多次生成的方法也有類似現象:TokUR 在長度控制後掉到比單純 predictive entropy 還差,代表多花幾倍生成成本,並沒有換到更好的效果。

可解釋性與因果驗證。 作者做了兩種檢驗:

• 用內容相符的 prompt 配對來隔離各類不確定性,再用 J-Lens 解碼,各方向確實推升對應詞彙:ambiguity 方向對應 "ambiguous"、incompleteness 對應 "cannot"、conflicting evidence 對應 "controversial"、general uncertainty 對應 "dubious"。
• 在生成過程中把類別方向加到 residual state 上(steering),模型在簡單問題上也會「失去信心」。這提供了因果證據,說明這些方向確實參與塑造模型行為,而不只是統計上相關。

遷移性。 摘要指出,U-Space 比 supervised 估計器遷移得更穩。附錄 D 還把同樣的建構方式用到情緒分類(GoEmotions,四個目標 AUROC 皆超過 70%)和 reward hacking 偵測(School of Reward Hacks),這兩個任務用的是從資料推導出的錨點。

意義

這篇論文有三個值得注意的地方。

第一,它把長度混淆當成正式的評估問題。過去很多 UQ 方法的成績,可能有一部分只是在偷偷反映輸出長度。論文的長度控制評估(以及「長度本身 68.6% → 52.4%」的對照)提供了一個更誠實的比較框架。未來要宣稱某個信心估計器有效,恐怕都該交代長度控制後的表現。

第二,成本與可解讀性兼得。U-Lens 只要單次 forward pass,不用多次取樣,也不用標註資料訓練。對實際部署的團隊來說,這比 Semantic Entropy 這類方法便宜很多。token 層級的地圖也能回答「不確定從哪裡開始」,對除錯推理過程、做人工審查都有直接用處。

第三,這套「錨點 → 對比方向 → 正交基底」的做法有通用性。附錄的情緒與 reward hacking 實驗暗示,它不限於不確定性,可能成為讀取模型內部概念的一般工具。

不過限制也要看清楚:

• 實際領先幅度不大,長度控制下 AUROC 領先 1.9 到 4.7 個百分點,絕對值也只有 66% 到 71%,離可靠的決策依據還有距離;
• 只在有長推理輸出的 reasoning benchmark 上評估,短回答或其他任務型態是否成立還不清楚;
• 讀取深度固定在三分之二處、需要 Jacobian、錨點受單一 token 限制,都是工程上的前提;
• 作者自己也強調,這個讀出「不是經過校準的安全偵測器」,在高風險情境仍不應取代獨立驗證。

整體來看,U-Space 是一個把可解釋性研究拉進實用 UQ 的嘗試。它的價值與其說是分數大幅超車,不如說是提供了一個無需訓練、可以看見的視角,並迫使這個領域用更嚴謹的方式檢驗「信心分數到底在量什麼」。程式碼已公開在 https://github.com/s2labres/U-Space。

背景

長程 agent(例如瀏覽網頁、操作終端機的 LLM agent)工作時間一長,對話歷史、工具輸出和中間推理就會超過 context window。目前最常見的做法是 context compaction:在快塞滿時,請模型把歷史寫成一段文字摘要,然後丟掉原始歷史,只帶著摘要繼續。

這個做法的問題在於,摘要是在「還不知道接下來要做什麼決策」的時候寫的。摘要寫得再好,也很難涵蓋後續每一個決策需要的細節。例如某個網頁已經查過、某個指令已經試過且失敗、某個結論出自哪份來源,這些資訊一旦沒被寫進摘要就永遠消失。結果是 agent 容易重複呼叫同一個工具、重複犯同樣的錯,或是講得出結論卻說不出出處。

REMORY(Residual Memory)的出發點就是這個缺口:文字摘要不一定足以支撐後續的每一個決策。論文因此提出,在摘要旁邊另外補上一小段「神經記憶」,把摘要沒有保留的資訊以另一種形式帶過去。

方法

REMORY 的核心是一個獨立的神經記憶網路(neural memory network),它的輸出是一串有上限長度的 soft memory tokens。軟 token 不是詞表裡的文字,而是連續向量,可以直接餵給 LLM 當作輸入 embedding,因此人類無法直接閱讀,資訊要透過 LLM 才能被還原。

整體流程可以這樣理解:

1. agent 的歷史變長,需要壓縮時,照常產生一份文字摘要。
2. 記憶網路同時看到完整歷史和這份摘要,產生固定上限數量的 soft tokens。
3. 這些 soft tokens 以摘要為條件(conditioned on the summary),接在摘要後面,一起放進新的 context,讓 agent 繼續工作。

論文把這個設計類比成 residual connection,只是方向不在網路深度上,而在序列維度上。摘要是主幹,soft tokens 則是疊加在主幹上的「殘差」,專門補摘要沒涵蓋的部分,而不是取代摘要。這也解釋了為什麼 token 要以摘要為條件:它不需要重複摘要已經說過的內容,只要學著補上差額。

訓練目標也很直接:主要的 LLM 保持凍結(frozen),只訓練記憶網路,讓「摘要加 soft tokens」餵給凍結 LLM 之後產生的後續輸出,盡量逼近同一個 LLM 看到完整歷史時會產生的輸出。換句話說,完整歷史扮演老師,記憶網路學的是把完整歷史中對後續行為有影響的部分蒸餾成少量向量。因為不動 LLM 本體,這個方法可以掛在現成模型上,不需要重新微調或預訓練。

需要說明的是,我讀到的 Hugging Face 頁面和 arXiv 摘要頁只提供摘要層級的資訊,記憶網路的具體架構、soft token 數量上限、訓練資料與損失函數細節都沒有揭露,這些得看全文。

實驗結果

論文的結果分成兩條線。

SummHay。 這是一個長文件摘要的 benchmark,評分同時看兩件事:insight coverage(摘要涵蓋了多少關鍵洞見)和 source attribution(能否正確指出洞見來自哪些來源文件)。依摘要,REMORY 在 insight coverage 幾乎不變的情況下,明顯改善了 source attribution,並讓 joint score 逼近使用完整 context 的表現。關鍵數字是它只用了輸入位置(input positions)的 5.2%,也就是以大約二十分之一的 context 長度,拿到接近全量輸入的成績。這符合直覺:純文字摘要通常會丟掉來源標記這類細節,而這正是軟記憶最適合補回的資訊。

長程 agent benchmark。 論文在 Qwen3.8-27B 和 GLM-5.3-Flash 兩個模型上測試,結果顯示加入 residual memory 後,兩者在各長程 agent benchmark 上都有一致的進步。行為層面的指標更有說服力:在 BrowseComp(網頁瀏覽搜尋)和 Terminal-Bench 2.1(終端機任務)上,兩個模型的重複工具輸出(repeated tool outputs)與工具錯誤(tool errors)都大幅減少,論文用的字眼是 substantially fewer。這正好對應前面提到的失敗模式:忘記自己做過什麼,所以重做一遍。

至於各 benchmark 的具體分數、提升幅度與消融實驗,摘要沒有給出,我也沒有取得全文,所以這裡不引用任何未經確認的數字。

意義

第一,這是對 context compaction 的一個務實補強。多數 agent 框架的壓縮策略都停留在「怎麼把摘要寫得更好」,REMORY 換了思路:承認文字摘要本身有資訊瓶頸,另外開一條不受文字形式限制的通道。因為它與摘要並存、不取代摘要,即使記憶網路表現不佳,agent 也還有原本的摘要可用,風險相對可控。

第二,凍結 LLM 的設計降低了導入門檻。只訓練一個外掛的記憶網路,就能套用到不同的基礎模型。論文在兩個不同家族的模型(Qwen 與 GLM)上都看到收益,暗示這個做法不太依賴特定模型,不過是否能直接遷移到沒測過的模型,仍需要驗證。

第三,它指出一個值得追蹤的方向:agent 的「記憶」不一定要是人類可讀的文字。軟 token 難以閱讀、難以除錯,這是它的代價。如果記憶出錯,你很難像檢查摘要那樣直接看出是哪裡漏了。另一個現實限制是,記憶網路要讀完整歷史才能產生 token,壓縮當下的計算成本,以及需要存取模型內部 embedding 的要求,都意味著它比較適合自己部署開源模型的團隊,對只能呼叫封閉 API 的使用者則不一定可行。

整體而言,REMORY 的主要訊息是:在固定的 context 預算下,以約 5.2% 的輸入位置換取接近全量 context 的表現,並實際減少 agent 的重複與出錯。要判斷它是否值得導入,還得回頭看全文裡的 token 預算、訓練成本與各 benchmark 的絕對分數。

背景

Diffusion Transformer(DiT)已經是影片生成與高解析度 3D 資產生成的主流架構。它的瓶頸也很明確:self-attention 的計算量隨序列長度呈平方成長。一段 720p 影片展開成 token 序列後長度驚人,attention 往往佔掉大部分的去噪(denoising)時間。

業界的常見解法是 sparse attention:只計算「重要」的 query–key 互動,其餘跳過。問題是在高稀疏度(例如只保留 15% 的互動)時,現有方法的輸出會和 dense attention 差越來越遠,畫面細節變糊、3D 幾何變形。

這篇論文沒有直接提出新的稀疏模式,而是先問一個問題:稀疏與稠密之間的落差到底來自哪裡?作者用「oracle 對照實驗」,在相同的 attention density 下逐項替換設計,把誤差拆成三個來源:

1. 結構綁定誤差(structural binding error):多數方法以 block 為單位選取 KV,被選中的 block 會連帶保留不相關的鄰近 token(KV binding)。同一組 query 還必須共用同一份選擇結果(query binding),這會犧牲精度。
2. 選擇誤差(selection error):常見做法用 mean-pooled 表示來估計重要性,這種近似沒辦法保住累積 attention 機率,會漏掉總體貢獻很大的 token。
3. 被丟棄尾部誤差(discarded tail error):被丟掉的 token 其實仍帶有非零的 attention 權重。就算選擇完全正確,這些「尾巴」加起來仍會讓輸出和 dense 有差距。

這個拆解讓後續設計有明確目標:三個誤差各打一個,而不是只調整稀疏模式。

方法

MC-Sparse(Meta-Cached Sparse Attention)是一個 training-free 的框架,不需要重新訓練或微調模型。它的設計對應上述三個誤差來源。

Token-level KV 選擇。 不再以 block 為單位,而是跨 block 邊界挑選個別 KV token,消除 KV binding。

Tile-aligned 的 query 分組。 完全不分組、每個 query 各自選 KV,GPU 執行效率會很差。MC-Sparse 的折衷是把相似的 query 分到大小相等、與 GPU tile 尺寸對齊的群組裡。分組使用 Fast PDDP(Principal Direction Divisive Partitioning),讓組內 query 彼此相似,同時維持硬體友善的規則形狀。

用精確 attention 機率做選擇。 不靠 mean-pooled 近似,而是在特定的 anchor step 計算真實的 dense attention 機率,為每個 query 群組挑出能最大化累積 attention 質量的 KV 集合。

Metadata 快取與跨步重用。 精確選擇的成本不低,所以作者利用「相鄰 denoising step 的 attention 模式相當穩定」這個性質,只在 anchor step 計算一次,並快取三樣東西:query 分組、選中的 KV 索引,以及 dense 與 sparse 輸出之間的殘差(residual)。在後續的 reuse step,用快取的分組與索引去處理新的 Q/K/V,再把快取的殘差加回輸出,補償被丟棄的尾部貢獻,對應第三個誤差來源。

系統層面。 作者另外寫了 token-sparse attention kernel,聲稱可達理想 1/s 加速的 96–99% 效率。Fast PDDP 相比 Flash-KMeans 有 3.96–6.14× 的攤提加速。anchor step 的額外開銷約等於把 density 提高 5–6%。快取的 offloading 與 prefetching 則用來隱藏 PCIe 傳輸延遲。

實驗結果

測試模型。 影片方面有 Minimax-H3-Base(768p)、HunyuanVideo-13B(720p)、Wan2.1-14B 的 T2V 與 I2V(720p)。3D 方面是內部的 HY3D-Internal(image-to-geometry)。

比較對象。 Sol-Attn、PISA、SVG2、SVG-EAR、SpargeAttention,dense 基準為 FlashAttention-3。

評估指標。 對 dense 輸出的保真度用 PSNR、SSIM、LPIPS;3D 用 Chamfer Distance、volumetric IoU、F1;品質用 VBench(ImgQual、BgCons)與 Uni3D-I、ULIP3D-I;效率看 attention density 與 DiT 去噪加速(不含 encoder/decoder)。

影片結果(相對 dense attention):

• Minimax-H3-Base:15% density 時 PSNR 27.30 dB、加速 1.80×;25% density 時 28.44 dB、1.61×。對照的 Sol-Attn 在 32.4% density 下只有 23.66 dB、1.59×,也就是 density 更高、保真度反而低了約 4 dB。
• HunyuanVideo-13B:15% density 為 30.78 dB、2.00×;25% density 為 32.89 dB、1.82×。
• Wan2.1-14B-T2V:25% density 為 28.81 dB、SSIM 0.912、1.53×;SVG-EAR(25.3% density)為 27.61 dB、0.897、1.38×。MC-Sparse 約高 1.2 dB,並更快。
• Wan2.1-14B-I2V:25% density 為 32.11 dB、SSIM 0.929、1.52×。

3D 結果。 在 HY3D-Internal、15% density 下達到 2.32× 的去噪加速。Chamfer Distance 為 0.177,PISA 則是 0.976,在稀疏方法中幾何保真度最佳。

消融實驗。 四個元件(精確選擇加重用、token-level KV 粒度、query 分組、殘差補償)疊加後,PSNR、SSIM、LPIPS 都逐步改善。其中殘差補償的增益最大,與細粒度選擇和分組搭配時尤其明顯。在相同的實際 density 下,token-level 選擇在所有分組策略中都優於 block-level。

需要留意:論文強調的是「對 dense 輸出的保真度」與「無可見品質劣化」。加速數字是 DiT 去噪階段的加速,不是端到端生成時間。anchor step 的開銷也已被計入有效 density。

意義

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

第一是分析方法。很多 sparse attention 論文是提出一種新的稀疏模式,再用 benchmark 證明它比較好。MC-Sparse 先用 oracle 實驗把誤差拆成綁定、選擇、尾部三塊,每個設計決策都對得上一個可量測的誤差來源。這種做法讓後人能判斷哪些問題已經解決、哪些還沒有。例如消融顯示殘差補償貢獻最大,代表「丟掉的 token 其實不能直接丟」是過去被低估的因素。

第二是工程取捨。精確的 token 級選擇本來太貴,無法實用。作者靠「跨 denoising step 的穩定性」把成本攤提掉,再用 tile-aligned 分組換回 GPU 效率。這個思路和 diffusion 領域常見的 feature caching 是同一脈絡,只是被套用到 attention 的選擇結果與殘差上。對已經部署大型影片模型的團隊來說,training-free 意味著可以直接套在現有 checkpoint 上,不必重新訓練。

限制也很清楚。1.5×–2.3× 的加速比起「理論上 15% density 的 6 倍以上」仍有距離,因為還有 anchor 開銷、非 attention 的計算與記憶體搬運。快取殘差與索引需要額外顯存或 offload 機制。論文也依賴「相鄰步驟 attention 穩定」的假設,若換成步數極少的蒸餾模型(few-step distillation),這個假設是否成立,論文摘要與頁面中並未說明。另外,3D 實驗使用的是內部模型,外部較難重現。

整體而言,MC-Sparse 更像是一份診斷加處方的系統性工作。如果後續有開源實作,它很可能成為影片與 3D DiT 推論加速的實用基準之一。