在Pixel上透過凍結的多token預測加速Gemini Nano模型

Google研究人員近日宣佈，他們成功將多token預測（Multi-Token Prediction，簡稱MTP）技術應用於已部署的Gemini Nano v3模型，顯著提升了Pixel裝置上AI功能的執行效率。這一成果使得諸如通知摘要和文本校對等基於AI的功能在移動裝置上執行得更快、更省電，同時無需犧牲使用者隱私或模型準確性。

移動裝置上的大語言模型（LLM）面臨著嚴格的能源和記憶體限制。傳統的自迴歸生成模式每次只生成一個token，導致處理瓶頸。而MTP透過讓模型一次預測多個未來token，將生成過程分解為草稿和驗證兩個階段，從而加速推理。然而，傳統的單獨草稿模型會佔用寶貴的記憶體，且無法利用主模型的計算狀態。

為了解決這些問題，研究團隊設計了一種全新的架構：在不修改已訓練好的Gemini Nano v3模型權重（即凍結主幹網路）的前提下，在模型頂層附加一個輕量級的Transformer頭部（MTP頭部）。這個頭部專門訓練來預測未來token，而主幹網路保持不變。由於凍結的主幹網路，這種方法純粹是一種效率最佳化，不會影響模型原有的能力和安全性。

更關鍵的是，他們實現了“零複製”架構。MTP頭部不維護自己的鍵值快取，而是直接交叉關注到主模型的凍結KV快取中。這一設計避免了重複計算，消除了草稿階段的預填充延遲，並將每個例項的記憶體佔用降低了高達130MB。實驗表明，在Pixel 9裝置上，相比引數數量相當的獨立草稿模型，MTP頭部在不同任務中實現了50%以上的速度提升。特別是在指令遵循任務（如複雜約束下的摘要或重寫）和結構可預測的任務（如智慧回覆）中，MTP的token接受率提升了55%。

目前，這一技術已在Pixel 9和Pixel 10系列上部署。在實際應用如AI通知摘要和校對中，MTP每次推理平均能正確預測近兩個額外token，從而減少了驗證步驟，降低了處理器喚醒時間，改善了電池續航。未來，研究團隊計劃將MTP整合到更多Pixel裝置中，並探索並行解碼等替代架構，以進一步提升邊緣裝置上的AI效率。