Decodificação Especulativa do Vizinho Mais Próximo para Geração e Atribuição em Modelos de Linguagem de Grande Escala

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

Resumo

Modelos de linguagem de grande escala (LLMs) frequentemente alucinam e carecem da capacidade de fornecer atribuição para suas gerações. Modelos de linguagem semi-paramétricos, como o kNN-LM, abordam essas limitações refinando a saída de um LM para um prompt específico usando suas correspondências de vizinhos mais próximos em um armazenamento de dados não paramétrico. No entanto, esses modelos frequentemente exibem velocidades de inferência lentas e produzem textos não fluentes. Neste artigo, introduzimos o Nearest Neighbor Speculative Decoding (NEST), uma nova abordagem de modelagem de linguagem semi-paramétrica que é capaz de incorporar trechos de texto do mundo real de comprimento arbitrário nas gerações do LM e fornecer atribuição às suas fontes. O NEST realiza recuperação em nível de token em cada etapa de inferência para calcular uma distribuição de mistura semi-paramétrica e identificar continuações de trechos promissores em um corpus. Em seguida, ele usa um procedimento de decodificação especulativa aproximada que aceita um prefixo do trecho recuperado ou gera um novo token. O NEST melhora significativamente a qualidade da geração e a taxa de atribuição do LM base em uma variedade de tarefas intensivas em conhecimento, superando o método convencional kNN-LM e competindo de forma competitiva com a ampliação de recuperação em contexto. Além disso, o NEST melhora substancialmente a velocidade de geração, alcançando um aumento de 1,8x no tempo de inferência quando aplicado ao 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.