K1 論文詳解

背景

Transformer 在長文本推論時,KV cache 會隨序列長度線性成長,這是服務端記憶體壓力的主要來源。近年的 linear attention 路線(例如 Qwen3.8-27B、Kimi-Linear-48B-A3B-Instruct 這類混合架構模型)改用固定大小的遞迴狀態(recurrent state)取代不斷長大的 KV cache,理論上記憶體不再隨長度增加。

但這個固定大小的狀態,在多請求並行服務時反而成了新的瓶頸。論文指出,當並行請求數達到 70 時,Qwen 的 FP32 狀態池占用的記憶體,已經超過 BF16 權重本身。每個請求都要有一份常駐狀態,請求一多,總量就很可觀。

直覺的做法是把狀態量化到低位元,但效果很差。原因在於 Delta-rule 的狀態是逐 token 遞迴更新的,每一步的量化誤差都會被帶進下一步,不斷累積。論文中的 baseline 顯示,uniform INT4 在 Qwen 的長生成推理任務上平均只剩 12.73%,而 FP32 狀態是 80.60%。

STEPQuant 的出發點是:這些誤差並不是「平均地」傷害模型,影響程度取決於兩件事。一是時間維度,誤差出現在壽命長的記憶上,會跨很多個 decoding step 持續存在。二是空間維度,不同 key row 上同樣大小的誤差,對輸出的影響不同,而且狀態數值在 row 與 column 兩個方向上的大小差異很大。

方法

STEPQuant 是一個 post-training quantization(PTQ)框架,分成兩部分。

時間維度:依記憶壽命分配位元。 狀態中的記憶壽命由遺忘門(retention gate)決定。作者量測到,壽命越長的單元累積的量化誤差越大(Spearman ρ 約 0.80)。因此 STEPQuant 在固定的位元預算下,從幾個離散的位元寬度候選中挑選,把較高的精度給「誤差大且壽命長」的單元。風險最高的那一小部分單元直接保留 FP16。論文提到,Qwen 約有 1.4% 的 head 屬於這類。

空間維度:key row 與 value column 的雙軸 scale。 作者發現狀態在 key row 方向的最大值與中位數相差約 10.3 倍,在 value column 方向則高達 19.4 倍,離群值很明顯。一般的 per-tensor 或單軸 scale 處理不好這種分佈。此外,同樣大小的誤差落在不同的 key row,對最終讀出(readout)的影響也不同。

因此 STEPQuant 為 key row 與 value column 各配一組 scale。key row 的 scale 同時考慮該 row 的數值大小和實測的「對輸出誤差的影響」,value column 的 scale 則在擬合時給高影響 row 較大的權重。兩組 scale 是聯合擬合的,而不是各自獨立決定。

這些分配與影響分數都在離線階段完成,只用 32 段 WikiText-2 做校準,推論時不需要額外訓練。

工程面上,作者把它整合進 SGLang,並寫了優化過的 GPU kernel,把量化、狀態更新等操作融合,以減少記憶體流量。

實驗結果

實驗在 SGLang 上跑 Qwen3.8-27B 與 Kimi-Linear-48B-A3B-Instruct,涵蓋兩類任務。長生成推理包含 7 項:LiveCodeBench、EvalPlus、AIME、MATH-500、HMMT、GPQA-Diamond、IFBench。短生成理解包含 6 項:MMLU、ARC-C、OpenBookQA、HellaSwag、WinoGrande、LAMBADA。

準確度(長生成任務平均):

• Qwen:FP32 狀態 80.60%;STEPQuant 6-bit 80.59%、4-bit 80.51%;uniform INT4 只有 12.73%。
• Kimi:FP32 狀態 61.52%;STEPQuant 6-bit 61.47%、4-bit 58.52%;uniform INT4 為 21.63%。uniform INT6 為 45.70%。

也就是說,6-bit 幾乎和 FP32 狀態打平。4-bit 在 Qwen 上只差約 0.1 個百分點,在 Kimi 上差約 3 個百分點。論文摘要也說,4-bit 配置的表現優於 uniform INT8。作者自己承認 Kimi 在 4-bit 的落差比 Qwen 明顯,推測與架構對量化的容忍度不同有關。

消融實驗(Qwen、6-bit): 只用空間維度的方法平均 80.19%,只用時間維度的方法平均 75.63%,兩者合併的完整 STEPQuant 為 84.51%。這組數字的基準與上面的長任務平均不同,而且完整版的分數高於 80.60% 的 FP32 長任務平均,所以應該是另一組任務或設定,不宜直接和上表比較。可以確定的結論是:兩個維度各有貢獻,合併後最好。

生成長度也是重要訊號。 uniform 量化不只是分數低,還會讓模型輸出異常拉長。例如 Kimi 的 uniform INT4 會生成約 64K token,逼近 65,536 的上限,準確度卻接近零。STEPQuant 的輸出長度則與 FP32 相近。這說明狀態誤差累積後,模型會「迷失」而不是單純變笨。

效率(NVIDIA A800):

• 6-bit 下遞迴狀態壓縮約 5.03×–5.08×。
• 含權重在內的總服務記憶體,Qwen 減少 68.7%、Kimi 減少 53.7%。
• 狀態更新操作加速約 2.91 倍。

意義

這篇論文的價值,在於把「linear attention 省了 KV cache,但狀態又成為新瓶頸」這個實際的服務問題講清楚,並給出一個可直接部署的解法。對於並行請求數高的線上服務,狀態池壓縮 5 倍以上、總記憶體最多少三分之二,代表同一張卡能多承載不少請求,或是能把剩下的空間讓給更長的上下文與更大的 batch。

從方法論看,它提醒我們:遞迴模型的量化不能沿用 Transformer 權重或 KV cache 量化的思路。在遞迴結構裡,誤差有「壽命」,位置不同(哪個 key row)影響也不同,所以要同時看時間與空間。「依誤差累積時間分配位元」這個想法,應該也能延伸到其他有狀態的架構,例如各種 state-space model 與 gated linear attention。

限制也要留意。第一,校準資料只用了 WikiText-2 的 32 段,對差異很大的領域(例如程式碼或多語言)是否穩健,論文沒有檢驗。第二,效益依賴 SGLang 整合與客製 GPU kernel,移植到其他推論框架需要額外工程。第三,極低位元時需要保留少量 FP16 的 pivot,會增加儲存與實作的複雜度。第四,誤差傳播的分析假設不同路徑的輸入與門控相同,真實 token 的變異可能偏離校準時的統計。另外,實驗只涵蓋兩個模型家族與 A800 硬體,結論能否推廣還需更多驗證。

整體來說,STEPQuant 證明了遞迴狀態可以在 6-bit 下近乎無損,在 4-bit 下仍可用。對於正在評估 linear attention 混合架構部署成本的團隊,這是一個值得參考的方向。作者已在 https://github.com/Dreamer-Toby/STEPQuant 釋出程式碼。

背景

機器人要做好操作任務,光看當下這一幀畫面通常不夠。物體正在往哪移動、任務做到哪一步、剛剛是不是已經抓過一次,這些資訊都藏在過去的影像裡。所以近來有一類 world-action model(WAM)很受關注。它先用影片模型預測未來畫面,再據此產生動作,等於讓機器人「在腦中先演練一次」。

這類方法遇到兩個互相拉扯的問題。第一,歷史拉得越長,要處理的 token 越多,推論延遲跟著上升。機器人控制是即時的,動作晚到幾百毫秒,抓動態物體就會失敗。第二,更微妙的是「模型能看到歷史」不等於「模型會用歷史」。許多系統把雙向(bidirectional)預訓練的影片模型拿來改成因果式控制器,餵給它更長的歷史,成效卻不見得變好。

Long-WAM 的核心主張就是這一點:access to history is not the same as using it。作者認為,影片基礎模型是用什麼方式預訓練的,決定了長上下文有沒有價值。他們的答案是自迴歸(autoregressive, AR)預訓練,讓模型從頭就學習「由歷史推未來」的因果結構。

方法

Long-WAM 是一套 model-system 框架,模型和系統兩邊一起設計。

1. AR 影片預訓練(LongLive2.0-Robot)。 作者先在不含動作標籤的機器人影片與第一人稱(egocentric)影片上,訓練一個因果的影片預測模型。資料來自 RoVid-X、AgiBot World、EgoDex、EgoVerse、VITRA 等五個資料集,共約 229 萬個樣本,折合約 10,000 個 window-equivalent 小時。預訓練用了 64 張 H100,約 30,720 GPU 小時。因為不需要動作標籤,可以直接吃大量便宜的影片。

2. Causal-to-causal 的 world-action 適配。 有了 AR 影片基礎後,下游訓練不把它改成雙向結構,而是保留「歷史到未來」的因果性,再接上動作生成,類似 inverse dynamics modeling(IDM)的思路。影片預測和動作生成在結構上解耦,但動作仍以預測出的未來為條件。下游訓練用 16 張 GB200,AdamW,學習率 1e-4。

3. 即時部署的系統設計。 這部分讓長上下文能實際跑在機器人上:

• Streaming VAE 編碼:觀測一到就逐 chunk 編碼,不用每次把整段歷史重新編碼一遍,省下大量重複計算。
• 非同步執行(asynchronous execution):推論和機器人動作執行重疊進行,不需要額外的 guidance 或 blending。
• 硬體專屬加速:NVFP4(4-bit 權重與激活)量化、CUDA Graph replay,加上各裝置的調校。

重點是這些加速沒有靠「砍掉未來影片預測」來換速度,每個 action chunk 仍然包含未來影片 latent 的預測。

實驗結果

上下文長度的影響(RoboCasa GR-1)。 這是全篇最關鍵的實驗。把上下文從 0.0 秒拉到 19.2 秒,Long-WAM 的成功率從 63.3% 升到 78.7%,提升 15.4 個百分點。換成雙向預訓練的初始化,同樣拉長上下文卻沒有淨增益。AR 版本領先雙向版本的幅度,在無歷史時只有 3.3 個百分點,到 19.2 秒時擴大為 17.1 個百分點。這說明差距主要來自模型有沒有能力利用歷史,而不是單純多了資料。此外,用機器人領域資料做 AR 預訓練,還能進一步提高 GR-1 與 LIBERO-Long 的峰值成功率。

標準 benchmark。

• LIBERO-Long:用 2.4 秒上下文,成功率從 94.5% 提升到 99.5%(+5.0 點),在比較的方法中最佳。
• RoboTwin 2.0:平均 94.4%,同樣在比較方法中名列最佳,Randomized 子集與最佳者持平。
• DOMINO(動態操作 benchmark):成功率 34.9%,超出比較基線 15 個百分點以上。

比較對象涵蓋 π₀、π₀.5、OpenVLA、GR00T、UniVLA、InternVLA-M1 等 VLA,以及 Motus、LingBot-VA、LingBot-VA 2.0、Fast-WAM、DreamZero 等 WAM。

延遲與硬體。 在 RTX 5090 上,每個 action chunk(含未來影片 latent 預測與完整 VAE 計算)耗時 107.4 ms,比未優化的 BF16 快約 3.3 倍。DGX Spark 為 328.2 ms,Jetson AGX Thor 為 378.7 ms。非同步執行幾乎無損,在 RoboTwin 2.0 上成功率下降不到 0.2 個百分點。

真實機器人。 在 Unitree G1 上做動態疊杯(杯子在移動),Long-WAM 成功率 95%,而 Pi0.5 與 Fast-WAM 在 20 次試驗中一次都沒成功。在 YAM 雙臂機器人上,40 秒以上的長時序任務(顏色分類、放餃子、疊碗等)平均成功率 81.7%。作者也指出,Long-WAM 可作為「帶記憶的執行器」,配合上層規劃器處理複合任務。

意義

這篇論文的價值不只是刷新幾個 benchmark,而是提出一個可驗證的設計原則:長上下文能不能帶來收益,取決於預訓練目標,而不是單純把視窗開大。對想做長時序機器人策略的團隊來說,這是很實用的提醒。如果基礎模型是雙向訓練的,硬塞歷史可能只是白白增加延遲。

第二,它展示了 world model 路線可以在邊緣裝置上落地。以前大家擔心 WAM 因為要生成未來影片而太慢,只能犧牲預測來提速。Long-WAM 在 RTX 5090 上做到 107.4 ms,在 Jetson AGX Thor 上也有 378.7 ms 且不砍掉未來預測,說明 streaming 編碼、非同步執行與低精度量化合起來,足以讓這類模型進入實機控制迴圈。

第三,動態疊杯 95% 對 0/20 的對比很有說服力。這類任務需要從歷史推斷物體速度和軌跡,正是「看過去」能發揮作用的場景,也是只看當前幀的 VLA 較難處理的。

也有一些限制值得留意。預訓練成本不低(64 張 H100、約 3 萬 GPU 小時),小團隊難以複現。DGX Spark 與 Jetson 的延遲仍在數百毫秒,高頻或極快速的任務可能不夠。DOMINO 的 34.9% 說明動態場景離解決還很遠。上述數字也都來自作者自己的設定,獨立複現仍有待觀察。總的來說,這項工作把「記憶」、「預測」與「即時性」三件事放進同一個可部署的框架,對後續的機器人基礎模型很有參考價值。

背景

影片擴散模型的主要成本來自 Diffusion Transformer(DiT)要處理的 token 數量。以 Wan2.1-14B 生成 480×832、81 幀的影片為例,video autoencoder 先把影片壓成 latent,DiT 再對 latent 做 attention。attention 的計算量隨 token 數增長得很快,所以把 autoencoder 的壓縮率拉高,是最直接的加速方法。

實務上這條路有三個麻煩:

1. 壓縮率越高,重建品質越差。要補回品質就得增加 channel 數,但 channel 變多又會讓 DiT 收斂變慢。
2. 新 latent 的分布和原本 DiT 訓練時看到的不同,所以 DiT 要從頭重訓,或花大量成本做 adaptation。
3. 即使保留原本的 autoencoder 並往下壓縮,若只用重建 loss 去最佳化,latent 還是會偏離 DiT 已經學會的分布。

換句話說,重建得好不代表生成得好。GRACE(Generation-Aware Latent Compression)的出發點就在這裡:壓縮 autoencoder 時,最佳化目標應該是「DiT 能不能繼續用」,而不只是「重建像不像」。

方法

GRACE 是兩階段框架,核心是把 latent 拆成 base 和 residual 兩部分。

階段一:訓練壓縮後的 autoencoder。

• Base latent:把輸入影片先做空間 2 倍、時間 2 倍的下採樣,再丟給凍結的原始 encoder(Wan 的 VAE),得到的 latent 保留了 DiT 熟悉的分布。
• Residual latent:額外學一個 16 channel 的 latent,負責補回下採樣時損失的空間細節與幀間動態。
• 兩者沿 channel 維度串接,共 32 channel。整體壓縮從原本的空間 f8、時間 t4,提高到空間 f16、時間 t8。搭配 patch size 2,token 數約減少 8 倍(空間 4 倍乘時間 2 倍)。

訓練目標除了一般的重建 loss(L1、LPIPS),還加了 generation-aware alignment。做法是把壓縮後的 latent 和原本的 latent 都送進凍結的 DiT,在中間層的 feature space 用 cosine similarity 對齊,各項 loss 的權重則依 gradient norm 比例自適應調整。這樣 autoencoder 學到的是「在 DiT 眼中看起來一樣」的表示。

階段二:輕量調整 DiT。

• 擴充輸入與輸出的 projection,讓 channel 數從 C 變成 C+C'。
• 對 transformer block 做 LoRA 微調,不必全量訓練。
• Asymmetric denoising:去噪時讓 base latent 比 residual 早一步收斂,偏移量為 δ。base 先穩定下來,residual 再以它為錨點補細節,全程仍是一次 forward pass,不增加額外步驟。

這個設計的好處是,原本的預訓練知識大部分被保留,新學的部分只負責補資訊。

實驗結果

訓練資料為 Panda-70M。訓練預算是 autoencoder 約 8.5 個 H200 天,加上 DiT 調整約 30 天,合計約 38.5 GPU-days,對這個規模的模型來說相當節省。

加速效果:在 Wan2.1-14B、480×832×81 的設定下,token 數減少約 8 倍,延遲從 851.5 秒降到約 76–78 秒,即 11.1 倍加速。在 736×1280×81 的更高解析度下,加速提升到 15.5 倍。

VBench 品質:

• I2V:GRACE 總分 87.90,預訓練 pipeline 為 87.92,幾乎持平。
• T2V:GRACE 為 85.81,預訓練版本為 83.93。
• 736×1280 解析度下,I2V 為 87.84(原版 87.88),T2V 為 84.74(原版 82.96)。

與基線的比較:

• 單一 latent 基線(同樣的壓縮設定,但沒有雙 latent 和 alignment)的 I2V 為 86.44,比 GRACE 低 1.46 分。
• Step-Video-VAE 重建比 GRACE 好 1.91 dB,但生成分數低 3.01 分。
• DC-Gen 同樣是把 Wan2.1 適配到 Video DC-AE,GRACE 在 I2V 上領先 7.23 分,速度也更快。論文的定性比較指出,DC-Gen 的輸出色彩偏飽和、像被過度增強,GRACE 則更接近原 pipeline 的風格。
• LTX-Video 的 I2V 總分較高,但 T2V 低 6.64 分,動作理解也較弱。

消融實驗(I2V / T2V 的增益):

• 只用雙 latent:+0.20 / +1.43。
• 再加 alignment loss:I2V 再增加 1.51,T2V 總分再增加 2.20。
• 再加 asymmetric denoising:+0.39 / +1.13。
• 對齊目標也有差異:對齊 V-JEPA 這類語意基礎模型,semantic score 為 80.04;對齊預訓練 DiT 本身的特徵,則達到 82.55。
• 若把去噪順序反過來(residual 先於 base),表現會低於同步去噪。

重建與生成脫鉤:GRACE-VAE 的重建 PSNR 為 32.63 dB,低於單一 latent 基線的 33.76 dB,但生成分數是壓縮型 autoencoder 中最高的。

意義

這篇論文最值得注意的是它證實了一個常被忽略的事:video autoencoder 的重建指標(PSNR、LPIPS)和下游生成品質並不等價。GRACE 讓重建 PSNR 略降 1 dB 左右,卻換來更好的生成結果,原因是它把「與 DiT 的相容性」直接寫進訓練目標。

在工程上,它提供了一條不必重訓大模型的加速路線。多數既有方法要從頭訓練 DiT 或做大規模 adaptation,GRACE 保留 base latent、只學 residual,再用 LoRA 微調,約 38.5 GPU-days 就讓 14B 模型達到與原版相近的 VBench 分數,同時把延遲降到約 1/11。對於要部署 Wan2.1 這類大型影片模型的團隊,這等於把單支 480p 影片的生成時間從十幾分鐘壓到一分多鐘。

這個「凍結 base、學 residual」的思路也不限於影片,同樣的邏輯可能適用於其他需要壓縮 latent、又想沿用預訓練 DiT 的場景,例如圖像或 3D 生成。

有幾點需要保留。論文沒有詳細討論失敗案例。評估主要依賴 VBench,而 VBench 本身並不能完全反映細節與動態的真實品質,尤其是在 8 倍 token 壓縮之下。T2V 分數超過預訓練版本,也有可能部分來自微調本身的影響,不全然是壓縮設計的功勞。另外,這套方法需要針對每個 DiT 各自做 alignment 與 adaptation,能否無痛搬到其他架構,還要等後續驗證。整體而言,GRACE 是影片生成加速上相當務實的一步。