BiTA: Ajuste Bidirecional para Aceleração sem Perdas em Modelos de Linguagem de Grande Escala

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

January 23, 2024

Autores: Feng Lin, Hanling Yi, Hongbin Li, Yifan Yang, Xiaotian Yu, Guangming Lu, Rong Xiao

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) comumente empregam geração autoregressiva durante a inferência, resultando em alta demanda de largura de banda de memória e, consequentemente, em latência prolongada. Para mitigar essa ineficiência, apresentamos o Bi-directional Tuning for Lossless Acceleration (BiTA), um método inovador que acelera LLMs por meio de geração semi-autoregressiva otimizada e verificação de rascunhos. Inspirados pelo conceito de ajuste de prompt, aprimoramos LLMs com um design eficiente em parâmetros chamado ajuste bidirecional para capacitar a geração semi-autoregressiva. Utilizando decodificação baseada em árvore eficiente, os modelos realizam a geração de candidatos a rascunho e a verificação em paralelo, garantindo saídas idênticas às suas contrapartes autoregressivas sob amostragem gananciosa. O BiTA serve como um módulo plug-in leve, aumentando de forma contínua a eficiência de inferência de LLMs existentes sem a necessidade de modelos de assistência adicionais ou custos significativos de memória extra. Aplicando o BiTA proposto, o LLaMA-2-70B-Chat alcança um aumento de velocidade de 2,7 vezes no benchmark MT-Bench. Experimentos extensivos confirmam que nosso método supera as técnicas de aceleração mais avançadas.

English

Large language models (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.7times speedup on the MT-Bench benchmark. Extensive experiments confirm our method surpasses state-of-the-art acceleration techniques.