Pixel上で凍結マルチトークン予測を用いたGemini Nanoモデルの高速化

Googleの研究チームは、デプロイ済みのGemini Nano v3モデルにマルチトークン予測（MTP）を後付けする新たな手法を発表しました。これにより、Pixel端末上でAI機能の高速化と省エネルギー化を実現しつつ、モデルの能力や安全性を維持できます。

モバイルデバイス上の大規模言語モデル（LLM）は、厳しいエネルギー予算とメモリ制限に直面しています。標準的な自己回帰生成は一度に1トークンずつ処理するためボトルネックとなります。MTPは、複数の将来トークンを同時に予測することで、生成プロセスをドラフトと検証の2段階に分割し、推論を加速します。しかし、従来の独立したドラフターモデルは貴重なメモリを消費し、メインモデルの計算状態を活用できません。

この問題を解決するため、研究チームは新しいアーキテクチャを設計しました。既に学習済みのGemini Nano v3モデルの重みを凍結し、その最終層に軽量なTransformerヘッド（MTPヘッド）を追加します。このヘッドのみをトレーニングして将来トークンの予測誤差を最小化します。バックボーンが凍結されているため、MTPは純粋な効率最適化となり、ベースモデルの性能や安全性に影響を与えません。

さらに、ゼロコピーアーキテクチャを採用しました。MTPヘッドは独自のKVキャッシュを持たず、メインモデルの凍結KVキャッシュに直接クロスアテンションします。これにより、ドラフト段階のプリフィルレイテンシを排除し、メモリ使用量をインスタンスあたり最大130MB削減します。実験の結果、Pixel 9デバイスにおいて、同等パラメータ数の独立したドラフターと比較して、MTPはタスクに応じて50%以上の速度向上を達成しました。特に、複雑な制約のある要約や書き換えといった指示追従タスク、およびスマート返信のような構造予測可能なタスクでトークン受入率が最大55%向上しました。

この技術はすでにPixel 9およびPixel 10シリーズに展開されています。実際のワークロード（AI通知要約や校正など）では、MTPは推論パスあたり平均で約2つの追加トークンを正しく予測し、検証ステップを削減します。これにより、重いプロセッサの起動時間が減り、バッテリー寿命が改善されます。今後は、さらなるPixelデバイスへの統合や、並列デコードなど代替アーキテクチャの探求を通じて、エッジデバイス上のAI効率をさらに高める計画です。