Google

In Large Language Model (LLM) inference, beyond the model computation problem of calculating $$\Pr(\mathbf{y}|\mathbf{x})$$, lies the search problem. This problem focuses on how to efficiently find the best output sequence $$\hat{\mathbf{y}}$$ for a given input sequence $$\mathbf{x}$$ (or the generated KV cache). A naive approach is exhaustive search, which considers every possible output sequence to select the one with the highest prediction probability. While this method guarantees a globally optimal solution, a direct exhaustive search is impractical for LLMs because the number of potential output sequences grows exponentially with the length of $$\mathbf{y}$$.

The Search Problem in LLM Inference

In the context of Large Language Model (LLM) inference, a hypothesis is formally defined as a tuple consisting of an input sequence $$\mathbf{x}$$ and an output sequence $$\mathbf{y}$$. Because the input sequence $$\mathbf{x}$$ remains fixed during the inference process, a hypothesis can be simplified to represent only the generated output sequence.

Hypothesis in LLM Inference

The search problem for Large Language Model (LLM) inference can be mathematically re-expressed as finding the optimal output sequence, $$\hat{\mathbf{y}}$$, from the entire search space $$\mathcal{Y}$$ that maximizes the conditional probability given the input sequence $$\mathbf{x}$$. This is formally defined by the equation:
$$\hat{\mathbf{y}} = \argmax_{\mathbf{y} \in \mathcal{Y}} \Pr(\mathbf{y}|\mathbf{x})$$

Mathematical Formulation of the Search Problem in LLM Inference

Solving the search problem in LLM inference requires managing the fundamental trade-off between exploration and exploitation. Exploration involves searching broadly across the vast space of possible output sequences to discover novel, high-quality options. Exploitation, on the other hand, involves focusing on and refining the most promising sequences already found. The central challenge is to devise an efficient search strategy that balances these two aspects to produce high-quality outputs without conducting an exhaustive search.

Exploration vs. Exploitation in LLM Search

In NLP, the search space $$\mathcal{Y}$$ for generating output sequences is commonly structured as a search tree. Each node in this tree represents a prefix subsequence, and its child nodes are formed by appending one token from the vocabulary, along with the probability of predicting that token. The tree is organized into levels, where nodes at the same level correspond to sequences of identical length. As generation progresses, new tokens are incrementally appended, causing the tree to grow deeper and wider, representing an increasing number of potential sequence extensions.

Search Tree Structure in Token Generation

Due to the computational infeasibility of an exhaustive search across an exponentially large space of sequences, LLM inference relies on heuristic algorithms. These practical methods use pruning strategies to navigate the vast output space efficiently. Pruning involves avoiding the exploration of low-quality or improbable sequences at each decoding step, which trades guaranteed optimality for a manageable computational cost. Common heuristic techniques include greedy search and sampling-based search.

Heuristic Search Algorithms for LLM Inference

Search algorithms provide an efficient method for producing a collection of potential solutions, known as candidate solutions, for a given problem.

Efficient Generation of Candidate Solutions via Search Algorithms

The search phase in LLM inference involves finding an output sequence `y` that is considered optimal or sub-optimal. The quality of any candidate sequence is determined by its conditional log-probability, `log Pr(y|x)`, which the search process aims to maximize.

Search for Optimal or Sub-optimal Sequences in LLM Inference

In the context of LLM inference, where predictions are conditioned on an input sequence `x`, the starting point or 'root' of the search process can be conceptualized as a representation of this input `x`. All subsequent predictions branch out from this initial state.

Root of the Search Space as a Representation of Input (x)

A text generation model has a vocabulary of 10,000 possible words it can choose from for each position in a sequence. If this model were to find the optimal output by evaluating every single possible sequence, how would the total number of sequences to check change if the desired output length is increased from 3 words to 5 words?

Based on the principles of sequence generation, identify the fundamental flaw in the engineer's proposed plan and explain why it is not a viable approach for this task.

Evaluating an Inference Strategy

A colleague proposes using an 'exhaustive search' method for generating text from a language model, arguing it's the only way to guarantee finding the sequence with the highest overall probability. Explain the fundamental computational problem with this approach, specifically addressing how the number of possibilities to check is affected as the desired length of the output text increases.

The Impracticality of Exhaustive Search

The objective of finding the output `ŷ` that maximizes `Pr(y|x)` is a foundational concept in NLP, with historical roots in early probabilistic models for speech recognition and machine translation. While solving this optimization problem can be straightforward for simple tasks, such as token prediction with a very small language model, it poses significant computational hurdles in most practical LLM scenarios. These difficulties, which stem from both computing the conditional probability `Pr(y|x)` and searching the vast output space for the `argmax`, are recognized as fundamental and well-studied problems in the broader field of artificial intelligence, allowing for the application of established techniques.

Historical Context and Computational Challenges of Maximum Probability Prediction

An output sequence, denoted by `y`, can be mathematically represented as an ordered list of elements, such as tokens or actions. A common notation is `$y = (a_1, a_2, ..., a_{n_r})$`, where `$a_{n_r}$` is the final element in a sequence of length `$n_r$`.

Learn Before

Related