Décodage spéculatif par plus proche voisin pour la génération et l'attribution de modèles de langage

Nearest Neighbor Speculative Decoding for LLM Generation and Attribution

May 29, 2024

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

cs.AI

papers.abstract

Les grands modèles de langage (LLM) ont souvent tendance à halluciner et manquent de capacité à fournir des attributions pour leurs générations. Les modèles de langage semi-paramétriques, tels que kNN-LM, abordent ces limitations en affinant la sortie d'un modèle de langage pour un prompt donné en utilisant ses correspondances de plus proches voisins dans un stockage de données non paramétrique. Cependant, ces modèles présentent souvent des vitesses d'inférence lentes et produisent des textes peu fluides. Dans cet article, nous introduisons le décodage spéculatif par plus proches voisins (NEST), une nouvelle approche de modélisation de langage semi-paramétrique capable d'intégrer des segments de texte du monde réel de longueur arbitraire dans les générations du modèle de langage et de fournir des attributions à leurs sources. NEST effectue une récupération au niveau des tokens à chaque étape d'inférence pour calculer une distribution de mélange semi-paramétrique et identifier des suites de segments prometteurs dans un corpus. Il utilise ensuite une procédure de décodage spéculatif approximative qui accepte un préfixe du segment récupéré ou génère un nouveau token. NEST améliore significativement la qualité de génération et le taux d'attribution du modèle de langage de base sur une variété de tâches nécessitant des connaissances, surpassant la méthode kNN-LM conventionnelle et rivalisant avec l'augmentation par récupération en contexte. De plus, NEST améliore considérablement la vitesse de génération, obtenant une accélération de 1,8x du temps d'inférence lorsqu'il est appliqué à 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.