強大的大型語言模型（LLM）現在已能透過 Gemini Nano 和 Gemma 等裝置端模型實現，讓 AI 隨身攜帶成為現實。這項技術讓手機上的日常功能得以實現，例如即時摘要大量通知或校對重要簡訊，且所有處理都在裝置上完成，無需將您的私人資料傳送出去。然而，要讓這些功能對日常使用者真正有用，它們必須非常有效率地運作。

在行動裝置上實現這種速度是一項重大挑戰。不同於龐大的伺服器環境，手機在嚴格的能源預算和硬體記憶體（RAM）限制下運作。此外，標準語言模型以「自迴歸」方式生成文字，這表示它們一次只處理並輸出一個詞（或詞元）。這種循序漸進的過程會產生瓶頸，未能充分利用手機的處理能力，同時也耗費記憶體頻寬，最終可能導致使用者體驗變慢並消耗電池電量。

為了解決這個瓶頸，我們宣布推出一種新架構，將多詞元預測（MTP）技術應用於現有的「凍結」Gemini Nano v3 模型。借鑒 EAGLE 框架和 Confident Adaptive Language Modeling (CALM) 等先前方法，我們設計了新的架構組件，以最大化這些專為行動環境設計的效率提升。我們最近的公告強調了使用 MTP 加速 Gemma 4 並將其提供給開發者。

今天的文章探討了邊緣運算獨特且極端的限制。這項方法已近期推廣至 Pixel 9 和 10 系列，提供開箱即用的速度提升。對於使用者而言，這表示 AI 通知摘要和校對等功能生成文字的速度顯著加快，且能耗更低。對於開發者而言，它消除了主要的痛點：無需為每個新任務微調獨立且佔用大量記憶體的草稿模型，即可提供高速的裝置端 AI。

MTP 建立在推測解碼（speculative decoding）的演進基礎上。在傳統設定中，生成 N 個詞元需要大型模型進行 N 次前向傳播。推測解碼將此過程分為兩部分：草稿（Draft）：一個較小、較快的近似模型（「草稿器」）生成一小段候選詞元序列（例如 3 個詞元）。

驗證（Verify）：一個大型模型（「驗證器」）並行處理這些候選詞元。如果候選詞元與大型模型預測的結果相符，則接受它們。如果不符，系統將回溯到第一個分歧點。

然而，這會導致一些效率低下的問題。運行一個獨立的「獨立」草稿器模型（例如 128M 參數）會與有限的 RAM 競爭。此外，獨立草稿器對主模型的豐富內部狀態「一無所知」，僅根據文字歷史預測下一個詞元，而沒有主模型已計算出的語義上下文。MTP 透過從獨立架構轉向整合架構來解決這些效率低下的問題。

我們不是訓練一個單獨的小型語言模型來生成草稿詞元，而是在主模型的最終層附加一個輕量級的 Transformer 頭，即 MTP 頭。

這種使用深度退出層進行草稿的架構，利用了主模型骨幹已經完成的工作。MTP 頭接收主模型最終的高維激活（隱藏狀態），並利用它們自迴歸地預測一系列未來的詞元。

雖然 MTP 頭通常與骨幹模型一起預訓練（例如我們最近發布的 Gemma 4 模型），但在利用已部署的裝置端基礎模型時，這種做法是不可行的。相反地，我們的工作重點是改造草稿器頭，使其獨立於預訓練流程運作。

我們取一個經過完整訓練的 Gemini Nano v3 模型，凍結其權重，並在最終層附加一個密集的 Transformer 堆疊，即 MTP 頭。我們僅訓練這些參數，以最小化未來詞元的預測誤差。由於骨幹模型是凍結的，MTP 純粹是一種效率優化，確保基礎模型的能力或安全對齊不會下降。

由於不正確的草稿在驗證過程中會被捨棄，最終輸出與主模型保持位元級別的完全一致，這使我們能夠推出具有完整向後相容性的效率更新。

雖然標準的 MTP 實作透過在主模型和草稿器之間共享靜態參數（如嵌入權重）來優化訓練效率，但裝置端推論面臨更嚴格的瓶頸：動態記憶體。即使共享權重，如果草稿器獨立處理上下文，它會透過生成和維護自己的鍵值（KV）快取，對記憶體造成「雙重負擔」。考慮到行動裝置上有限的記憶體，避免這種冗餘至關重要。

為了解決這個問題，我們設計了一種零複製架構，其中 MTP 頭有效地利用了主模型的狀態。MTP 頭不是維護自己的歷史記錄，而是設計成直接交叉注意力到主模型的凍結 KV 快取。這使得草稿器能夠查詢骨幹模型已經計算出的「記憶」和上下文，而無需重複。

這種設計帶來了兩項效率提升。首先，它消除了草稿器預填充延遲：透過利用現有快取，MTP 頭無需額外時間處理提示詞。其次，它減少了運行時的記憶體佔用。與獨立草稿器相比，我們觀察到每個實例節省了 130MB 的記憶體，這是透過節省草稿器嵌入查找表、預填充點注意力變體和應用程式特定調整參數實現的。

在我們的實驗中，我們發現 MTP 草稿器始終能產生更準確的詞元預測，這使得 Pixel 9 裝置上的速度提升達到 50% 或更多，具體取決於任務，相較於參數數量相當的「獨立草稿器」。

這種性能差距源於 MTP 能夠存取更豐富的表示。與將主模型視為黑盒的獨立草稿器不同，MTP 頭直接利用大型骨幹模型已處理的最終激活：指令遵循：在摘要或具有複雜限制的重寫等任務中，MTP 顯著優於獨立微調的草稿器。可預測的文字結構：對於具有高度結構可預測性的任務（例如智慧回覆），MTP 頭有效地學習了主模型的語法模式，詞元接受率提高了高達 55%。

為了在 Pixel 9 和 10 裝置上部署 MTP，我們重新設計了裝置端推論堆疊，以處理驗證和草稿階段之間複雜的依賴關係。

結果驗證了架構選擇的正確性。在生產工作負載中，例如 AI 通知摘要和校對，MTP 在每次推論傳遞中平均能正確預測近兩個額外的詞元。此外，更少的驗證步驟意味著喚醒重型處理器的時間更少，從而降低能耗並延長電池續航力。

我們期待將 MTP 整合到未來的 Pixel 裝置中，並探索替代架構，包括並行解碼和無需輔助頭的範式，以在嚴格的行動裝置限制下進一步降低草稿延遲並增加同時詞元驗證。我們也在研究如何更有效地處理語言生成固有的模糊性。雖然標準的推測解碼假設單一的最佳未來路徑，但我們正在開發允許模型並行探索分支可能性的技術。

這旨在即使在不確定的情況下也能最大化接受長序列的可能性。此外，我們正在研究驗證寬容度：針對特定用例放寬草稿和驗證之間嚴格的精確詞元匹配，以進一步提高邊緣運算的效率。

這項工作是我們優化裝置端 LLM 效率努力的一部分，由 Filippo Galgani、Omri Homburger、Pooja Consul、Matthew Markwell 和 Vivek Kumar 共同完成。某些元素建立在 Google DeepMind 的 Gemini 團隊的開發成果之上：Tal Schuster、Ziwei ji、Ivan Korotkov 和 Ganesh Jawahar。

我們還要向 Nadav Bar、Utku Evci、Nir Shabat、Joe Zou 以及 Google Research、Google Deepmind 和 Platforms & Devices 的團隊致以誠摯的感謝，感謝他們的審查、寶貴的回饋和支持。