Decodificación Especulativa del Vecino Más Cercano para Generación y Atribución en Modelos de Lenguaje Grande

Nearest Neighbor Speculative Decoding for LLM Generation and Attribution

May 29, 2024

Autores: Minghan Li, Xilun Chen, Ari Holtzman, Beidi Chen, Jimmy Lin, Wen-tau Yih, Xi Victoria Lin

cs.AI

Resumen

Los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) a menudo alucinan y carecen de la capacidad de proporcionar atribución para sus generaciones. Los modelos de lenguaje semi-paramétricos, como kNN-LM, abordan estas limitaciones refinando la salida de un LM para una instrucción dada utilizando sus coincidencias más cercanas en un almacén de datos no paramétrico. Sin embargo, estos modelos suelen exhibir velocidades de inferencia lentas y producen textos poco fluidos. En este artículo, presentamos Nearest Neighbor Speculative Decoding (NEST), un novedoso enfoque de modelado de lenguaje semi-paramétrico que es capaz de incorporar fragmentos de texto del mundo real de longitud arbitraria en las generaciones del LM y proporcionar atribución a sus fuentes. NEST realiza recuperación a nivel de token en cada paso de inferencia para calcular una distribución de mezcla semi-paramétrica e identificar continuaciones prometedoras de fragmentos en un corpus. Luego, utiliza un procedimiento de decodificación especulativa aproximada que acepta un prefijo del fragmento recuperado o genera un nuevo token. NEST mejora significativamente la calidad de la generación y la tasa de atribución del LM base en una variedad de tareas intensivas en conocimiento, superando el método convencional kNN-LM y compitiendo con la mejora de recuperación en contexto. Además, NEST mejora sustancialmente la velocidad de generación, logrando una aceleración de 1.8x en el tiempo de inferencia cuando se aplica a Llama-2-Chat 70B.

English

Large language models (LLMs) often hallucinate and lack the ability to provide attribution for their generations. Semi-parametric LMs, such as kNN-LM, approach these limitations by refining the output of an LM for a given prompt using its nearest neighbor matches in a non-parametric data store. However, these models often exhibit slow inference speeds and produce non-fluent texts. In this paper, we introduce Nearest Neighbor Speculative Decoding (NEST), a novel semi-parametric language modeling approach that is capable of incorporating real-world text spans of arbitrary length into the LM generations and providing attribution to their sources. NEST performs token-level retrieval at each inference step to compute a semi-parametric mixture distribution and identify promising span continuations in a corpus. It then uses an approximate speculative decoding procedure that accepts a prefix of the retrieved span or generates a new token. NEST significantly enhances the generation quality and attribution rate of the base LM across a variety of knowledge-intensive tasks, surpassing the conventional kNN-LM method and performing competitively with in-context retrieval augmentation. In addition, NEST substantially improves the generation speed, achieving a 1.8x speedup in inference time when applied to Llama-2-Chat 70B.