Drafting Ricorrente per Decodifica Speculativa Veloce nei Modelli Linguistici di Grande Dimensione

Recurrent Drafter for Fast Speculative Decoding in Large Language Models

March 14, 2024

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

cs.AI

Abstract

In questo articolo, presentiamo un approccio migliorato di decodifica speculativa volto a incrementare l'efficienza nel servire modelli linguistici di grandi dimensioni. Il nostro metodo sfrutta i punti di forza di due tecniche consolidate: il classico approccio di decodifica speculativa a due modelli e il più recente approccio a modello singolo, Medusa. Traendo ispirazione da Medusa, il nostro approccio adotta una strategia a modello singolo per la decodifica speculativa. Tuttavia, il nostro metodo si distingue per l'utilizzo di una singola testa di draft leggera con un design di dipendenza ricorrente, simile in sostanza al piccolo modello di draft utilizzato nella decodifica speculativa classica, ma senza le complessità dell'architettura completa del transformer. Grazie alla dipendenza ricorrente, possiamo utilizzare la ricerca a fascio per filtrare rapidamente i candidati indesiderati con la testa di draft. Il risultato è un metodo che combina la semplicità del design a modello singolo ed evita la necessità di creare una struttura di attenzione ad albero dipendente dai dati solo per l'inferenza in Medusa. Dimostriamo empiricamente l'efficacia del metodo proposto su diversi modelli linguistici open source di rilievo, accompagnati da un'analisi completa dei compromessi coinvolti nell'adozione di questo approccio.

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.