CITER: Collaborative Inference for Efficient Large Language Model Decoding with Token-Level Routing

要約

大規模言語モデルは、様々なタスクにおいて目覚ましい成功を収めているが、推論時に高い計算コストに悩まされ、リソースに制約のあるアプリケーションへの導入が制限されている。この問題に対処するために、我々はトークンレベルのルーティング戦略を通じて、小規模言語モデルと大規模言語モデル（SLMs ⇄ LLMs）間の効率的な連携を可能にする、新しいCollaborative Inference with Token-lEvel Routing (CITER)フレームワークを提案する。具体的には、CITERは効率化のためにクリティカルでないトークンをSLMにルーティングし、汎化品質のためにクリティカルなトークンをLLMにルーティングする。我々はルータの学習をポリシーの最適化として定式化し、ルータは予測の品質と生成の推論コストの両方に基づいて報酬を受け取る。これにより、ルータはトークン・レベルのルーティング・スコアを予測することを学習し、現在のトークンとその決定の将来の影響の両方に基づいてルーティング決定を行うことができる。報酬評価プロセスをさらに加速するために、報酬推定のコストを大幅に削減し、我々のアプローチの実用性を向上させるショートカットを導入する。5つのベンチマークデータセットでの広範な実験により、CITERが高品質な生成を維持しながら推論コストを削減し、リアルタイムでリソースに制約のあるアプリケーションに有望なソリューションを提供することが実証された。我々のデータとコードはhttps://github.com/aiming-lab/CITER。

要約(オリジナル)

Large language models have achieved remarkable success in various tasks but suffer from high computational costs during inference, limiting their deployment in resource-constrained applications. To address this issue, we propose a novel Collaborative Inference with Token-lEvel Routing (CITER) framework that enables efficient collaboration between small and large language models (SLMs \& LLMs) through a token-level routing strategy. Specifically, CITER routes non-critical tokens to an SLM for efficiency and routes critical tokens to an LLM for generalization quality. We formulate router training as a policy optimization, where the router receives rewards based on both the quality of predictions and the inference costs of generation. This allows the router to learn to predict token-level routing scores and make routing decisions based on both the current token and the future impact of its decisions. To further accelerate the reward evaluation process, we introduce a shortcut which significantly reduces the costs of the reward estimation and improving the practicality of our approach. Extensive experiments on five benchmark datasets demonstrate that CITER reduces the inference costs while preserving high-quality generation, offering a promising solution for real-time and resource-constrained applications. Our data and code are available at https://github.com/aiming-lab/CITER.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL