Constrained generation: JSON-mode for tinygpt

Reliable JSON output is the most common production failure mode for small (1-3B) LLMs in agent loops. Even when the model “knows” how to produce JSON, its sampler picks a wrong byte ~1 in N times — and one wrong byte (a missing quote, a stray comma, an invalid character) poisons the entire output for downstream parsers.

Constrained generation eliminates this class of bug by enforcing the JSON grammar (and, optionally, a schema) at the logits level: at every decode step the sampler can ONLY pick a token whose bytes extend the output into a state that’s still consistent with the grammar. Invalid tokens are masked to -inf before softmax. The result: the produced output is byte-for-byte guaranteed to be a valid JSON value that matches the schema.

This is the same family of algorithm used by outlines, vLLM’s guided_json mode, and llama.cpp’s GBNF — we ship the character-level variant which is simpler to reason about than the token-level FSM but a touch slower per step.

Usage

tinygpt sample <model.tinygpt> \
    --prompt "Output JSON:" \
    --tokens 200 \
    --temperature 0.5 \
    --json-schema path/to/schema.json

The output is guaranteed to be:

• A complete JSON value of the type pinned by the schema’s root
• Matching every property/enum/type constraint in the schema
• Stopping (by default) at the closing token of the root value

Toggle --no-json-stop-on-complete to keep sampling past the closing token (useful for streaming protocols that emit JSON followed by a sentinel).

How it works

Two-piece pipeline

  schema.json
      │
      ▼
  JSONSchemaNode  ──── (parse, one-time)
      │
      ▼
  JSONSchemaFSM  ──── (character-level state machine)
      ▲
      │
  LogitsMasker  ──── (per-vocab byte table, built once at setup)

The state machine is a stack of Frames, each representing one level of nesting (a value being emitted, an object in some sub-phase, an array, a string body, a number with sub-state). At every byte the top-of-stack frame transitions; if it completes, it pops and the parent advances.

The mask, step by step

At decode step t:

1. The model produces logits over the full vocabulary.
2. For every token id, we look up its pre-computed UTF-8 byte sequence (built once at setup by decoding each id in isolation).
3. We clone the current FSM and try to feed the bytes through. If every byte is accepted, the token is valid; otherwise it’s not.
4. We build a [Float] mask of length vocab, with 0 for valid tokens and -inf for invalid.
5. We add the mask to the logits, then sample (argmax for greedy or categorical for temperature > 0).
6. We commit the sampled token’s bytes to the FSM. The mask guarantees this succeeds.

When the FSM reaches a terminal accepting state (isComplete), only EOS (if known) is allowed; the caller typically stops here.

Grammar enforced

The FSM enforces a minified subset of JSON: no insignificant whitespace between structural tokens. Whitespace inside string bodies is content and behaves normally. The rationale: small LLMs trained on prose have a strong whitespace prior — if the FSM permitted leading or inter-token whitespace, the model would emit whitespace forever without ever advancing into the structure. This matches the behavior of JSON.stringify(x) and what every constrained-generation library produces in practice.

Trailing commas ({...,} or [...,]) are rejected — we split the object/array afterOpen and afterComma phases so a close immediately after a comma is impossible.

Schema subset supported

The schema parser handles the JSON Schema Draft-07 features below. Anything else is silently treated as .any (no constraint at that position) — the model produces valid JSON, just not necessarily matching the unsupported feature.

Performance

Setup cost: O(vocab) tokenizer decode calls — under 1 second even for 128k vocab.

Per-step cost: O(vocab × avg-token-bytes) byte-level FSM probes. On M-class Apple Silicon this is ~5-10% overhead at vocab≤32k for typical JSON token-byte distributions. The mask is built on CPU, then shipped to GPU as one MLXArray add before sampling — the GPU sees a regular softmax-and-categorical.

Measured on the demo byte-level Shakespeare model (vocab=256, 12-layer, M1):

(Byte-level vocab is the cheap end of the spectrum. At 32k-128k vocab expect 5-10% overhead, which is in line with the published numbers for outlines and vLLM on equivalent grammars.)

Examples

Tool-call schema

{
  "type": "object",
  "properties": {
    "tool_name": {
      "type": "string",
      "enum": ["read_file", "run_test", "edit_file"]
    },
    "ok": { "type": "boolean" }
  },
  "required": ["tool_name", "ok"]
}

Run:

tinygpt sample demo.tinygpt \
    --prompt "JSON:" --tokens 200 --temperature 0.6 \
    --json-schema schema_toolcall.json

Output (real run on the demo model):

JSON:{"tool_name":"edit_file","ok":false}

Simple object

{
  "type": "object",
  "properties": {
    "x": {"type": "boolean"},
    "y": {"type": "integer"}
  },
  "required": ["x", "y"]
}

Output:

Out:{"y":-3,"x":true}

Boolean

{"type": "boolean"}

With --prompt "t" --tokens 50 --temperature 0:

t true

Enum string

{"type": "string", "enum": ["yes", "no", "maybe"]}

Forces the output to be one of the three quoted enum values.

Caveats

1. Token-string approximation. We pre-compute each token’s “rendered string” by decoding the id in isolation. For most tokenizers this matches the contextual rendering for ASCII-rich tokens, but BPE merges that depend on neighbouring tokens (rare for JSON-relevant bytes) can mismatch. If this becomes a problem in practice the next step is a token-level FSM (Outlines-style) that walks merges explicitly.

2. No insignificant whitespace. The FSM emits minified JSON — see “Grammar enforced” above. If a downstream consumer requires pretty-printed JSON, post-process the output.

3. Model can still get stuck. If the only valid continuation is one the model puts very low logit on (e.g., a Shakespeare model asked to produce a digit), greedy sampling may hit the --tokens cap on a single-byte run. Set --temperature > 0 to break ties.

4. The prompt is not part of the FSM. The FSM sees only model- generated tokens. You can prefill the prompt with anything; only the output is constrained.

5. Speculative decoding paths skip the mask. --draft / --heads generate K tokens in a burst then verify with the target — wiring the schema mask into that path requires masking K positions in parallel, which is a follow-up.

6. Required-key ordering is not enforced. The schema’s required set is checked at close-time, not at each key’s open-time. So {"b":2,"a":1} is accepted as long as both a and b are in properties and both are emitted before }. JSON-Schema doesn’t constrain key order so this matches the spec.

Files

• native-mac/Sources/TinyGPTModel/JSONSchema.swift — schema parser
• native-mac/Sources/TinyGPTModel/ConstrainedGen.swift — FSM + logits masker
• native-mac/Sources/TinyGPT/Sample.swift — --json-schema wiring
• native-mac/Tests/TinyGPTModelTests/ConstrainedGenTests.swift — FSM unit tests

Follow-ups

1. Token-level FSM. Pre-compute a trie of valid token-id sequences at setup; per-step lookup becomes O(log vocab) instead of O(vocab). ~5-10× lower overhead on large-vocab models.

2. Wire to --draft / --heads. Speculative-decode paths need the mask applied at all K positions simultaneously.

3. additionalProperties, oneOf, anyOf, $ref. Cover the rest of Draft-07. $ref is the highest-value because real-world schemas use it heavily.

4. Pattern / format strings. Compile regex / format constraints into nested FSMs and slot them into the string-body state. Same approach as outlines’s pattern grammar.