Prefix Language Modeling › The Prefix LM Objective

Now that we understand the intuition behind splitting sequences into prefix and target regions, we need a formal training objective. The key question is: how do we mathematically express what we want the model to learn?

Think about what happens when you complete someone's sentence. You take in the context they've provided, then generate a continuation that fits. Prefix LM captures this behavior: the model should fully understand the prefix before generating the target.

This leads to two requirements that our objective must encode:

1. Full prefix visibility: When predicting any target token, the model should have access to the complete prefix as well as earlier target tokens.
2. Causal target generation: Target tokens should be generated left-to-right, each depending on previous target tokens but not future ones.

Consider a sequence $x = (x_1, \ldots, x_n)$ divided at position $k$ into two parts: the prefix $x_{1:k} = (x_1, \ldots, x_k)$ and the target $x_{k+1:n} = (x_{k+1}, \ldots, x_n)$. The training objective minimizes the negative log-likelihood of predicting target tokens:

$$\mathcal{L}_{\text{Prefix-LM}} = -\sum_{t=k+1}^{n} \log P_\theta(x_t | x_{1:k}, x_{k+1:t-1})$$

where:

• $\mathcal{L}_{\text{Prefix-LM}}$: the prefix language modeling loss to minimize
• $n$: the total sequence length
• $k$: the position where the prefix ends and the target begins (the split point)
• $x_{1:k}$: the prefix tokens $(x_1, \ldots, x_k)$, which are fully visible bidirectionally
• $x_{k+1:t-1}$: the target tokens generated before position $t$, specifically $(x_{k+1}, \ldots, x_{t-1})$
• $P_\theta(x_t | x_{1:k}, x_{k+1:t-1})$: the probability the model with parameters $\theta$ assigns to token $x_t$, conditioned on the full prefix and all preceding target tokens
• The summation $\sum_{t=k+1}^{n}$ iterates only over target positions, from $k+1$ to $n$

The heart of this formula lies in the conditioning term $x_{1:k}, x_{k+1:t-1}$. Let's unpack what this means for a concrete example.

Suppose we have the sequence "The cat sat on the mat" with $k=3$ (prefix is "The cat sat"). When predicting the fourth word "on":

• The model sees: all of "The cat sat" (bidirectionally) + nothing yet from the target
• When predicting "the": sees "The cat sat" + "on"
• When predicting "mat": sees "The cat sat" + "on the"

Notice the asymmetry: prefix tokens are always fully visible, while target tokens accumulate causally. This is the mathematical expression of the L-shaped attention mask we'll implement shortly.

The loss summation starts at $k+1$, not at 1. This means we compute loss only on target tokens, ignoring prefix tokens entirely. Why?

The prefix is given as input at inference time. Asking the model to predict prefix tokens would be like testing whether someone can recite back what you just told them. That's not the skill we want to develop. Instead, we focus all learning signal on the generation task: given this context, what comes next?

This differs from causal LM, where every position contributes to the loss. In prefix LM, the prefix serves purely as conditioning context, its representation is refined through gradient flow, but no direct prediction loss applies to it.