LLM2D
协作推断以实现高效的大语言模型解码与 tokens 级路由
CITER: Collaborative Inference for Efficient Large Language Model Decoding with Token-Level Routing
作者: Wenhao Zheng, Yixiao Chen, Weitong Zhang, Souvik Kundu, Yun Li, Zhengzhong Liu, Eric P. Xing, Hongyi Wang, Huaxiu Yao
发布日期: 2/5/2025
arXiv ID: 2502.01976

摘要

arXiv:2502.01976v1 交叉类型: cross 摘要: 大型语言模型在各种任务中取得了显著的成功,但在推理过程中面临着高计算成本的问题,这限制了它们在资源受限的应用中的部署。为了解决这一问题,我们提出了一种新的CITER(Collaborative Inference with Token-level Routing)框架,该框架通过令牌级别路由策略使小型和大型语言模型(SLMs & LLMs)能够高效协作。具体而言,CITER将非关键令牌路由给SLM以提高效率,将关键令牌路由给LLM以获得泛化质量。我们将路由器训练视为策略优化,其中路由器根据预测的质量和生成的推理成本获得奖励。这使路由器能够学习预测令牌级别的路由分数,并基于当前令牌及其决策对未来影响做出路由决策。为了进一步加速奖励评估过程,我们引入了一种快捷方式,该快捷方式显著减少了奖励估计的成本并提高了我们方法的实际可行性。在五个基准数据集上的广泛实验表明,CITER在保持高质量生成的同时减少了推理成本,为实时和资源受限的应用提供了一个有前景的解决方案。