DEV Community

Devanshu Biswas
Devanshu Biswas

Posted on

Don't ask the model for valid JSON — delete the wrong tokens, force every illegal logit to , and it's valid by construction

#ai #llm #json #machinelearning

You ask a model for JSON and it mostly complies — then prepends "Sure!", wraps the answer in a

``json fence, quotes a number, or adds a trailing comment, andJSON.parse()` throws. Prompting ("return ONLY JSON") lowers the failure rate but never zeroes it, because sampling is probabilistic: there's always some probability mass on the wrong token. Constrained decoding kills that at the source — not by asking nicely, but by making the wrong tokens impossible to sample. I built the whole engine client-side to watch a JSON object generate token by token; here's how it works.

A decoding step is logits → softmax → sample

Every step, the model emits one logit (a raw score) for every token in its vocabulary. Softmax turns those into probabilities; then you sample or take the argmax. The key realisation: constraining doesn't touch the model's weights. It edits the logits, in the tiny gap between "model produced scores" and "we pick a token".

The grammar is a state machine

The set of legal tokens depends on where you are in the output — and that "where" is a state. A grammar (a JSON schema, a regex, a context-free grammar) compiles into a state machine; each state knows which tokens keep the output valid and where each one leads. Deterministic states allow exactly one token; value states allow a small typed set. After "age": only number tokens are legal — a quote is forbidden — and that's how types are guaranteed.


js
const FSM = {
  VAL_NAME:  { allow:['Ada','Bob','Lin'], next:{'*':'S7'} },   // a STRING
  VAL_AGE:   { allow:['42','7','100'],    next:{'*':'S14'} },  // a NUMBER, no quote!
  VAL_ADMIN: { allow:['true','false'],    next:{'*':'S20'} },  // a BOOLEAN
};
Enter fullscreen mode Exit fullscreen mode

The mask: forbidden logits go to −∞

Ask the machine for the legal set, then set every other token's logit to −∞. After softmax, exp(−∞) = 0, so those tokens get probability exactly zero — not merely discouraged, physically un-sampleable. The remaining legal tokens are renormalised, so the model still expresses its preference, but only among valid choices. This single line is the entire trick.


js
function applyMask(logits, mask){
  return logits.map((l, i) => mask[i] ? l : -Infinity);   // ← the whole idea
}
Enter fullscreen mode Exit fullscreen mode

The loop is valid by construction

Wire it together: get the logits, mask by the current state, pick the best surviving token, emit it, advance the machine. Because every step can only emit a legal token, the finished string parses on the first try — no validation, no retry.


js
function generate(){
  let state = 'S0', out = '';
  while (state !== 'DONE'){
    const logits = model.step(out);                  // raw scores over VOCAB
    const masked = applyMask(logits, maskFor(state));// −∞ on illegal tokens
    const tok    = pick(masked);                     // can ONLY be legal
    out += tok;  state = advance(state, tok);
  }
  return out;   // {"name":"Ada","age":42,"admin":true}
}
JSON.parse(generate());   // never throws
Enter fullscreen mode Exit fullscreen mode

That's the difference from the validate-and-retry approach (generate freely, check against the schema, re-ask on failure). Retry is provider-agnostic — it works over any plain text API — but it wastes calls, can loop, and never guarantees success. Constrained decoding is valid on the first try with zero retries, but it constrains while the model writes instead of validating after, so it needs logit access.

The real catch: tokenizer alignment

My demo vocab was tidy words. Real grammars are over characters, but models emit sub-word tokens that don't line up with grammar rules. So "which tokens are legal" really means "which tokens are a legal continuation of the partial string" — a prefix check over a trie of the vocabulary. Doing it naively (rescan every token every step) is slow; the production libraries compile the grammar into a token-level automaton once, then mask in near-constant time per step.

Because it needs the logits, constrained decoding lives where you control decoding — local runtimes (llama.cpp GBNF), inference servers (vLLM), and libraries like Outlines, xgrammar and llguidance, or a provider's JSON/grammar mode. One caution: an over-tight grammar can distort the distribution and hurt quality, because the model can't back out of a corner — so constrain the shape and keep the content free. Flip the constrained/unconstrained toggle and watch the free model drift into garbage while the constrained one stays valid:

https://dev48v.infy.uk/ai/days/day50-constrained-decoding.html

Top comments (0)