Vaibhav Mittal

Posted on Jul 30

I Expected JSON Grammar Masks to Kill Sampling Diversity. The Prompt Got There First.

#ai #llm #opensource #machinelearning

I started this project with a fairly simple suspicion.

When an LLM is forced to follow a JSON schema, the grammar removes most of
the vocabulary at every step. If I keep using ordinary sampling settings such
as temperature=0.7 and top_p=0.9, is there enough distribution left for
those settings to matter?

I expected the grammar mask to make the sampler act almost greedily. After
capturing the logits, checking the masks, and running 6,144 sampled outputs, I
got a less convenient answer.

For Qwen2.5-1.5B-Instruct, the sampler did collapse to one token very often,
but the grammar mask was not causing it in my main run. The model had already
made up its mind after seeing the full schema in the prompt.

A mask can remove 99% of token IDs while removing almost none of the
model's probability mass.

That sentence explained the first result. Then I withheld the schema from the
prompt and ran the comparison again. Legal probability mass fell from 0.9996
to 0.7962, and the mask started changing nuclei. The prompt had been doing
more work than I had given it credit for.

The short version: Qwen's grammar mask changed the nucleus size on 0 of
427 measured free-choice steps when the schema was shown. In a matched
follow-up with the schema withheld, it changed 120 of 556. TinyLlama's mask
changed 38 of 504 in the original run.

Measuring collapse

What I counted

Top-p sampling sorts tokens by probability and keeps the smallest prefix whose
total probability reaches p. Sampling then happens inside that prefix.

Suppose the next-token probabilities are:

token A: 0.93
token B: 0.04
token C: 0.02
token D: 0.01

At top_p=0.9, the nucleus contains only token A. The call is technically
sampling, but there is nothing to sample between.

I called that a collapsed step. More precisely:

collapse = nucleus size is 1

There is an important exception. If the grammar permits only one token, the
step is forced by syntax. A closing brace or quote should not count as evidence
that top-p removed diversity. I separated those steps before calculating any
headline rate.

The quantity I cared about was therefore:

collapsed free-choice steps
---------------------------
all free-choice steps

A free-choice step has at least two legal tokens before top-p is applied.

The same-logit comparison

The normal generation path is not enough for this question. A constrained
generation and a free generation can choose different tokens early, so their
later logits no longer describe the same prefix.

Instead, I captured one raw logit vector at each constrained-generation step
and evaluated two distributions from it:

the original model distribution
the same distribution after applying the grammar mask and renormalizing

For each one I recomputed nucleus size across a grid of temperatures and top-p
values. This gave me a same-logit counterfactual. I could ask whether the mask
changed the nucleus without confusing that effect with a different generated
trajectory.

I also recorded the legal probability mass:

Z_t = probability mass assigned to grammar-legal tokens

When Z_t is close to one, the model already wants to stay inside the grammar.
When it is small, the mask has to move a lot of probability.

The main run used:

Qwen2.5-1.5B-Instruct
TinyLlama-1.1B-Chat-v1.0
six JSON schemas
two prompts per schema
temperatures from 0.0 to 1.5
top-p values from 0.5 to 1.0

The schemas covered a tight enum, an integer, a free string, a patterned
string, a mixed object, and a mixed object with a reasoning string first.

One prompt detail is essential: _prompt_text appended the complete JSON
schema to the user message. Qwen saw the field names, types, enum values, and
required properties before I measured its legal probability mass. That setup
is common, but it is not the only way structured-output systems are built.

Checking the mask

Every result depends on the legal-token mask being right, so I started with a
schema small enough to check by hand:

{
  "type": "object",
  "properties": {
    "a": {"type": "integer"}
  },
  "required": ["a"]
}

I wrote a ten-state byte-level recognizer for the compact form {"a":12} and
compared its allowed tokens with XGrammar at every prefix. I also compared
XGrammar with Outlines.

The first comparison failed by exactly 271 token IDs at every prefix. That
number turned out to be the clue. Qwen's model head has 151,936 rows, while
its tokenizer has 151,665 entries. My checker was treating the padded model
rows as empty byte strings and accidentally allowing them as no-op tokens.

After excluding those rows, the hand-written recognizer, XGrammar, and
Outlines agreed at every checked prefix.

TinyLlama exposed another problem. SentencePiece can encode an opening brace
with a token whose grammar bytes include a leading space, even though the
ordinary decoded string looks like the same brace. The original hand check
was using the wrong beginning-of-sequence tokenization.

That bug was worth finding early. Later, the same leading-space behavior
created a misleading result when I compared generated token IDs with
re-encoded token IDs. All twelve TinyLlama sequences changed their first token
ID on the round trip, but the decoded bytes did not change at all. I now treat
the ID change as a tokenizer diagnostic, not as a changed output.

Results

Qwen: the two histograms are the same

With the full schema in the prompt, at T=0.7 and top_p=0.9, Qwen had 427
free-choice steps across 12 constrained generations. The nucleus contained
one token on 336 of those steps, a pooled rate of 78.7%.

The step count is not an independent sample size. Steps inside one generation
share a prompt, a schema, and a prefix. Giving each of the 12 generations
equal weight, the mean collapse rate was 82.5%, with a 10.0 percentage-point
standard deviation and a range from 70.0% to 100%.

The surprising part was the counterfactual:

Qwen measurement	result
pooled collapse rate	78.7%
mean legal probability mass	0.9984
mask changed nucleus size	0 of 427 steps
mask changed top token	0 of 427 steps

The free and constrained histograms sit directly on top of each other.

Figure 1. Free and constrained nucleus sizes at T=0.7 and top_p=0.9.
The two distributions overlap.

Qwen had the schema in its prompt and was already assigning almost all of its
probability to legal continuations. The mask removed vocabulary entries, but
hardly any probability. Removing low-probability tokens did not change the
top-p prefix.

The wider grid behaved in the expected direction. Lower temperature and lower
top-p produced more one-token nuclei. At top_p=1.0, there is no truncation,
so a free-choice step cannot count as collapsed under this definition.

Figure 2. Qwen collapse rate over the complete temperature and top-p grid.
At top_p=1.0, there is no nucleus truncation.

The prompt got there first

I ran a paired follow-up on the mixed-object schema. It includes strings, an
enum, an integer, a boolean, an array, and a nested object. I used one prompt
that still made sense without the schema:

Return a plausible record for a failed API request.

The model, grammar, sampler settings, and seeds 101 through 108 stayed fixed.
The only change was whether the full schema was appended to the user message.

measurement	schema shown	schema withheld
generations	8	8
prompt tokens	306	30
free-choice steps	543	556
mean generation `Z_t`	0.9996	0.7962
SD of generation `Z_t`	0.0001	0.0173
mask changed nucleus size	0	120
free counterfactual collapse	82.9%	59.5%
constrained collapse	82.9%	72.3%
mask-added collapse	0.0 points	12.8 points

Figure 3. Eight matched seeds on the mixed-object schema. Points in the left
panel are generation means; black markers show the mean and standard
deviation.

All 16 constrained generations finished and parsed as JSON. Every matched seed
had lower mean legal mass with the schema withheld. The mean paired drop was
0.2034. In the withheld condition, the mask changed 21.6% of free-choice
nuclei and changed the top token on 106 steps.

This is the mechanism I had missed. The mask's effect depends partly on how
much the prompt has already concentrated probability on legal tokens. It is
not determined by the model and grammar alone.

TinyLlama moved, but not by much

TinyLlama had 504 free-choice steps across 12 constrained generations. Its
pooled collapse rate was 76.2%. The generation-level mean was 76.7%, with an
11.6 percentage-point standard deviation and a range from 54.5% to 90.9%.

This time the mask had a measurable effect:

TinyLlama measurement	result
free counterfactual collapse rate	71.2%
constrained collapse rate	76.2%
mean legal probability mass	0.9239
mask changed nucleus size	38 of 504 steps
mask changed top token	36 of 504 steps

So the grammar added about five percentage points of collapse for TinyLlama.
That is real in this sample, but it is a much narrower result than saying
structured decoding generally turns top-p into greedy decoding.

The separate free control for TinyLlama was not clean enough to rescue a
larger claim. None of its 12 free generations parsed as JSON, and both prompts
hit the 128-token cap in four of the six schemas. The same-logit comparison is
the useful one here.

Schema shape still mattered

My original ordering prediction mostly survived. Tight enums and patterned
strings collapsed more often than ordinary string values. The gap was smaller
than I expected, and even free strings often had a one-token nucleus.

For Qwen at the reference setting, the schema-pooled rates ranged from 73.5%
for the free-string schema to 94.4% for the tight enum. TinyLlama ranged from
56.9% to 90.9% on those two schemas.

Figure 4. Qwen collapse rates by schema and approximate grammar position.

The position labels in this figure come from the decoded prefix. They are
useful for a rough breakdown, but a token that crosses two grammar positions
or an escaped string can fool the tagger.

Constraint pressure did not predict collapse

I expected low legal probability mass to line up with one-token nuclei. It
did not.

Pearson correlation between Z_t and the collapse flag was:

Qwen:      -0.020
TinyLlama: -0.056

Both are effectively zero in this sample.

Figure 5. Constraint pressure against the collapse flag for Qwen. The
correlation is -0.020.

The Qwen plot is mostly a pile of points near zero pressure because Z_t is
usually close to one. A few unusual steps do not explain the overall collapse
rate.

Whole-output diversity

The per-step measurement told me that the nucleus was often one token at
T=0.7, but it did not tell me whether sampler settings still changed whole
outputs.

For the second experiment I generated 6,144 Qwen outputs:

3 schemas
x 2 prompts
x 4 temperatures
x 2 top-p values
x 2 conditions
x 64 seeds
= 6,144 outputs

At top_p=0.9, the mean constrained distinct-output rate rose with
temperature:

temperature	constrained	free
0.3	14.3%	14.3%
0.7	42.4%	42.4%
1.0	49.5%	51.8%
1.5	60.7%	65.9%

Figure 6. Mean distinct-output rate across three schemas and two prompts,
with 64 seeds in each cell.

The constrained curve is not flat. Temperature plainly still did something.
For the free-string and mixed-object schemas, constrained and free distinct
counts matched in every top_p=0.9 cell on both prompts. The visible gap at
high temperature came mainly from the enum schema, where the grammar limits
the field to five values. That gap is the schema doing its job, not evidence
that the sampler settings had been disabled.

There is a length caveat at T=1.5. Some outputs hit the 128-token cap, and a
truncated string still counts as distinct. I would not treat the final points
as a clean measure of useful JSON diversity.

I also compared each step's probability distribution at T=0.7 and T=1.3
using total variation distance. Many steps barely moved because one token was
already dominant, while a smaller tail of steps remained temperature
sensitive. Once again, Qwen's free and constrained distributions looked
nearly the same.

Figure 7. Total variation distance between each Qwen step's distribution at
T=0.7 and T=1.3.

Order of operations can fail completely

I also checked where grammar masks sit in real sampling paths. In the pinned
versions I inspected, Hugging Face Transformers, vLLM, and SGLang all apply
the grammar mask before top-p.

For the Hugging Face path, the relevant order is:

penalties -> grammar mask -> temperature -> top-p

I deliberately reversed the last part and ran twelve seeds:

mask -> temperature -> top-p
temperature -> top-p -> mask

Mask-first completed all twelve generations with valid JSON. Top-p-first
failed all twelve because the top-p set and the legal-token set had an empty
intersection. Sampling then received an invalid probability row.

This was not a quality comparison. The reversed path simply failed. It also
confirmed that sampler order is part of the behavior, not an implementation
detail to ignore.

Implementation notes

I computed the metric grid while the full logit vector was still in memory.
My first storage plan saved only the top 4,096 logits and tried to reconstruct
everything later. That loses entropy information at high temperature and
cannot represent top_p=1.0, where the full support matters.

The first batched diversity run was also too slow because I moved full logits
to the CPU before sorting for top-p. Keeping logits, masks, and seeded
samplers on the GPU reduced the complete E2 run to 37 minutes and 24 seconds
with a batch size of eight.

I kept raw JSONL runs out of Git and committed the smaller aggregate CSVs,
figures, configs, source, and notes. The full repository is available here:

nucleus-collapse-constrained-decoding on GitHub

What this changes in practice

If the complete schema is already in the prompt, my Qwen results suggest that
a high one-token nucleus rate may be the model being confident, not the
grammar secretly turning sampling into greedy decoding. Temperature still
changed whole-output diversity in that setup.

If the schema is enforced by an API or gateway but withheld from the model,
the mask can do much more work. That is the setup where I would measure legal
probability mass and compare free and constrained nuclei on my own prompts
before assuming ordinary sampling settings behave the same way.

The practical question includes both the grammar engine and how much of the
output contract the model saw before decoding began.

What I would claim now

I would not say that JSON grammar masks generally collapse top-p sampling.
Two small models and twelve constrained generations per model are nowhere near
enough for that claim.

The prompt ablation is smaller still: one model, one schema, one prompt, and
eight matched seeds per condition. I read it as a mechanism check, not an
estimate of what every hidden-schema deployment will do.

What these runs do show is more specific:

A one-token nucleus can be common even after forced syntax steps are removed.
For Qwen2.5-1.5B-Instruct with the full schema in the prompt, the grammar mask did not cause that collapse at T=0.7 and top_p=0.9.
Withholding one mixed-object schema dropped mean generation Z_t from 0.9996 to 0.7962 and produced a 12.8-point mask effect on collapse.
TinyLlama showed a modest, model-dependent mask effect of about five percentage points.
Within the original schema-shown runs, constraint mass did not predict collapse in either model.
Temperature continued to change whole-output diversity in the larger Qwen run.

I began with a story about grammar masks silently disabling sampling. The
first measurements did not support it because the prompt had already pushed
Qwen onto legal tokens. When I removed that prompt information, the grammar
became visible in both probability mass and nucleus size.

That is not the result I planned to write up, but it is the one the logits
gave me.

DEV Community