Uncertainty-Aware Hybrid Inference with On-Device Small and Remote Large Language Models

TL;DR

This paper introduces an uncertainty-aware hybrid inference framework that combines small on-device models with remote large models, significantly reducing communication and computation costs while maintaining high accuracy.

Contribution

The paper proposes U-HLM, a novel structure that uses uncertainty measurement to skip unnecessary LLM inferences, improving efficiency in hybrid language models.

Findings

U-HLM reduces uplink and LLM computation by 45.93%.

Achieves up to 97.54% of LLM accuracy.

Doubles token throughput compared to non-skipping HLM.

Abstract

This paper studies a hybrid language model (HLM) architecture that integrates a small language model (SLM) operating on a mobile device with a large language model (LLM) hosted at the base station (BS) of a wireless network. The HLM token generation process follows the speculative inference principle: the SLM's vocabulary distribution is uploaded to the LLM, which either accepts or rejects it, with rejected tokens being resampled by the LLM. While this approach ensures alignment between the vocabulary distributions of the SLM and LLM, it suffers from low token throughput due to uplink transmission and the computation costs of running both language models. To address this, we propose a novel HLM structure coined Uncertainty-aware opportunistic HLM (U-HLM), wherein the SLM locally measures its output uncertainty and skips both uplink transmissions and LLM operations for tokens that are likely to be accepted. This opportunistic skipping is enabled by our empirical finding of a linear correlation between the SLM's uncertainty and the LLM's rejection probability. We analytically derive the uncertainty threshold and evaluate its expected risk of rejection. Simulations show that U-HLM reduces uplink transmissions and LLM computations by 45.93%, while achieving up to 97.54% of the LLM's inference accuracy and 2.54$\times$ faster token throughput than HLM without skipping.