Rascunhador Recorrente para Decodificação Especulativa Rápida em Modelos de Linguagem de Grande Escala

Recurrent Drafter for Fast Speculative Decoding in Large Language Models

March 14, 2024

Autores: Aonan Zhang, Chong Wang, Yi Wang, Xuanyu Zhang, Yunfei Cheng

cs.AI

Resumo

Neste artigo, apresentamos uma abordagem aprimorada de decodificação especulativa, com o objetivo de aumentar a eficiência na execução de grandes modelos de linguagem. Nosso método aproveita os pontos fortes de duas técnicas consolidadas: a abordagem clássica de decodificação especulativa com dois modelos e a abordagem mais recente de modelo único, Medusa. Inspirados pelo Medusa, nossa abordagem adota uma estratégia de modelo único para decodificação especulativa. No entanto, nosso método se distingue por empregar um único cabeçalho de rascunho leve com um design de dependência recorrente, semelhante em essência ao pequeno modelo de rascunho usado na decodificação especulativa clássica, mas sem as complexidades da arquitetura completa do transformer. E, devido à dependência recorrente, podemos usar a busca em feixe para filtrar rapidamente candidatos indesejados com o cabeçalho de rascunho. O resultado é um método que combina a simplicidade do design de modelo único e evita a necessidade de criar uma estrutura de atenção em árvore dependente de dados apenas para inferência, como no Medusa. Demonstramos empiricamente a eficácia do método proposto em vários modelos de linguagem de código aberto populares, juntamente com uma análise abrangente das compensações envolvidas na adoção dessa abordagem.

English

In this paper, we introduce an improved approach of speculative decoding aimed at enhancing the efficiency of serving large language models. Our method capitalizes on the strengths of two established techniques: the classic two-model speculative decoding approach, and the more recent single-model approach, Medusa. Drawing inspiration from Medusa, our approach adopts a single-model strategy for speculative decoding. However, our method distinguishes itself by employing a single, lightweight draft head with a recurrent dependency design, akin in essence to the small, draft model uses in classic speculative decoding, but without the complexities of the full transformer architecture. And because of the recurrent dependency, we can use beam search to swiftly filter out undesired candidates with the draft head. The outcome is a method that combines the simplicity of single-model design and avoids the need to create a data-dependent tree attention structure only for inference in Medusa. We empirically demonstrate the effectiveness of the proposed method on several popular open source language models, along with a comprehensive analysis of the trade-offs involved in adopting this approach.