Constrained Decoding › The Core Idea: Token Masking

To understand constrained decoding, we need to think carefully about what happens at each generation step. A language model produces a probability distribution over its entire vocabulary, perhaps 50,000 tokens or more. From this distribution, we sample a single token to continue the sequence. The challenge is that the model has no built-in knowledge of our structural requirements. It might assign high probability to tokens that would break our desired format.

The solution is simple: before sampling, we modify the distribution to exclude impossible choices. This is called token masking, and it forms the foundation of all constrained decoding techniques.

Consider a concrete scenario. You're generating JSON, and so far you've produced {"name": "Alice". What tokens could come next in valid JSON?

• A comma , to add another field
• A closing brace } to end the object

What tokens would break the JSON?

• A letter like a or b (no string delimiter)
• A number without a preceding colon
• Another opening brace { (not a valid position)

The model might assign probability to any of these tokens based on patterns it learned during training. Our job is to intercept the sampling process and remove the invalid options before the model can choose them.

At each step $t$ of generation, given the tokens produced so far $(x_1, \ldots, x_{t-1})$, we follow a four-step process:

1. Query the model: Compute the probability distribution $P(x_t | x_1, \ldots, x_{t-1})$ over the full vocabulary
2. Evaluate constraints: Determine which tokens are valid given our structural requirements and the current partial output
3. Mask invalid tokens: Set the probability of every invalid token to zero
4. Renormalize and sample: Adjust the remaining probabilities to sum to 1, then sample

The critical question is: how do we "adjust" the probabilities in step 4? We can't simply zero out invalid tokens and sample, because the remaining probabilities no longer form a valid distribution (they don't sum to 1). We need a principled approach.

Let $V_{\text{valid}}$ denote the set of tokens that satisfy our constraint given the current context. We want to construct a new probability distribution, $P_{\text{constrained}}$, that:

• Assigns zero probability to invalid tokens
• Preserves the relative ranking among valid tokens
• Forms a proper probability distribution (sums to 1)

The solution is to divide each valid token's probability by the total probability mass of all valid tokens:

$$P_{\text{constrained}}(x_t) = \begin{cases} \frac{P(x_t)}{\sum_{x \in V_{\text{valid}}} P(x)} & \text{if } x_t \in V_{\text{valid}} \\ 0 & \text{otherwise} \end{cases}$$

where:

• $x_t$: the candidate token at generation step $t$
• $P(x_t)$: the original probability assigned by the model to token $x_t$
• $V_{\text{valid}}$: the set of tokens that satisfy the constraint given the current context
• $\sum_{x \in V_{\text{valid}}} P(x)$: the sum of probabilities over all valid tokens, serving as the normalization constant

The formula accomplishes exactly what we need through a technique called renormalization. Let's trace through why each component matters.

The numerator $P(x_t)$ preserves the model's original preference for token $x_t$. If the model strongly preferred this token before masking, it will still be favored afterward.

The denominator $\sum_{x \in V_{\text{valid}}} P(x)$ is the total probability mass that was originally assigned to all valid tokens. By dividing by this sum, we scale up the valid probabilities so they sum to 1.

The conditional structure ensures invalid tokens get exactly zero probability. They cannot be sampled, period.

Consider a worked example. Suppose the model produces these probabilities:

The valid tokens (, and }) together have probability mass $0.30 + 0.10 = 0.40$. The invalid tokens have mass $0.60$.

After applying the constrained formula:

• $P_{\text{constrained}}(\text{comma}) = \frac{0.30}{0.40} = 0.75$
• $P_{\text{constrained}}(\text{brace}) = \frac{0.10}{0.40} = 0.25$
• $P_{\text{constrained}}(\text{a}) = 0$, and so on for all invalid tokens

Notice that the 3:1 ratio between , and } is preserved (0.30:0.10 becomes 0.75:0.25). The model originally preferred comma over closing brace by this ratio, and constrained decoding respects that preference. We've filtered out the impossible options without overriding the model's learned intuitions about likely continuations.

This is the key insight that makes constrained decoding powerful. We're not overriding the model's preferences. We're filtering them. The model brings its understanding of language and context, including likely continuations. The constraint brings the structural requirements. Together, they produce output that is both fluent (because the model guides token selection among valid options) and valid (because invalid options are impossible to select).

If the model assigns probability 0.3 to token A and 0.1 to token B, and both are valid, the constrained distribution preserves this 3:1 ratio. Token A remains three times more likely than token B. We only intervene when the model would choose something structurally invalid.