Compressing Sequences in the Latent Embedding Space: $K$-Token Merging for Large Language Models
AWS 團隊提出 K-Token Merging 技術，透過嵌入空間壓縮讓模型輸入長度減少 75%，最高削減 94% 運算量。
- 針對潛在嵌入空間進行壓縮，取代傳統的字詞刪除，保留完整上下文資訊。
- 透過輕量級編碼器將連續 K 個詞元融合為單一向量，迴避擴充詞表造成的記憶體耗損。
- 當 K=4 時可減少 75% 輸入長度，在短文本生成情境下，能降低高達 94% 的預填充運算量。
當我們輸入長篇提示詞給大型語言模型時，運算成本正以平方級數飆升，但最新的研究顯示，透過「潛在嵌入空間」的壓縮技術，模型能在幾乎不掉效能的情況下減少 75% 的輸入長度。這項由 AWS AI Labs、Mistral AI 與 UCLA 等機構共同提出的 $K$-Token Merging 技術，打破了過去只能刪減字詞的限制，在特定情境下甚至能削減高達 94% 的預填充運算量，為長文本處理帶來全新架構解法。
Qwen 2.5 詞表解析：挖掘二萬八千位元的冗餘
大型語言模型（LLM，Large Language Model）的應用已從基礎客服延伸至複雜的軟體生成與代理系統，這使得輸入的上下文長度不斷增加。然而，在標準的 Softmax 注意力機制下，記憶體與運算需求會隨著輸入長度呈現二次方增長。為了解決這個難題，業界長期投入「提示詞壓縮（Prompt compression）」，但多數方法都專注於減少 token 數量，忽略了潛在嵌入空間（Latent embedding space）中本身的巨大冗餘。
以擁有 151,936 個詞彙量的 Qwen 2.5 模型為例，每個 token 大約消耗 28,672 位元（896 維度乘以 32 位元）。從資訊理論的角度來看，要從詞表中識別出一個特定 token，理論上只需要大約 18 位元。即使是連續的 $K$ 個 token（$K$-gram），理論下限也僅約為 $18K$ 位元。這種理論下限與實際消耗之間的巨大落差，明確指出了目前的嵌入表示法存在極大浪費，同時意味著在特徵維度上有著極高的壓縮潛力。
K-Token 融合：迴避 29 倍擴充記憶體災難
為了解決上述的低效問題，研究團隊提出了 $K$-Token Merging 框架。這是一種全新的軟性壓縮方法，其核心概念是利用一個輕量級編碼器（Encoder），將連續的 $K$ 個 token 嵌入向量，壓縮成「單一個」嵌入向量。這意味著在模型的輸入端，多個 token 可以共用一個濃縮後的表徵，接著再透過 LoRA（Low-Rank Adaptation，低秩適應微調技術）讓模型適應這些壓縮後的輸入。
這種作法與傳統的「擴充詞表」有著本質上的差異。如果嘗試把所有常見的連續字詞組合直接加入詞表，會面臨長尾 token 爆炸的問題。研究指出，所有不重複的 4-gram 數量可能會高達原始詞表的 29 倍，導致儲存成本高昂且難以最佳化。更為致命的是，固定詞表無法泛化到未曾見過的 token 組合。相反地，$K$-Token Merging 透過編碼器即時運算，完全避免了詞彙量暴增的記憶體危機，且能自然處理任何未知的文本組合。
跳脫硬性刪減瓶頸：對比 LLMLingua2 差異
過去的 token 壓縮技術主要分為兩大陣營。第一類是「硬性壓縮（Hard prompt compression）」，例如 SelectiveContext 透過自我資訊量來刪除多餘字詞，或是 LLMLingua2 訓練分類器來標記並捨棄冗餘 token。這類方法雖然直觀，但若遇到數學推理、程式碼翻譯等高密度資訊任務，只要漏掉一個關鍵字就會導致模型生成錯誤，造成效能大幅衰退。
第二類則是「軟性壓縮（Soft prompt compression）」，例如 Gist 將指令壓縮成少量的前置中繼 token，並強制限制注意力機制的運作；或是 LTSC 找出常出現的子序列並替換為學習過的中繼 token。雖然表現較佳，但這些方法仍然侷限在「Token 空間」內操作。$K$-Token Merging 則是直接在「潛在嵌入空間」進行降維，不僅沒有捨棄任何一個原始 token，更將它們重新編碼成高密度的壓縮向量，在文本結構化推理、情感分類與程式碼編輯等實驗中，展現出極佳的效能保留能力。
預填充運算量減 94%：平均初始化編碼器細節
在具體架構設計上，模型運作分為預填充（Prefill）與生成（Generation）兩個階段。在預填充階段，輕量級的編碼器（約僅佔 50 MB 的小型多層感知機）會讀取每 $K$ 個 token 的原始嵌入，並輸出一個壓縮版本。到了生成階段，模型則一律使用原始未壓縮的詞表進行自迴歸輸出，確保生成的文本品質不受影響。訓練時的目標函數也僅針對這些未壓縮的生成 token 計算負對數概似損失，確保模型能正確解讀壓縮向量。
實務上為了加速訓練收斂，團隊特別設計了 Average-Initialized Encoder（平均初始化編碼器）。該編碼器在訓練初期，輸出的壓縮向量會非常接近這 $K$ 個原始向量的平均值，為最佳化過程提供一個穩定的語義起點。實驗證實，當設定 $K=4$ 時，輸入長度縮減了 75%。由於注意力機制的運算量與長度平方成正比，壓縮後的預填充運算量僅剩原本的 6.25%。在輸出長度遠小於輸入長度的情境下，整體運算量等於降低了驚人的 94%，為大型語言模型邁向超長文本處理提供了極具實用價值的底層技術。
不依賴容易遺漏資訊的硬性刪字，而是直接在嵌入空間將多個連續詞元降維成單一向量，為 LLM 處理超長輸入帶來了兼顧算力與精準度的新思維。