Google

A prevalent strategy for utilizing the retrieved reference tokens in k-NN Language Models involves creating a new probability distribution over the vocabulary. This distribution is derived from the nearest neighbors and serves to guide the model's final prediction by incorporating context from the datastore.

Using Reference Tokens to Define a Vocabulary Distribution in k-NN LM

In determining the retrieval-based distribution for a $$k$$-nearest neighbors ($$k$$-NN) language model, a distance metric, $$d_v$$, is defined relative to the vocabulary, $$V$$. For a query's hidden state, $$\mathbf{h}_i$$, and a retrieved datastore key-value pair, $$(\mathbf{z}_j, w_j)$$, the value $$d_v$$ equals the distance between $$\mathbf{h}_i$$ and $$\mathbf{z}_j$$ if the token $$w_j$$ corresponds to the $$v$$-th entry of the vocabulary $$V$$. If $$w_j$$ does not match the $$v$$-th entry, $$d_v$$ is set to $${}0$$.

Aggregated Distance Calculation for k-NN Vocabulary Distribution

The final output probability distribution in a $${}k$$-NN language model is formed by combining the retrieval-based distribution from the $${}k$$-nearest neighbors with the standard output distribution of the base large language model (LLM). This combination is performed using a linear interpolation, regulated by a coefficient $$\lambda$$, which adjusts the relative weight of the two distributions.

Linear Interpolation of k-NN and LLM Distributions

A language model uses a retrieval mechanism to improve its predictions. To predict the next word for the prefix 'The cat chased the mouse under the...', the model finds the 5 most similar contexts from a large text collection. The words that followed these 5 contexts were: `sofa`, `chair`, `sofa`, `sofa`, `table`. Based *only* on this set of 5 retrieved words, describe the key characteristics of the new probability distribution that would be formed over the vocabulary. Specifically, which words would have high, low, and zero probability?

Characterizing a Retrieval-Based Probability Distribution

Analyze the provided scenario. Explain how the set of retrieved reference tokens is used to create a new probability distribution over the vocabulary. Based on this, predict which token would likely have the highest probability in this new distribution and justify your reasoning.

Evaluating a k-NN LM's Intermediate Output

A k-Nearest Neighbors Language Model (k-NN LM) is generating text and needs to predict the next token. It queries its datastore and retrieves the 5 nearest reference tokens, along with their corresponding distances: {"river": 0.1}, {"stream": 0.2}, {"river": 0.3}, {"ocean": 0.8}, {"river": 0.9}. How are these retrieved tokens and their distances used to construct a new probability distribution over the model's vocabulary?

Learn Before

Related