Google

The chain rule of probability is a core mathematical concept used in language modeling to determine the joint probability of a sequence of tokens, denoted as $$x_0, x_1, \dots, x_m$$. This rule breaks down the overall sequence probability into a product of individual conditional probabilities. Specifically, the probability of the entire sequence is calculated as: $$\Pr(x_0, \dots, x_m) = \Pr(x_0) \cdot \Pr(x_1|x_0) \cdot \Pr(x_2|x_0, x_1) \cdots \Pr(x_m|x_0, \dots, x_{m-1})$$. Using the compact product notation, this equates to: $$\Pr(x_0, \dots, x_m) = \prod_{i=0}^{m} \Pr(x_i|x_0, \dots, x_{i-1})$$. Furthermore, to enhance computational efficiency and numerical stability, this product is frequently transformed into an alternative logarithmic form.

Chain Rule for Sequence Probability

In the chain rule for sequence probability, the base case is the first token, $x_0$. Since there are no preceding tokens, its probability is its marginal probability, $\text{Pr}(x_0)$. In many language models, this initial token is a deterministic start-of-sequence symbol, meaning its probability is fixed at 1, i.e., $\text{Pr}(x_0) = 1$. This assumption simplifies the joint probability calculation for the rest of the sequence. Specifically, the probability of the sequence following the initial token, $\text{Pr}(x_1, ..., x_m|x_0)$, is unaffected when multiplied by $\text{Pr}(x_0)$, as shown by the equation: $$ \text{Pr}(x_0) \text{Pr}(x_1, ..., x_m|x_0) = \text{Pr}(x_1, ..., x_m|x_0) $$ This effectively means the calculation can start from the second token, conditioned on the first.

Base Case for Sequence Probability

The joint probability of generating a specific sequence of tokens, such as `⟨s⟩ a b c d`, is calculated by applying the chain rule. This method decomposes the joint probability into a product of conditional probabilities, where each token's probability is conditioned on all the tokens that precede it. The expanded formula for this example is: $$Pr(⟨s⟩) \cdot Pr(a|⟨s⟩) \cdot Pr(b|⟨s⟩ a) \cdot Pr(c|⟨s⟩ a b) \cdot Pr(d|⟨s⟩ a b c)$$

Joint Probability of a Generated Sequence using the Chain Rule

The relationship between the joint, conditional, and marginal probabilities of token sequences is defined by the chain rule of probability. In logarithmic form, this allows the conditional log-probability of an output sequence `y` given an input `x` to be expressed in terms of the joint and marginal log-probabilities. The formula is: $$ \log \text{Pr}_{\theta}(\mathbf{y}|\mathbf{x}) = \log \text{Pr}_{\theta}([\mathbf{x}, \mathbf{y}]) - \log \text{Pr}_{\theta}(\mathbf{x}) $$ This identity can be rearranged to show how the joint log-probability of the concatenated sequence `[x, y]` is composed of the marginal and conditional parts: $$ \log \text{Pr}_{\theta}([\mathbf{x}, \mathbf{y}]) = \log \text{Pr}_{\theta}(\mathbf{x}) + \log \text{Pr}_{\theta}(\mathbf{y}|\mathbf{x}) $$ This relationship is fundamental for training language models, where the objective is often to maximize the conditional log-probability `log Pr_θ(y|x)`.

Relationship Between Joint, Conditional, and Marginal Log-Probabilities of Sequences

The log-probability of a sequence $\mathbf{x} = (x_0, \dots, x_m)$ is derived by applying the logarithm to the product form of the chain rule of probability. This key step transforms the product of conditional probabilities into a more computationally stable sum. The derivation proceeds as follows: 

$ \log \text{Pr}(\mathbf{x}) = \log \text{Pr}(x_0 \dots x_m) $ 

$ = \log [\text{Pr}(x_0) \text{Pr}(x_1|x_0) \cdots \text{Pr}(x_m|x_0 \dots x_{m-1})] $ 

$ = \log \text{Pr}(x_0) + \sum_{j=1}^{m} \log \text{Pr}(x_j|\mathbf{x}_{<j}) $ 

This decomposition is a foundational step for formulating the log-likelihood objective in language models.

Derivation of Sequence Log-Probability via Chain Rule

The chain rule for calculating the joint probability of a sequence can be expressed in an alternative logarithmic form. This is achieved by taking the logarithm of the entire probability expression, which transforms the product of conditional probabilities into a sum. This summation form is computationally more stable, especially for long sequences, as it mitigates the risk of numerical underflow from multiplying many small fractions. The formula is: $$\log \Pr(x_0,...,x_m) = \sum_{i=0}^{m} \log \Pr(x_{i}|x_0,...,x_{i-1})$$. In this formulation, it is assumed that for the initial token where $$i=0$$, the probability is $$\Pr(x_{i}|x_0,...,x_{i-1}) = \Pr(x_0) = 1$$. As a consequence of this assumption, the overall probability of the sequence simplifies to $$\Pr(x_0,...,x_m) = \Pr(x_0)\Pr(x_1,...,x_m|x_0) = \Pr(x_1,...,x_m|x_0)$$.

Logarithmic Form of the Chain Rule for Sequence Probability

The formula for an initial event or token $x_0$ in a sequence being an impossible event is: $$Pr(x_0) = 0$$

Formula for an Impossible Initial Event

A language model is tasked with calculating the total probability of the three-token sequence 'the cat sat'. The model provides the following probability estimates:

*   The probability of the first token is `Pr("the") = 0.1`
*   The probability of the second token, given the first, is `Pr("cat" | "the") = 0.5`
*   The probability of the third token, given the first two, is `Pr("sat" | "the", "cat") = 0.8`

Using the principle that the joint probability of a sequence is the product of the conditional probabilities of its components, what is the joint probability `Pr("the", "cat", "sat")`?

When calculating the joint probability of a very long sequence of tokens by multiplying many conditional probabilities (each between 0 and 1), a computational issue known as numerical underflow can occur. Explain what numerical underflow is in this context and describe how the formula for joint probability can be adapted to mitigate this problem.

Computational Stability of Sequence Probability

Which of the following expressions correctly decomposes the joint probability of a four-token sequence `(x₁, x₂, x₃, x₄)` using the chain rule of probability?

Learn Before

Related