在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效率。