Constrained decoding usually means masking invalid next tokens so a language model can only emit strings that satisfy a grammar, schema, or vocabulary. For small schemas this is often cheap: a parser or finite-state machine computes the legal token set, the sampler masks the logits, and decoding continues. Generative retrieval changes the scale. Instead of choosing from a handful of JSON fields, the model may need to output one valid semantic identifier from millions of candidate item IDs while running beam search on a GPU or TPU.
STATIC, short for Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding, addresses that accelerator bottleneck. The core idea is to keep the trie semantics but remove the pointer-chasing implementation. A set of valid token sequences is compiled offline into a hybrid dense-plus-sparse representation: dense lookup tables for the high-fanout early layers, and a compressed sparse row transition matrix for deeper levels. At runtime the decoder uses vectorized kernels to retrieve valid continuations for every beam in parallel.
The technique matters whenever constrained decoding is part of the hot inference path. In large recommender systems, business rules such as freshness, inventory availability, category filters, or regional eligibility can make unconstrained item generation unsafe. STATIC shows how to enforce those rules without round-tripping through CPU-side trie traversal or compiling data-dependent branching into an accelerator graph.
In generative retrieval, each retrievable item is represented by a semantic ID: a short sequence of discrete tokens. The model retrieves by autoregressively generating the semantic ID instead of searching an embedding index directly. This makes retrieval compatible with LLM-style decoding machinery, but it also creates a validity problem. The model can generate token sequences that do not correspond to any item, or that correspond to items outside the currently allowed constraint set.
A prefix tree is the natural data structure for valid token sequences. Given the prefix already generated, the current trie node tells the decoder which tokens may appear next. Traditional tries are poor matches for GPUs and TPUs because child pointers scatter memory reads, branch counts vary by node, and traversal depends on runtime data. Those properties fight memory coalescing and static graph compilation.
STATIC uses a hybrid representation. The first few trie levels typically have the largest branching factor, so they are handled with dense bit-packed masks and state tables. Deeper levels are represented as a CSR matrix, which stores non-empty transitions compactly while still allowing fixed-shape vectorized reads. The result preserves exact trie constraints while moving the work into accelerator-native array operations.
STATIC treats the valid-output set as data that can be compiled into tensors before inference. The online decoder does not walk a recursive object graph. It indexes precomputed tensors, builds a valid-token mask, and applies that mask to model logits before selecting the next token.
The offline phase starts with the set of valid token sequences. In a recommender, those may be semantic IDs for currently eligible items. The builder analyzes prefix structure, assigns trie states, emits first-token masks, materializes dense lookup tables for the configured number of early layers, and packs the remaining transitions into CSR arrays.
def build_static_index(valid_sequences, dense_layers=2):
trie = build_prefix_tree(valid_sequences)
start_mask = tokens_allowed_at_root(trie)
dense_mask, dense_states = build_dense_prefix_tables(
trie,
depth=dense_layers,
)
csr_indices, csr_indptr, csr_next_state = pack_sparse_tail_as_csr(
trie,
start_depth=dense_layers,
)
return StaticIndex(
start_mask=start_mask,
dense_mask=dense_mask,
dense_states=dense_states,
csr_indices=csr_indices,
csr_indptr=csr_indptr,
csr_next_state=csr_next_state,
dense_layers=dense_layers,
)During decoding, the model produces logits for each active beam. STATIC computes the legal next tokens for each beam prefix and masks all other logits before sampling or beam expansion. For early decoding steps this is a dense table lookup. For later steps a vectorized transition kernel reads a fixed-size slice of sparse transitions for all beams and constructs the masks in parallel.
def constrained_step(model, index, beam_prefixes, beam_states):
logits = model.next_token_logits(beam_prefixes)
step = beam_prefixes.shape[1]
if step == 0:
valid_mask = index.start_mask
elif step <= index.dense_layers:
valid_mask, beam_states = dense_transition(
index.dense_mask,
index.dense_states,
beam_states,
beam_prefixes[:, -1],
)
else:
valid_mask, beam_states = sparse_transition_kernel(
index.csr_indices,
index.csr_indptr,
index.csr_next_state,
beam_states,
beam_prefixes[:, -1],
)
masked_logits = where(valid_mask, logits, -float("inf"))
return select_next_beams(masked_logits, beam_states)CSR is a standard sparse-matrix layout: one array stores column indices, another stores row offsets, and optional value arrays store transition payloads. For STATIC, a row corresponds to a trie state and entries correspond to valid child-token transitions. The decoder can issue burst reads from contiguous arrays instead of chasing object pointers. That turns the masking step into predictable tensor work that fits GPU and TPU execution much better.
A video or product recommender may need to retrieve only items that are fresh, in stock, safe for a region, or eligible for a particular user. STATIC can compile the allowed semantic IDs into a constraint index so the model cannot generate invalid identifiers. The source paper reports a production deployment for enforcing a recent-video constraint in a large-scale video recommendation setting.
Generative retrieval models may struggle with items that were not frequent or present during training. Constraining decoding to a cold-start candidate set lets the system ask the model to choose among new valid items rather than relying on unconstrained generation. The constraint layer does not make the model understand the new item by itself, but it prevents impossible outputs and focuses probability mass on the current candidate pool.
The same pattern applies beyond recommenders when valid outputs are large token vocabularies: catalog IDs, ontology codes, route names, product SKUs, or database object identifiers. If the allowed strings are enumerable and latency-sensitive, compiling them into accelerator-friendly transition tensors can be more reliable than asking a CPU-side parser to participate in every decoding step.
Build the STATIC index outside the per-token inference loop. If constraints change frequently, maintain versioned indexes and switch requests to the newest ready version atomically. The decoder should see an immutable tensor bundle for the lifetime of a request; rebuilding the index in the hot path defeats the purpose of vectorized masking.
class ConstraintIndexRegistry:
def __init__(self):
self.current = None
def refresh(self, eligible_item_ids):
sequences = encode_semantic_ids(eligible_item_ids)
next_index = build_static_index(sequences, dense_layers=2)
self.current = next_index
def decode(self, model, prompt, beam_size):
index = self.current
if index is None:
raise RuntimeError("constraint index is not ready")
return beam_search(model, prompt, beam_size, constraint_index=index)Dense tables make early lookups fast but consume memory proportional to the combinations they cover. Sparse CSR tables save memory but require a transition kernel. A common design is to use dense tables only for the first one or two levels, where fanout is highest and every request pays the cost, then shift the rest of the trie into CSR.
The constraint index should depend on tokenized valid sequences, not on a particular request or model invocation. This allows reuse across batches and beams. It also makes correctness tests straightforward: every generated sequence should be accepted by the original valid-sequence set, and every prefix state should produce exactly the child tokens represented in the trie.