Спекулятивное декодирование ближайшего соседа для генерации и атрибуции языковых моделей на основе машинного обучения

Nearest Neighbor Speculative Decoding for LLM Generation and Attribution

May 29, 2024

Авторы: Minghan Li, Xilun Chen, Ari Holtzman, Beidi Chen, Jimmy Lin, Wen-tau Yih, Xi Victoria Lin

cs.AI

Аннотация

Большие языковые модели (LLM) часто генерируют галлюцинации и не способны обеспечить атрибуцию для своих генераций. Полупараметрические языковые модели, такие как kNN-LM, решают эти ограничения, улучшая вывод LM для заданного запроса с использованием его ближайших соседей из непараметрического хранилища данных. Однако эти модели часто характеризуются медленной скоростью вывода и порождают несвязные тексты. В данной статье мы представляем метод ближайших соседей спекулятивного декодирования (NEST), новый полупараметрический подход к языковому моделированию, способный включать текстовые фрагменты произвольной длины из реального мира в генерации LM и обеспечивать атрибуцию их источникам. NEST выполняет поиск токенов на уровне каждого шага вывода для вычисления полупараметрического смешанного распределения и определения перспективных продолжений фрагментов в корпусе. Затем он использует приближенную процедуру спекулятивного декодирования, которая принимает префикс извлеченного фрагмента или генерирует новый токен. NEST значительно повышает качество генерации и уровень атрибуции базовой LM на различных задачах, требующих большого объема знаний, превосходя традиционный метод kNN-LM и конкурентно справляясь с увеличением извлечения в контексте. Кроме того, NEST существенно улучшает скорость генерации, достигая ускорения в 1,8 раза во времени вывода при применении к 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.