GoogleのAI Edge Torch Generative APIでデバイス上のLLMを実現

Google AI Edge Torch Generative APIのご紹介

近年、生成AIモデルの進化は目覚ましく、特にモバイルデバイス上での処理能力が飛躍的に向上しています。Googleは、開発者が高性能の生成AIモデルをエッジデバイスにシームレスに導入できるように、AI Edge Torch Generative APIを発表しました。この新しいAPIは、PyTorchでのカスタムリーダブルな大規模言語モデル(LLM)を、TensorFlow Lite(TFLite)ランタイムを利用して展開できる機能を提供します。

AI Edge Torch Generative APIの主な特徴

このAPIを利用することで、開発者は様々な新しい機能をエッジデバイス上で実現可能です。ここでは、具体的な機能とそのメリットをいくつか紹介します。

  • カスタムトランスフォーマーのサポート: 開発者は、自分自身で設計したモデルを簡単に作成できます。
  • CPUでの高性能処理: 現在はCPU向けだけがサポートされていますが、今後はGPUおよびNPU向けのサポートも予定されています。
  • TFLiteの展開フローに完全対応: 既存のTFLiteフローに統合することで、量子化やランタイムの設定がシームレスに行えます。
  • 人気モデルとの互換性: TinyLlama、Phi-2、Gemma 2Bといったモデルと互換性があります。
  • クロスプラットフォームサポート: Android、iOS、Webで利用可能です。

パフォーマンスの向上について

AI Edge Torch Generative APIは、手書きのトランスフォーマーアルゴリズムの研究を元に最適化されています。その結果、多くの市場で使用されている人気の高いLLMを手軽にデバイスに導入でき、最高性能を維持しながら開発者の作業効率を向上させることができます。

以下の表は、現在知られているデバイス上での3つのモデルのベンチマークの一例です。これにより、どれくらいの処理速度向上が見込めるかを示しています。

開発者エクスペリエンスの充実

このAPIは、開発者が自分のニーズに合わせてモデルを設計できるよう、基本的な構造を提供しています。既存のモデルを再設計して新しい機能を追加したり、まったく新しいモデルを一から作成することが可能です。主な利点は以下の通りです:

  • 再マッピング機構: モデルの重みを簡単に合わせることができます。
  • 直感的な量子化API: 重みの削減をシンプルに行えます。
  • マルチシグネチャエクスポート: プレフィルやデコードなど複数のエントリーポイントを簡単に定義できます。

変換と量子化のプロセス

ai_edge_torchの変換APIを使用して、モデルを高効率のTensorFlow Liteフォーマットに変換することができます。プロセスは以下のステップで構成されます:

  • StableHLOへのエクスポート: PyTorchモデルを安定したグラフ形式に変換します。
  • 最適化されたTFLiteモデルの生成: グラフの最適化を行い、性能を最大化します。
  • 量子化手法の適用: モデルのサイズを削減し、処理速度を向上させます。

展開とクロスプラットフォームのサポート

LLMの推論は通常、多くの前処理・後処理を含むため、GoogleはMediaPipe LLM Inference APIという高レベルのAPIを提供しています。これを使うことで、開発者は複雑な推論パイプラインを簡単に実装できます。

また、TensorFlow Liteのコアランタイムを利用することで、AndroidやiOSアプリケーションへの統合が容易になります。必要な変更は最小限で済み、迅速に新機能を追加できます。

今後の展望とモデルエクスプローラーの機能

AI Edge Torch Generative APIのアップデートには、Webサポートの追加や、より良い量子化オプションの拡充が含まれています。デベロッパーコミュニティとの連携を重視し、APIの改良や新機能の要望に応えることを目指しています。また、次回のシリーズ記事では、モデル視覚化・デバッグツールであるModel Explorerの詳細にも触れる予定です。

Google AI Edge Torch Generative APIは、モバイルデバイス上での独自の生成AIモデルを実現したい開発者にとって大きな助けとなるでしょう。これにより、多様なニーズに応じたAIソリューションを簡単に提供できる時代が到来しています。