手機跑大模型聽起來很酷但真的跑起來之后問題就接踵而來了。我們常用的云端模型可以依賴服務(wù)器和數(shù)據(jù)中心但手機上的模型就沒這么強大的后盾。不僅如此它還要顧及電池、內(nèi)存、發(fā)熱還有盡量別讓用戶等太久。Google 最近分享了一篇文章介紹 Gemini Nano 在 Pixel 設(shè)備上的一次推理加速優(yōu)化Frozen Multi-Token Prediction。這個方法已經(jīng)用在 Pixel 9 和 Pixel 10 系列的 Gemini Nano v3 中服務(wù)于 AI 通知摘要AI Notification Summaries、文本校對Proofread 等端側(cè) AI 功能。本文就借這篇文章簡單看看手機端大模型為什么容易慢Frozen Multi-Token Prediction 大概怎么工作以及它為什么能讓 Gemini Nano 在 Pixel 上生成得更快一些。大模型為什么會慢大語言模型生成文字時通常是一個 token 一個 token 往外吐。你可以把它理解成一個很謹慎的打字員。每寫下一個詞都要看一遍前面的內(nèi)容再判斷下一個詞是什么。這個流程看著很穩(wěn)但放到手機上就很容易變慢。因為手機的算力、內(nèi)存和功耗都有限。模型每次只生成一個 token就意味著要反復調(diào)用推理流程也會更頻繁地占用內(nèi)存帶寬。Google 在原文里提到移動設(shè)備有嚴格的能耗預算和 RAM 限制而傳統(tǒng)自回歸生成方式會形成瓶頸影響體驗和電池消耗。所以端側(cè) AI 要提速一個很自然的方向就是能不能一次多生成幾個 token多猜 Token 再統(tǒng)一驗證Google 的 Frozen Multi-Token Prediction 的核心思路就是先讓一個輕量模塊提前猜幾個后續(xù) token再交給主模型檢查。這和 speculative decoding 的實現(xiàn)思路有點像。一般來說生成 N 個 token 需要大模型跑 N 次speculative decoding 把過程拆成兩步先由更快的小模塊生成候選 token再由主模型并行驗證。如果候選內(nèi)容和主模型判斷一致就可以一次接受多個 token如果中間不一致就從分歧處重新生成下一個 Token。這樣一來模型生成文字就不用永遠一步一步地往前挪。只要猜對了它就可以一次往前推進幾個 token。需要提一下的是Google 沒有為 Gemini Nano 單獨放一個很重的 drafter 模型。它在 Gemini Nano v3 的主模型后面接了一個輕量的 MTP head讓這個小模塊負責提前預測。圖 1Frozen MTP 的 zero-copy 架構(gòu)MTP Head 復用主模型已經(jīng)計算出的 hidden states 和 KV cache生成候選 token 后再由主模型統(tǒng)一驗證。Frozen Multi-Token Prediction 是什么這里 Frozen 的意思是主模型權(quán)重被凍結(jié)。Google 拿已經(jīng)訓練好的 Gemini Nano v3把它的權(quán)重固定住然后接上一個新的 MTP head。訓練時只訓練這個新增模塊主模型本身不動。Google 認為這樣 MTP 更像是一個效率優(yōu)化層不會影響基礎(chǔ)模型原有能力和安全對齊如果提前預測錯了也會在驗證階段被丟掉最終輸出仍由主模型決定。這點很適合端側(cè)部署。因為手機上的模型已經(jīng)進入正式環(huán)境重新訓練和重新適配都很麻煩。Frozen MTP 的做法更像是在原有模型旁邊加一個“提前預判器”主要目標是減少等待時間和資源浪費。端側(cè)優(yōu)化內(nèi)存也很關(guān)鍵圖 1 里還有一個很關(guān)鍵的細節(jié)zero-copy architecture。如果單獨放一個 drafter 模型它也要讀上下文、維護自己的 KV cache還要占用額外內(nèi)存。對手機來說這些都是成本。Google 的方案是讓 MTP Head 直接利用主模型已經(jīng)算好的狀態(tài)包括 hidden states 和 KV cache。這樣一來小模塊不用重新處理完整 prompt也不用重復維護一份上下文緩存。Google 提到這種設(shè)計可以消除 drafter 的額外 prefill 延遲并且相比獨立 drafter每個實例最多節(jié)省約 130MB 運行內(nèi)存。這也是手機端 AI 和云端 AI 很不一樣的地方。云端更容易通過堆算力解決問題手機端要算得快還要盡量少占內(nèi)存、少喚醒重處理器、少消耗電量。優(yōu)化后的性能在 Pixel 9 設(shè)備上的實驗中MTP drafter 相比參數(shù)量接近的獨立 drafter在部分任務(wù)上能帶來 50% 以上的加速。對于 smart replies 這類結(jié)構(gòu)更可預測的任務(wù)token 接受率最高提升到 55%。在 AI Notification Summaries 和 Proofread 等生產(chǎn)負載中MTP 平均每次推理可以正確多預測接近 2 個 token。圖 2MTP 結(jié)果圖上圖展示了 Gemini Nano 在 Pixel 9 不同應(yīng)用場景中引入 MTP 后的 token generation 效果對比。小結(jié)這篇文章可以看作 Google 對端側(cè)大模型推理的一次工程優(yōu)化記錄。它講的重點并不復雜手機里的大模型生成內(nèi)容時如果每次只生成一個 token就容易慢如果能提前預測幾個 token并且讓主模型統(tǒng)一驗證就有機會減少推理步數(shù)。Frozen MTP 的做法是在 Gemini Nano v3 后面加一個輕量預測模塊。主模型保持不動小模塊負責提前猜主模型負責檢查。猜中了生成更快猜錯了丟掉重來。簡而言之手機端大模型要變快除了模型本身要輕生成方式也要更省步驟、更省內(nèi)存。