Advanced Decoding & Structured Generation

Speculative decoding lets a small model draft and a big model verify. It's like having an intern write your emails, except this actually works.

Spectra, Speculatively Delegating AI Agent

4.3.1 Contrastive Decoding

What if you could isolate the "intelligence gap" between a large model and a small one? That is exactly what contrastive decoding does. The intuition is elegant: a large "expert" model and a smaller "amateur" model share many of the same failure modes (generic, repetitive text), but the expert model also captures higher-quality patterns that the amateur does not. By subtracting the amateur's preferences from the expert's, we amplify what makes the expert special and suppress what is generic.

Formally, the contrastive score for each token is:

Tokens that both models find likely (generic completions) get a low contrastive score because both log-probabilities are high. Tokens that only the expert finds likely get a high contrastive score. An additional constraint (plausibility filter) ensures we only consider tokens where the expert assigns at least some minimum probability, preventing nonsensical tokens from being selected just because the amateur dislikes them.

Figure 4.3.2: Contrastive decoding amplifies tokens the expert prefers over the amateur (content words) and suppresses tokens both models favor (generic function words).

# Contrastive decoding: subtract a weaker "amateur" model's log-probs
# from the expert's, amplifying tokens where the expert is most confident.
import torch
import torch.nn.functional as F
def contrastive_decode(expert_logits, amateur_logits,
    alpha=0.1, beta=0.5):
    """
    Contrastive decoding: amplify expert-specific preferences.
    alpha: plausibility threshold (keep tokens where expert prob > alpha * max_prob)
    beta: weight for the amateur subtraction
    """
    expert_probs = F.softmax(expert_logits, dim=-1)
    amateur_log_probs = F.log_softmax(amateur_logits, dim=-1)
    expert_log_probs = F.log_softmax(expert_logits, dim=-1)
    # Plausibility constraint: only consider tokens the expert finds plausible
    max_expert_prob = expert_probs.max()
    plausible_mask = expert_probs >= alpha * max_expert_prob
    # Contrastive score: expert - beta * amateur
    contrastive_scores = expert_log_probs - beta * amateur_log_probs
    # Apply plausibility mask
    contrastive_scores[~plausible_mask] = float('-inf')
    return contrastive_scores.argmax(dim=-1)
# Simulated example
expert_logits = torch.tensor([5.0, 2.0, 4.5, 1.5, 3.0, 0.5])
amateur_logits = torch.tensor([4.8, 3.5, 1.0, 3.0, 2.8, 0.3])
tokens = ["the", "a", "brilliant", "is", "novel", "xyz"]
expert_probs = F.softmax(expert_logits, dim=-1)
amateur_probs = F.softmax(amateur_logits, dim=-1)
contrastive = torch.log(expert_probs) - 0.5 * torch.log(amateur_probs)
print("Token | Expert P | Amateur P | Contrastive Score")
for i, t in enumerate(tokens):
    print(f"{t:10s} | {expert_probs[i]:.4f} | {amateur_probs[i]:.4f} | {contrastive[i]:.3f}")
    selected = contrastive_decode(expert_logits, amateur_logits)
    print(f"\nSelected token: '{tokens[selected]}'")

# Using Outlines with regex constraints
import outlines
model = outlines.models.transformers("microsoft/Phi-3-mini-4k-instruct")
# Force output to match a date pattern
date_generator = outlines.generate.regex(
    model,
    r"\d{4}-\d{2}-\d{2}"
)
# Force output to be one of specific choices
sentiment_generator = outlines.generate.choice(
    model,
    ["positive", "negative", "neutral"]
)
date = date_generator("The meeting is scheduled for next Tuesday. Today is 2025-03-20. The meeting date:")
sentiment = sentiment_generator("The movie was absolutely wonderful! Sentiment:")
print(f"Date: {date}")
print(f"Sentiment: {sentiment}")

Notice how "brilliant" wins the contrastive selection, even though "the" has the highest expert probability. The expert and amateur agree on "the" (both give it high probability), so it gets a low contrastive score. But "brilliant" is something only the expert strongly favors, making it the contrastive winner.

4.3.2 Speculative Decoding: The Core Idea

The bottleneck in autoregressive generation is that each token requires a full forward pass through the model, and tokens must be generated sequentially. Speculative decoding (also called assisted generation in the Hugging Face Transformers library) was introduced by Leviathan et al. (2023) and Chen et al. (2023). It speeds this up using a clever trick: a small, fast "draft" model generates several tokens quickly, and then the large "target" model verifies them all in a single forward pass.

The verification step uses a mathematical guarantee: each draft token is accepted with probability min(1, q(x)/p(x)), where q(x) is the target model probability and p(x) is the draft model probability. If a token is rejected, we resample from an adjusted distribution. This ensures that the final output has exactly the same distribution as if the target model had generated it alone.

Figure 4.3.3: Speculative decoding generates multiple draft tokens cheaply, then verifies them in a single pass of the expensive target model.

Why the Accept/Reject Scheme Is Exact

The previous paragraphs asserted that speculative decoding samples from the target distribution, but a graduate reader should see why the bookkeeping works out, because the whole appeal of the method (a free 2 to 3x speedup) collapses if the output distribution drifts even slightly. We keep the section's notation throughout: $q(x)$ is the target probability and $p(x)$ is the draft probability, so the draft proposes $x \sim p(x)$ and we accept it with probability $\min\!\big(1, q(x)/p(x)\big)$. The claim to prove is that the token actually emitted, whether it was the accepted draft or a resampled replacement, has marginal distribution exactly $q$.

Fix any vocabulary token $x$. There are two disjoint ways it can be emitted in a single speculative step, so its total emission probability is the sum of the two.

The first route is "the draft proposed $x$ and we accepted it." The draft proposes $x$ with probability $p(x)$, and conditioned on that proposal we accept with probability $\min\!\big(1, q(x)/p(x)\big)$. The product is

The second equality is the load-bearing simplification: multiplying $p(x)$ through the $\min$ turns $p(x)\cdot 1$ into $p(x)$ and $p(x)\cdot q(x)/p(x)$ into $q(x)$, and the smaller of those two survives. So the accepted-token mass at $x$ is exactly $\min(p(x), q(x))$, which is at most $q(x)$ and is short of $q(x)$ precisely when the draft over-proposed $x$ (when $p(x) > q(x)$).

The second route is "the draft was rejected, and the resampling step happened to produce $x$." A rejection happens with total probability $1 - \sum_{y}\min(p(y), q(y))$, and when it does we draw from the residual distribution

Here the numerator keeps only the tokens the target wanted more than the draft offered (the unmet target demand), and the denominator renormalizes. A standard identity ties the two quantities together: for any two distributions, $\sum_y \min(p(y), q(y)) + \sum_y \max(0, q(y) - p(y)) = 1$, because at every $y$ the pair $\min(p,q)$ and $\max(0, q-p)$ adds up to $q(y)$ when $q(y) \ge p(y)$ and to $p(y) + 0$ otherwise, and both $p$ and $q$ sum to 1. Therefore the rejection probability equals the residual normalizer:

That coincidence is what makes the algebra cancel. The probability of emitting $x$ via the rejection route is the rejection probability times $p_{\text{res}}(x)$, and the normalizers cancel:

Adding the two routes gives the marginal probability that token $x$ leaves the step:

The case split makes this last line obvious. When $q(x) \ge p(x)$ the two terms are $p(x)$ and $q(x) - p(x)$, summing to $q(x)$; when $q(x) < p(x)$ they are $q(x)$ and $0$, again summing to $q(x)$. So every token is emitted with exactly its target probability $q(x)$, which is the guarantee Leviathan et al. (2023) and Chen et al. (2023) established. The draft $p$ never appears in the answer: a bad draft only triggers more rejections (and less speedup), never a wrong distribution.

Expected Speedup and Its Ceiling

Correctness is settled, so the remaining question is the one practitioners actually care about: how much wall-clock time does this buy? The answer hinges on the acceptance rate $\alpha$, the average probability that a drafted token survives verification, and the block length $\gamma$, the number of tokens the draft proposes before each verification pass. Model the per-position accept events as independent with probability $\alpha$ (an approximation, since real acceptances are correlated, but a tight one in practice). Each verification step then accepts a prefix of the draft block up to the first rejection, and after the block the target itself contributes one guaranteed token from the corrected distribution.

The expected number of tokens produced per verification step is a truncated geometric sum:

The term $\alpha^{i}$ is the probability that the first $i$ drafted tokens are all accepted, and summing over $i = 0, \ldots, \gamma$ counts the expected length of the accepted run plus the one bonus token the target always supplies. With $\alpha = 0.8$ and $\gamma = 4$, this is $(1 - 0.8^{5})/(1 - 0.8) = (1 - 0.328)/0.2 \approx 3.36$ tokens emitted for a single target forward pass that, run normally, would have produced just one. As $\alpha \to 1$ the expression approaches its ceiling of $\gamma + 1$ (every drafted token plus the bonus); as $\alpha \to 0$ it approaches $1$, recovering plain autoregressive decoding with the draft as pure overhead.

The wall-clock intuition follows directly. The expensive resource is the target forward pass, and speculative decoding amortizes a single target pass over $\mathbb{E}[\text{tokens per step}]$ emitted tokens instead of one. Because a transformer can verify the $\gamma + 1$ candidate positions in one batched forward pass (the same memory-bandwidth-bound weight load serves all positions at once), the marginal cost of checking extra draft tokens is nearly free until compute, rather than bandwidth, becomes the limit. The speedup ceiling is therefore set by two things: the per-token verifier cost (which you cannot avoid paying at least once per block) and the acceptance rate (which caps how many tokens each paid pass can cover).

The assumptions worth stating plainly are that the draft and target share a tokenizer (otherwise drafted token IDs are not even comparable under verification), and that the verifier evaluates all $\gamma + 1$ positions in one batched pass (otherwise the per-position cost reappears and the speedup evaporates). The failure modes mirror the formula: when the draft and target disagree often, $\alpha$ is low and the geometric sum stays near 1, so the draft's forward passes become pure overhead; and a cheap draft only pays off when $\alpha$ is high, because for a very large target the wasted draft compute is tolerable only if most drafted tokens are accepted. This is exactly why later draft architectures such as EAGLE and Medusa, covered in Section 9.4, work to push $\alpha$ as close to 1 as possible.

The proof above is short enough to run end to end on toy distributions, which is the surest way to convince yourself the residual trick is exact rather than approximate. The following snippet implements one speculative step over a small vocabulary using numpy, drawing many proposals and tabulating the emitted tokens so the empirical histogram can be compared against the target $q$.

# One speculative-decoding step on toy distributions: accept with
# probability min(1, q/p), else resample from the residual max(0, q - p).
# We sample many steps and check the emitted histogram matches the target q.
import numpy as np

rng = np.random.default_rng(0)
q = np.array([0.5, 0.3, 0.2])   # target distribution
p = np.array([0.4, 0.5, 0.1])   # draft distribution

def speculative_step(rng):
    x = rng.choice(len(p), p=p)                 # draft proposes x ~ p
    accept_prob = min(1.0, q[x] / p[x])
    if rng.random() < accept_prob:
        return x                                # accepted draft token
    residual = np.maximum(0.0, q - p)
    residual = residual / residual.sum()        # renormalize the unmet target mass
    return rng.choice(len(p), p=residual)       # resample from the residual

emitted = np.array([speculative_step(rng) for _ in range(200_000)])
histogram = np.bincount(emitted, minlength=len(q)) / len(emitted)
print("target q   :", q)
print("empirical  :", np.round(histogram, 3))

Hugging Face Transformers exposes speculative decoding as one extra argument to generate(): pass an assistant_model and the runtime will draft and verify automatically.

# Speculative decoding via Hugging Face "assisted generation".
# The small assistant drafts; the large target verifies in a single pass.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

target_id = "meta-llama/Llama-3.1-8B-Instruct"
draft_id  = "meta-llama/Llama-3.2-1B-Instruct"   # small, same tokenizer family

tok    = AutoTokenizer.from_pretrained(target_id)
target = AutoModelForCausalLM.from_pretrained(target_id, torch_dtype=torch.bfloat16, device_map="auto")
draft  = AutoModelForCausalLM.from_pretrained(draft_id,  torch_dtype=torch.bfloat16, device_map="auto")

prompt = tok("Explain speculative decoding in one paragraph:", return_tensors="pt").to(target.device)

# assistant_model=draft enables speculative decoding inside generate()
out = target.generate(
    **prompt,
    max_new_tokens=120,
    assistant_model=draft,           # the draft / "apprentice" model
    num_assistant_tokens=5,          # how many tokens to draft per round
    do_sample=False,                 # greedy verification keeps output deterministic
)
print(tok.decode(out[0], skip_special_tokens=True))

4.3.3 Grammar-Constrained Decoding

One of the most practical advances in text generation is grammar-constrained decoding, which forces the model to produce output that conforms to a formal grammar (JSON, XML, SQL, regular expressions, or any context-free grammar). This is achieved by masking invalid tokens at the logit level before sampling or argmax.

How Grammar-Constrained Decoding Works

At each generation step, a grammar parser tracks the current state of the partially generated output. Based on this state, it computes which tokens from the vocabulary are valid continuations according to the grammar. All other tokens have their logits set to negative infinity, making them impossible to select. The model then samples or argmax over only the valid tokens.

# Using the Outlines library for structured generation
import outlines
# Define a JSON schema for the expected output
schema = """{
 "type": "object",
 "properties": {
 "name": {"type": "string"},
 "age": {"type": "integer", "minimum": 0, "maximum": 150},
 "city": {"type": "string"},
 "interests": {
 "type": "array",
 "items": {"type": "string"}
 }
 },
 "required": ["name", "age", "city"]
}"""
# Create a generator that enforces the schema
model = outlines.models.transformers("microsoft/Phi-3-mini-4k-instruct")
generator = outlines.generate.json(model, schema)
# The model MUST produce valid JSON matching the schema
prompt = "Extract person info: John Smith is a 34-year-old from Chicago who enjoys hiking and photography."
result = generator(prompt)
print(result)

The output is guaranteed to be valid JSON conforming to the schema. Without grammar constraints, a language model might produce almost-correct JSON with missing quotes, trailing commas, or type mismatches. Grammar-constrained decoding eliminates these failure modes entirely.

Tools and Libraries

A growing ecosystem of libraries implements grammar-constrained decoding. The table below provides a quick overview; each tool is covered in depth later in this section, and a comprehensive comparison table (including model support and best-fit scenarios) appears in the "Comprehensive Tool Comparison" subsection below.

# Outlines: enforce a regex constraint on generated text.
# Useful when you need a strictly formatted answer (e.g. a date, a
# version number, a phone number) without any post-processing.
from outlines import models, generate

model = models.transformers("gpt2")
# Match a US ZIP code: 5 digits, optional -4 suffix
generator = generate.regex(model, r"\d{5}(-\d{4})?")
print(generator("Customer's billing ZIP: "))
# Output is GUARANTEED to match the regex

# pip install guidance
from guidance import models, gen, select
# Load a local model (also supports OpenAI, Anthropic backends)
lm = models.Transformers("microsoft/Phi-3-mini-4k-instruct")
# Template-based structured generation
result = lm + f"""\
Analyze this review: "The battery life is amazing but the screen is dim."
Sentiment: {select(["positive", "negative", "mixed"], name="sentiment")}
Battery: {gen(name="battery_opinion", max_tokens=15, stop="\\n")}
Screen: {gen(name="screen_opinion", max_tokens=15, stop="\\n")}
Overall score: {gen(name="score", regex="[1-5]/5")}
"""
print(result["sentiment"]) # "mixed"
print(result["score"]) # "3/5"

Under the Hood: Finite-State Machines for Token Masking

Understanding how grammar-constrained decoding achieves its guarantees requires a brief excursion into formal language theory. The key insight, formalized by Willard and Louf (2023), is that regular expressions and JSON schemas can be compiled into finite-state machines (FSMs), where each state represents a position in the grammar and each transition corresponds to a set of valid tokens. At each decoding step, the system looks up the current FSM state, retrieves the set of allowed token IDs, and masks everything else to negative infinity.

For JSON schemas specifically, the compilation works as follows. The schema is first converted into a regular expression that matches all valid JSON strings conforming to the schema. That regex is then compiled into a deterministic finite automaton (DFA). Finally, the DFA transitions are mapped to vocabulary token IDs by checking which tokens can extend a partial match from each state. This mapping is precomputed once per schema, so the per-token overhead at generation time is just a lookup and a mask application.

Context-free grammars (CFGs) require a pushdown automaton rather than a simple FSM, because they need a stack to track nested structures (e.g., matching braces in JSON). Libraries like llama.cpp use GBNF (a variant of BNF notation) to specify arbitrary CFGs and enforce them during generation. The overhead is slightly higher than for regular expressions, but still negligible compared to the cost of the model forward pass itself.

LMQL: A Query Language for Constrained LLM Generation

LMQL (Language Model Query Language, developed by Beurer-Kellner et al. at ETH Zurich) takes a fundamentally different approach from token-level masking. Instead of compiling grammars into FSMs, LMQL provides a Python-embedded query language that lets you express constraints declaratively: type constraints, length limits, regex patterns, stopping conditions, and even arbitrary Python predicates. The LMQL runtime then translates these constraints into token masks at each decoding step through eager constraint evaluation, checking whether partial completions can still satisfy the declared constraints and pruning branches that cannot.

# LMQL: declarative constraints on LLM output
# pip install lmql
import lmql
@lmql.query
async def analyze_product(review: str):
 '''lmql
 argmax
 "Review: {review}\n"
 "Analyze this product review.\n"
 "Sentiment: [SENTIMENT]\n"
 "Rating: [RATING]/5\n"
 "Key issues: [ISSUES]\n"
 from
 "openai/gpt-4o-mini"
 where
 SENTIMENT in ["positive", "negative", "mixed"] and
 INT(RATING) >= 1 and INT(RATING) <= 5 and
 len(ISSUES) < 200
 '''
# The constraints are enforced at the token level during generation
result = await analyze_product("Battery dies after 2 hours but the camera is superb.")
print(result.variables)
# {"SENTIMENT": "mixed", "RATING": "3", "ISSUES": "Short battery life limits usability..."}

The key advantage of LMQL over FSM-based approaches is expressiveness. You can write arbitrary Python predicates in the where clause, not just patterns that compile to regular expressions. The tradeoff is that arbitrary predicates may require speculative evaluation (trying partial completions to see if they can still satisfy the constraint), which adds computational overhead. For simple type and choice constraints, LMQL's performance is comparable to Outlines. For complex predicates, it may be slower but is often the only practical option.

SGLang: Structured Generation with RadixAttention

SGLang (Structured Generation Language, from the LMSYS team at UC Berkeley) addresses a different bottleneck: when you run many structured generation requests that share common prefixes (system prompts, few-shot examples, shared context), the KV cache for the shared prefix is recomputed for every request. SGLang introduces RadixAttention, a radix-tree-based KV cache that automatically detects and reuses shared prefixes across requests, dramatically reducing redundant computation.

Beyond caching, SGLang provides a frontend DSL (domain-specific language) for expressing structured generation programs as Python functions. You compose primitives like gen() for unconstrained generation, select() for choosing among options, and regex() for pattern-constrained generation. These primitives are compiled into an optimized execution plan that minimizes the number of forward passes.

# SGLang: structured generation with prefix caching
# pip install sglang
import sglang as sgl
@sgl.function
def classify_and_extract(s, text):
    s += sgl.system("You are a structured data extractor.")
    s += sgl.user(f"Extract information from: {text}")
    s += sgl.assistant(
        "Category: " + sgl.gen("category", regex=r"(tech|finance|health|sports|other)") + "\n"
        + "Confidence: " + sgl.gen("confidence", regex=r"0\.\d{2}") + "\n"
        + "Summary: " + sgl.gen("summary", max_tokens=50, stop="\n")
        )
    # Run with a local model via SGLang's serving backend
    runtime = sgl.Runtime(model_path="meta-llama/Llama-3.1-8B-Instruct")
    sgl.set_default_backend(runtime)
    # Multiple calls with the same system prompt share the KV cache via RadixAttention
    texts = [
        "Apple announced record Q4 earnings driven by iPhone 16 sales.",
        "New study links sleep quality to gut microbiome diversity.",
        "Lakers secure playoff berth with overtime victory against Celtics.",
        ]
    results = [classify_and_extract.run(text=t) for t in texts]
    for r in results:
        print(r["category"], r["confidence"], r["summary"])
        runtime.shutdown()

RadixAttention's payoff is easiest to see in terms of compute. Standard serving recomputes the prefill cost $O(L_{\text{prefix}}^{2} \cdot d)$ for every request of context length $L_{\text{prefix}}$. With prefix sharing across $N$ requests that share the first $L_s$ tokens, the amortized prefill cost per request becomes:

$$ C_{\text{shared}} \;=\; \underbrace{\frac{O(L_s^{2} \cdot d)}{N}}_{\text{shared prefix paid once}} \;+\; O\bigl((L_{\text{prefix}} - L_s)^{2} \cdot d\bigr) $$

When $L_s$ dominates $L_{\text{prefix}}$ and $N$ is large (e.g., a 4K system prompt reused across thousands of API calls), the first term collapses toward zero and total throughput rises proportionally.

jsonformer: Lightweight Structured Generation

For simpler use cases, jsonformer takes a minimalist approach to structured JSON generation. Rather than masking logits via an FSM, jsonformer generates JSON structurally: it produces the JSON skeleton (braces, brackets, colons, commas, key names) deterministically, and only invokes the LLM to fill in the values. This means the LLM never generates structural tokens at all, eliminating the possibility of malformed JSON without any logit manipulation. The tradeoff is that jsonformer only supports JSON (not arbitrary grammars) and requires the schema to be known at generation time.

Comprehensive Tool Comparison

Performance Implications of Constrained Decoding

Grammar-constrained decoding is not free. Several sources of overhead deserve consideration when deciding whether to use it in production.

Precomputation cost. Compiling a JSON schema or regex into an FSM and mapping its transitions to vocabulary token IDs is a one-time cost, but it can be significant for large vocabularies (50,000+ tokens) and complex schemas. Outlines caches these compiled masks, so the cost is amortized across requests using the same schema. For a typical JSON schema with 10 to 20 fields, precomputation takes 1 to 5 seconds; for deeply nested schemas it can take longer.

Per-token overhead. At each decoding step, the system must look up the current FSM state and apply the token mask. This is a fast operation (microseconds) compared to the model forward pass (milliseconds), so it adds less than 1% latency in practice. The exception is LMQL with complex Python predicates, where constraint evaluation may take longer.

Batching challenges. When serving multiple requests with different schemas in a batch, each request needs its own FSM state and token mask. This means the logit masking step cannot be fully vectorized across the batch, which can reduce throughput on GPU-based serving systems. SGLang and vLLM handle this by applying masks per-request within a continuous batch, but the overhead scales linearly with the number of distinct schemas in a batch.

Output quality. Constraining the token space can occasionally degrade output quality if the schema is very restrictive and forces the model into low-probability regions of its distribution. For example, if the schema requires an integer between 1 and 5 but the model's probability mass is spread across textual representations ("one", "two") rather than digits, the constraint forces selection from low-probability digit tokens. In practice, this is rarely a problem for well-prompted models, but it is worth monitoring.

4.3.4 Watermarking Generated Text

As LLM-generated text becomes more prevalent, the ability to detect whether text was produced by an AI is increasingly important. Watermarking embeds a statistical signal into the generated text that is invisible to humans but detectable by an algorithm.

The Kirchenbauer et al. (2023) Method

The most influential watermarking approach works as follows:

1. At each generation step, use the previous token as a seed to a hash function, partitioning the vocabulary into a "green list" and a "red list"
2. Add a small bias δ to the logits of all green-list tokens before sampling
3. This nudges (but does not force) the model to prefer green-list tokens

To detect the watermark, apply the same hash function to a piece of text and count how many tokens fall on the green list. Unwatermarked text will have roughly 50% green-list tokens (random chance). Watermarked text will have significantly more, detectable via a simple z-test.

# Text watermarking: hash the previous token to split the vocabulary into
# "green" and "red" lists, then bias logits toward green tokens by delta.
import torch
import hashlib
def watermark_logits(logits, prev_token_id, vocab_size, delta=2.0, gamma=0.5):
    """Apply watermark bias to logits based on previous token."""
    # Use previous token to seed the green/red list partition
    seed = hashlib.sha256(str(prev_token_id).encode()).hexdigest()
    rng = torch.Generator()
    rng.manual_seed(int(seed[:8], 16))
    # Random permutation determines green list (first gamma fraction)
    perm = torch.randperm(vocab_size, generator=rng)
    green_list_size = int(gamma * vocab_size)
    green_tokens = perm[:green_list_size]
    # Add bias to green-list tokens
    watermarked_logits = logits.clone()
    watermarked_logits[green_tokens] += delta
    return watermarked_logits
def detect_watermark(token_ids, vocab_size, gamma=0.5):
    """Detect watermark by counting green-list tokens."""
    green_count = 0
    total = len(token_ids) - 1 # skip first token (no previous)
    for i in range(1, len(token_ids)):
        prev_id = token_ids[i - 1]
        seed = hashlib.sha256(str(prev_id).encode()).hexdigest()
        rng = torch.Generator()
        rng.manual_seed(int(seed[:8], 16))
        perm = torch.randperm(vocab_size, generator=rng)
        green_set = set(perm[:int(gamma * vocab_size)].tolist())
        if token_ids[i] in green_set:
            green_count += 1
            green_fraction = green_count / total
            # Z-test: under null hypothesis, green fraction ~ gamma
            import math
            z_score = (green_fraction - gamma) / math.sqrt(gamma * (1 - gamma) / total)
            return green_fraction, z_score
            # Simulate detection
            print("Watermarked text: green_frac=0.78, z_score=5.6 => WATERMARKED")
            print("Human-written text: green_frac=0.51, z_score=0.2 => NOT watermarked")

4.3.5 Minimum Bayes Risk (MBR) Decoding

MBR decoding takes a fundamentally different approach from the methods discussed so far. Instead of using a single decoding strategy to produce one output, MBR generates multiple candidate outputs and then selects the best one according to a quality metric. Think of it like an election: each candidate "votes" for the others by scoring their similarity, and the candidate with the highest total votes wins. In other words, MBR finds the "center of gravity" among a set of diverse outputs.

The MBR Selection Algorithm

1. Sample N candidates: Generate N different outputs using any stochastic sampling method (e.g., temperature sampling with T=0.8)
2. Score each candidate: For each candidate, compute its average "utility" against all other candidates using a metric (ROUGE, BERTScore, or an LLM judge)
3. Select the best: Return the candidate with the highest average utility

The intuition is that the best candidate is the one that is most "central" among the samples: it is the output that other good outputs most agree with. This is more robust than picking the highest-probability output, which might be an outlier.

# Minimum Bayes Risk decoding: generate N candidates, score each by
# average pairwise utility (word overlap here), and pick the consensus winner.
import numpy as np
def mbr_decode(candidates, utility_fn):
    """Select the candidate with highest average utility against all others."""
    n = len(candidates)
    scores = np.zeros(n)
    for i in range(n):
        for j in range(n):
            if i != j:
                scores[i] += utility_fn(candidates[i], candidates[j])
                scores[i] /= (n - 1)
                best_idx = np.argmax(scores)
                return candidates[best_idx], scores

# Example with simple word-overlap utility
def word_overlap(a, b):
    """Simple utility: fraction of words in a that also appear in b."""
    words_a = set(a.lower().split())
    words_b = set(b.lower().split())
    if not words_a:
        return 0.0
        return len(words_a & words_b) / len(words_a)
        candidates = [
            "The cat sat quietly on the warm mat.",
            "A cat was sitting on a mat in the sun.",
            "The cat sat on the mat.",
            "Purple elephants danced wildly in space.", # outlier
            "The cat rested on the warm mat nearby.",
            ]
        best, scores = mbr_decode(candidates, word_overlap)
        print("MBR Scores:")
        for i, (c, s) in enumerate(zip(candidates, scores)):
            marker = " <-- BEST" if c == best else ""
            print(f" [{s:.3f}] {c}{marker}")

The MBR selection correctly identifies the most "central" candidate ("The cat sat on the mat.") and rejects the outlier. In practice, using BERTScore or an LLM-as-judge for the utility function produces much stronger results than simple word overlap. The main cost is computational: N samples times N utility evaluations gives O(N²) cost, though in practice N values of 10 to 50 offer a good tradeoff between quality and compute.

4.3.6 Structured Output via JSON Schema

Grammar-constrained decoding (Section 3 above) provides the foundation, but the most widespread application of this technique in production is structured JSON output: constraining the model to produce valid JSON that conforms to a specific schema. This pattern has become so important that every major LLM provider now offers it as a first-class feature.

The motivation is straightforward. When an LLM is part of a larger software system, its output must be machine-readable. Prompt engineering alone (e.g., "respond in JSON format") works most of the time, but "most of the time" is not good enough for production. A 2% malformed-output rate across millions of API calls creates thousands of failures per day. Schema-constrained generation eliminates this failure mode entirely by making it structurally impossible for the model to produce invalid output.

4.3.6.1 How Providers Implement It

Each major provider takes a slightly different approach, but the core idea is the same: the JSON schema is converted into a state machine (typically a finite-state automaton or pushdown automaton), and at each decoding step, only tokens that are valid continuations according to the current parser state are allowed.

4.3.6.2 The Instructor Library: A Practical Interface

The Instructor library (by Jason Liu) has become the de facto standard for extracting structured data from LLMs in Python. It wraps provider APIs and uses Pydantic models to define the expected output schema, then leverages each provider's native structured-output capability to guarantee compliance. This gives you type-safe, validated Python objects directly from LLM calls.

# Structured output with Instructor + Pydantic: force the LLM to return
# JSON matching a SentimentResult schema instead of free-form text.
import instructor
from pydantic import BaseModel, Field
from openai import OpenAI

# Define the desired output structure with Pydantic
class SentimentResult(BaseModel):
    sentiment: str = Field(description="positive, negative, or neutral")
    confidence: float = Field(ge=0.0, le=1.0)
    reasoning: str = Field(description="Brief explanation of the sentiment")

def classify_sentiment(text: str) -> SentimentResult:
    """One Instructor call returns a validated SentimentResult."""
    client = instructor.from_openai(OpenAI())
    return client.chat.completions.create(
        model="gpt-4o",
        response_model=SentimentResult,
        messages=[{"role": "user", "content": f"Analyze: '{text}'"}],
    )

result = classify_sentiment("The new update broke my workflow entirely.")
print(result.sentiment)   # "negative"
print(result.confidence)  # 0.92
print(result.reasoning)   # "The user expresses frustration about a broken workflow"

4.3.6.3 Function Calling as Constrained Decoding

Function calling (also called "tool use") is the pattern where an LLM decides to invoke an external function and generates the appropriate arguments. From a decoding perspective, function calling is simply another form of constrained generation. When the model decides to call a tool, it must produce a JSON object whose structure matches the tool's parameter schema. The same grammar-constrained decoding machinery that enforces JSON schema compliance is used to ensure the function call arguments are valid.

In practice, the provider converts the list of available tools into a combined schema, and the model's generation is constrained to either produce a regular text response or a valid function call matching one of the tool definitions. This is why function calling is so reliable in modern APIs: it is not just the model "trying" to output correct JSON through next-token prediction; the decoding infrastructure guarantees that the output conforms to the tool schema at every token.

The connection between structured output and function calling reveals a deeper pattern: both are instances of schema-constrained generation, where a formal specification governs the space of valid outputs. As LLMs become integrated into more software systems, this pattern of treating generation as a constrained optimization problem (rather than free-form text completion) will only grow more important. We will explore tool use and function calling extensively in Section 26.1.

Exercises