让小模型和大模型协作,按词元高效分配计算任务。
CITER: Collaborative Inference for Efficient Large Language Model Decoding with Token-Level Routing
- 根据词元重要性,动态分配给小模型或大模型处理。
- 在五个数据集上降低推理成本,生成质量保持高位。
- 适合实时、资源受限场景的部署需求。
大语言模型在各类任务中取得显著成果,但推理阶段计算开销高,限制了其在资源受限场景的应用。为此,我们提出一种新型协同推理框架CITER,通过词元级路由策略实现小模型(SLM)与大模型(LLM)的高效协作。具体而言,CITER将非关键词元路由至小模型以提升效率,关键词元则路由至大模型以保障泛化质量。我们将路由器训练建模为策略优化问题,路由器依据预测质量与生成成本获得奖励,从而学习词元级路由评分,并基于当前词元及未来影响做出决策。为进一步加速奖励评估,我们引入捷径机制,显著降低奖励估算成本,提升方法实用性。在五个基准数据集上的大量实验表明,CITER在保持高质量生成的同时有效降低推理开销,为实时与资源受限应用提供了可行方案。代码与数据可在https://github.com/aiming-lab/CITER 获取。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。