Diseñador Recurrente para Decodificación Especulativa Rápida en Modelos de Lenguaje a Gran 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

Resumen

En este artículo, presentamos un enfoque mejorado de decodificación especulativa destinado a mejorar la eficiencia en el servicio de modelos de lenguaje grandes. Nuestro método aprovecha las fortalezas de dos técnicas establecidas: el enfoque clásico de decodificación especulativa con dos modelos y el enfoque más reciente de un solo modelo, Medusa. Inspirándonos en Medusa, nuestro enfoque adopta una estrategia de un solo modelo para la decodificación especulativa. Sin embargo, nuestro método se distingue al emplear una única cabeza de borrador ligera con un diseño de dependencia recurrente, similar en esencia al uso de un modelo de borrador pequeño en la decodificación especulativa clásica, pero sin las complejidades de la arquitectura completa de transformadores. Y debido a la dependencia recurrente, podemos utilizar búsqueda por haz para filtrar rápidamente candidatos no deseados con la cabeza de borrador. El resultado es un método que combina la simplicidad del diseño de un solo modelo y evita la necesidad de crear una estructura de atención en árbol dependiente de datos solo para la inferencia en Medusa. Demostramos empíricamente la efectividad del método propuesto en varios modelos de lenguaje de código abierto populares, junto con un análisis exhaustivo de las compensaciones involucradas al adoptar este enfoque.

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.