arXiv 2502.019762025-09-10cs.CLcs.AIcs.LGcs.PF

CITER：基于词元级路由的高效大语言模型解码协同推理

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

The University of North Carolina at Chapel Hill（北卡罗来纳大学教堂山分校）；Intel（英特尔）；Carnegie Mellon University（卡内基梅隆大学）；Mohamed bin Zayed University of Artificial Intelligence（穆罕默德·本·拉希德人工智能大学）；Rutgers University（罗格斯大学）

Wenhao Zheng, Yixiao Chen, Weitong Zhang, Souvik Kundu, Yun Li, Zhengzhong Liu, Eric P. Xing, Hongyi Wang, Huaxiu Yao

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

本文提出CITER框架，通过词元级路由协同小模型与大模型，以策略优化训练路由器，在降低推理成本的同时保持生成质量，适用于资源受限场景。

摘要：

大型语言模型在各种任务中取得了显著成功，但在推理过程中面临高计算成本，限制了其在资源受限应用中的部署。为解决这一问题，我们提出了一种新颖的基于词元级路由的协同推理（CITER）框架，通过词元级路由策略实现小型语言模型（SLMs）与大型语言模型（LLMs）之间的高效协作。具体而言，CITER将非关键词元路由至SLM以提高效率，并将关键词元路由至LLM以保证泛化质量。我们将路由器的训练形式化为策略优化，路由器根据预测质量和生成推理成本获得奖励。这使得路由器能够学习预测词元级路由分数，并基于当前词元及其决策的未来影响做出路由决策。为进一步加速奖励评估过程，我们引入了一种捷径，显著降低了奖励估计的成本，提高了我们方法的实用性。在五个基准数据集上的大量实验表明，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.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。