BiTA: Bi-Directional Tuning for Lossless Acceleration in Large Language Models

Abstract: LLMs commonly employ autoregressive generation during inference, leading to high memory bandwidth demand and consequently extended latency. To mitigate this inefficiency, we present Bi-directional Tuning for lossless Acceleration (BiTA), an innovative method expediting LLMs via streamlined semi-autoregressive generation and draft verification. Inspired by the concept of prompt tuning, we enhance LLMs with a parameter-efficient design called bi-directional tuning for the capability in semi-autoregressive generation. Employing efficient tree-based decoding, the models perform draft candidate generation and verification in parallel, ensuring outputs identical to their autoregressive counterparts under greedy sampling. BiTA serves as a lightweight plug-in module, seamlessly boosting the inference efficiency of existing LLMs without requiring additional assistance models or incurring significant extra memory costs. Applying the proposed BiTA, LLaMA-2-70B-Chat achieves a 2.7$\times$ speedup on the MT-Bench benchmark. Extensive experiments confirm our method surpasses state-of-the-art acceleration techniques.

• The paper introduces a bi-directional tuning method that leverages semi-autoregressive generation with learnable embeddings to accelerate LLM inference.
• The paper demonstrates a consistent speedup between 2.1× and 3.3× on benchmarks with only a 0.01% increase in trainable parameters.
• The paper employs an efficient tree-based decoding mechanism to verify draft outputs, ensuring quality comparable to autoregressive generation.

Introduction

In recent advancements within the field of LLMs, there has been a marked shift towards enhancing the efficiency of these models. Despite their powerful generative capabilities, LLMs face challenges regarding inference latency, particularly in resource-constrained environments. The Bi-directional Tuning for lossless Acceleration (BiTA), a pioneering solution, has been introduced to expedite the inference process for LLMs using a methodology rooted in semi-autoregressive generation and a verification paradigm that maintains output fidelity to autoregressive (AR) generation.

Acceleration Techniques

LLM acceleration techniques range from model compression and architecture simplification to more intricate algorithmic modifications. A significant strand within these techniques involves efficient decoding methods that aim for speed without conceding output quality. Amongst these, semi-autoregressive (SAR) decoding has emerged as a promising paradigm for reducing inference executions. SAR decoding diverges from the conventional AR generation by decoding output tokens in parallel, but it introduces the challenge of SAR models often suffering from quality degradation when compared to their AR counterparts.

Methodology

BiTA introduces a dual-component system to address these challenges. Firstly, a parameter-efficient tuning method inspired by prompt tuning is employed to enable SAR generation, metaphorically described as learnable prefix and suffix embeddings in token sequences. Secondly, an efficient tree-based decoding mechanism facilitates both generation and verification of draft candidates. This allows for swift and simultaneous operations without necessitating additional validation steps or exterior models.

The paper's proposed method magnificently achieves a substantial speedup, as demonstrated with the LLaMA-2-70B-Chat model which sees a 2.7× acceleration on the MT-Bench benchmark. This is attained with a negligible increase in trainable parameters, introducing as few as 0.01% additional parameters, which acts as a testament to the efficiency and innovativeness of the proposed approach.

Experimental Results

BiTA's impact was measured across a spectrum of LLMs and tasks, showcasing a consistent speedup ranging from 2.1× to 3.3×. The performance benefits were found to be particularly pronounced in larger models, possibly due to richer embedding contexts which enhance prediction capabilities. The paper also outlines a stark improvement over state-of-the-art speculative decoding techniques, substantiating BiTA's potential in practical application scenarios. Furthermore, through a series of ablation studies, the researchers elucidated the impact of various prompting designs and configurations on the speedup performance, establishing the superiority of the bi-directional tuning and efficient tree-based decoding strategies.

Conclusion

BiTA's methodology holds considerable potential for advancing the application of LLMs in real-time and constrained resource scenarios, offering a compelling acceleration solution without compromising the generativity and integrity of the model outputs. This work not only contributes to the ongoing endeavor of improving LLM efficiency but also extends the utility of these models, reinforcing their applicability across various domains and applications.