Chimera: A Lossless Decoding Method for Accelerating Large Language Models Inference by Fusing all Tokens

TL;DR

Chimera tackles the latency bottleneck of large language model decoding by introducing a lossless speculative-decoding framework that fuses a lightweight draft model with the backbone LLM. It employs a three-part draft stack—a trigram encoder for short-range context, a full-context encoder that leverages existing backbone representations, and residual decoding heads to predict future tokens—trained with a next-word and distillation objective. The method uses tree-structured verification and flexible decoding strategies to draft and then verify multiple candidate sequences in parallel, achieving up to 2.7x average speedups on Vicuna and LLaMA-2 models with competitive accuracy relative to Medusa. These results demonstrate Chimera’s practical impact for accelerating LLM inference in multi-turn dialogues and other latency-sensitive applications without sacrificing performance.

Abstract

Large language models (LLMs) have demonstrated remarkable capabilities across various tasks. However, their widespread application is hindered by the resource-intensive decoding process. To address this challenge, current approaches have incorporated additional decoding heads to enable parallel prediction of multiple subsequent tokens, thereby achieving inference acceleration. Nevertheless, the accuracy of these decoding heads falls short of the auto-regressive decoding approach. In light of these limitations, we propose Chimera, a novel framework specifically designed for speculative sampling. Within this framework, we introduce a lightweight draft model that effectively utilizes previously generated tokens to predict subsequent words. To ensure both accuracy and efficiency, we present two strategies within the lightweight draft model. Firstly, we focus on capturing short-range dependencies at the bottom layer. Secondly, we leverage the readily available representations from the original LLM.Through empirical evaluation on the Vicuna and LlaMA-2 series, Chimera demonstrates impressive results, achieving an average latency speedup ratio of 2.7x compared to the vanilla auto-regressive decoding approach. This highlights the potential of our proposed framework in significantly improving the efficiency of large language models during the decoding process.

Figures (8)

• Figure 1: The architecture of the Chimera model . The embeddings of prefix sequence are denoted as $x_1, x_2, x_3, x_4, x_5$. The last hidden states are signified as $h^{LLM}_{1}$, $h^{LLM}_{2}$ . The trigram calculated by formula \ref{['trigram_calculations']} are represented by $h^{trigram}_{1}$, $h^{trigram}_{2}$ and the full context hidden states calculated by formula \ref{['fullcontext_calculations']} are signified as $h^{full}_{2}$, $h^{full}_{3}$ , with subscripts indicating their positions in the sequence .
• Figure 2: The architecture of trigram encoder. $x_{i-3}, x_{i-2} ... x_{i+1}$ represent the embeddings of token at each corresponding position. $h^{trigram}_{i-1},h^{trigram}_{i},h^{trigram}_{i+1}$ indicate the outputs of the trigram encoder.
• Figure 3: The provided illustration depicts the attention masks employed in the Full Context Encoder during the training phase. The grey rectangles represent the tokens that should be masked during the attention calculation. The input to the full context encoder is the concatenation of the last hidden states from the original language model and the outputs of the trigram encoder.
• Figure 4: The architecture of full context encoder and residual decoding heads. $h^{LLM}_{1}$ and $h^{LLM}_{2}$ represent the last hidden states of original LLM. $h^{full}_{2}$ and $h^{full}_{3}$ indicate the outputs of full context encoder.
• Figure 5: Inference process.The input sequence to the chimera model is "Deep learning has revolutionized". For each position in the input sequence, the output of every chimera head is a probability distribution over the draft tokens. From these distributions, we sample the top-k tokens with the highest probabilities.