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

• URL: http://arxiv.org/abs/2502.01976v3
• Date: Sun, 09 Feb 2025 17:47:41 GMT
• Title: CITER: Collaborative Inference for Efficient Large Language Model Decoding with Token-Level Routing
• Authors: Wenhao Zheng, Yixiao Chen, Weitong Zhang, Souvik Kundu, Yun Li, Zhengzhong Liu, Eric P. Xing, Hongyi Wang, Huaxiu Yao,
• Abstract summary: CITER enables efficient collaboration between small and large language models (SLMs & LLMs) through a token-level routing strategy. 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. Our experiments show that CITER reduces the inference costs while preserving high-quality generation, offering a promising solution for real-time and resource-constrained applications.
• Score: 56.98081258047281
• License:
• Abstract: 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 CITER (Collaborative Inference with Token-lEvel Routing) 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.

Related papers

• Universal Model Routing for Efficient LLM Inference [72.65083061619752]
  We consider the problem of dynamic routing, where new, previously unobserved LLMs are available at test time. We propose a new approach to this problem that relies on representing each LLM as a feature vector, derived based on predictions on a set of representative prompts. We prove that these strategies are estimates of a theoretically optimal routing rule, and provide an excess risk bound to quantify their errors.
  arXiv  Detail & Related papers  (2025-02-12T20:30:28Z)
• Confident or Seek Stronger: Exploring Uncertainty-Based On-device LLM Routing From Benchmarking to Generalization [61.02719787737867]
  Large language models (LLMs) are increasingly deployed and democratized on edge devices. One promising solution is uncertainty-based SLM routing, offloading high-stakes queries to stronger LLMs when resulting in low-confidence responses on SLM. We conduct a comprehensive investigation into benchmarking and generalization of uncertainty-driven routing strategies from SLMs to LLMs over 1500+ settings.
  arXiv  Detail & Related papers  (2025-02-06T18:59:11Z)
• Reward-Guided Speculative Decoding for Efficient LLM Reasoning [80.55186052123196]
  We introduce Reward-Guided Speculative Decoding (RSD), a novel framework aimed at improving the efficiency of inference in large language models (LLMs) RSD incorporates a controlled bias to prioritize high-reward outputs, in contrast to existing speculative decoding methods that enforce strict unbiasedness. RSD delivers significant efficiency gains against decoding with the target model only, while achieving significant better accuracy than parallel decoding method on average.
  arXiv  Detail & Related papers  (2025-01-31T17:19:57Z)
• FTP: A Fine-grained Token-wise Pruner for Large Language Models via Token Routing [17.01412432658081]
  Large language models (LLMs) have demonstrated superior performance across various tasks by adhering to scaling laws. We propose a fine-grained token-wise pruning approach for the LLMs, which presents a learnable router to adaptively identify the less important tokens. Our approach achieves state-of-the-art (SOTA) pruning results, surpassing other existing pruning methods.
  arXiv  Detail & Related papers  (2024-12-16T07:09:46Z)
• Pruning All-Rounder: Rethinking and Improving Inference Efficiency for Large Vision Language Models [42.124670377223175]
  We propose a novel framework for inference acceleration called the Pruning All-Rounder (PAR) With a self-supervised learning manner, our method achieves a superior balance between performance and efficiency. Notably, PAR is highly flexible, offering multiple pruning versions to address a range of pruning scenarios.
  arXiv  Detail & Related papers  (2024-12-09T13:02:35Z)
• RouteLLM: Learning to Route LLMs with Preference Data [41.687640419561504]
  Large language models (LLMs) exhibit impressive capabilities across a wide range of tasks, yet the choice of which model to use often involves a trade-off between performance and cost. We propose several efficient router models that dynamically select between a stronger and a weaker LLM during inference. We develop a training framework for these routers leveraging human preference data and data augmentation techniques to enhance performance.
  arXiv  Detail & Related papers  (2024-06-26T18:10:22Z)
• OverPrompt: Enhancing ChatGPT through Efficient In-Context Learning [49.38867353135258]
  We propose OverPrompt, leveraging the in-context learning capability of LLMs to handle multiple task inputs. Our experiments show that OverPrompt can achieve cost-efficient zero-shot classification without causing significant detriment to task performance.
  arXiv  Detail & Related papers  (2023-05-24T10:08:04Z)
• MARLIN: Soft Actor-Critic based Reinforcement Learning for Congestion Control in Real Networks [63.24965775030673]
  We propose a novel Reinforcement Learning (RL) approach to design generic Congestion Control (CC) algorithms. Our solution, MARLIN, uses the Soft Actor-Critic algorithm to maximize both entropy and return. We trained MARLIN on a real network with varying background traffic patterns to overcome the sim-to-real mismatch.
  arXiv  Detail & Related papers  (2023-02-02T18:27:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.