Provably Tractable NFA-Constrained Language Generation via HMMs
Listen
Radio episode about this paper
Transcript
Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.
Tom: Today's paper: "Provably Tractable NFA-Constrained Language Generation via HMMs".
Jane: Constrained generation aims to sample from language models conditioned on hard constraints, and this paper introduces NFA-LM, a polynomial-time engine for NFA-constrained generation with theoretical guarantees under mild assumptions.
Tom: First, who's behind it and why it matters.
Title and authors: Tom: So, we’re looking at the title "Provably Tractable NFA-Constrained Language Generation via HMMs," which tells us right away that they are promising a way to handle language model generation under NFA rules in a way that has mathematical proof of performance bounds. It’s not just another tweak; it suggests a solid framework for this kind of constrained sampling.
Jane: Exactly, Tom. The authors Jialiang Sun and Kuldeep Meel are tackling the challenge head-on by suggesting they can achieve this with a polynomial-time engine, which means the computation scales reasonably well as the language model gets bigger or the constraints get more intricate.
Lu: The real implication here is how they generalize existing work; they take concepts from efficient randomized approximation schemes for #NFA and extend them to this specific constrained generation setting, which is a significant theoretical extension of prior research.
Meng: I’m thinking about the practical side—when we talk about polynomial time, we’re talking about something that can actually run on real hardware within reasonable time limits, unlike methods that might just be fast in theory but crawl in practice.
Lalam: For me, this title implies a future where complex generative tasks aren't just approximations; they are provably constrained generation processes, which means we gain much more confidence when deploying these models for specific applications.
The paper's summary: Tom: To summarize, the core of this paper is proposing NFA-LM, which first distills a Hidden Markov Model from the base language model and then uses that HMM to assign tractable probability weights to future continuations based on those NFA constraints. It’s like creating a shortcut for calculating probabilities without having to count every single possible path through the constraints.
Jane: That distillation step is key, Tom; by using an HMM as a proposal, they are simplifying the complex task of estimating the conditional probability into something that is much easier to handle mathematically and computationally. They then use this HMM structure to reweight the next-token distribution according to what the NFA allows.
Lu: The mechanism they use involves defining unrolled NFAs, Au, which resembles a directed acyclic graph with layers of state copies, and then they introduce canonical runs to count every suffix exactly once by fixing a total order on the states. That's how they make the counting tractable in the first place.
Meng: That idea of organizing the structure into these layers and defining those canonical runs seems very structured, which is good for engineering because it gives us a clear path for implementation rather than just an abstract mathematical concept.
Lalam: I think what’s most important to grasp is that they are moving away from methods that just distort the distribution; this approach aims to approximate the true conditional distribution of the language model given those hard rules, which is much more faithful.
The paper's improvements: Tom: The authors introduce several key components for this engine, including a distillation step to get an HMM, a precomputation stage that generates reusable suffix samples based on certain weights p j(q, b), and then the main generation process that estimates the HMM completion weight using those precomputed results.
Jane: The improvements focus on making the system modular: you distill first, then you precompute for reuse, and finally, you generate by combining those two parts to get a tractable estimate of how likely a sequence is under the constraints. This separation helps manage the complexity across different stages of generation.
Lu: A crucial part they highlight is defining canonical product runs on the unrolled NFA Au and analyzing transitions between product states using a weight function psi(a, b, b'), which sets up the machinery for counting those sequences efficiently.
Meng: From an implementation view, I’m interested in how much memory this precomputation stage adds. If we can reuse those HMM-weighted suffix samples S r,j effectively, that could drastically reduce the run-time cost during the actual generation phase.
Lalam: The theoretical guarantees they provide are what really sell this up; specifically Theorem three establishes a time complexity of O(n 2m 3h two epsilon-two (nmh) (delta-one)), which gives us a concrete measure of how much compute we can expect.
Conclusion: Tom: So we’ve seen how this NFA-LM framework uses HMM distillation and precomputation to tackle the counting problem associated with constrained generation in a way that offers polynomial time complexity and theoretical error bounds. It seems like a solid path forward for handling complex structural requirements in language models.
Jane: I agree, Tom; the main implication is that we can achieve distribution-aware constrained generation where the error is bounded by n epsilon with high probability, which means the generated text will respect the rules while still sounding very natural. It moves us past just getting an answer to getting a reliable guarantee about that answer's quality.
Lu: I see huge potential here for applying this to generating highly specific, structured content where NFA constraints are essential; it’s a powerful tool for formalizing complex linguistic requirements that were previously too hard to handle effectively.
Meng: Practically, the polynomial time complexity and the ability to quantify the approximation error give us a lot of assurance when we start building systems that rely on these constrained outputs. We can design our architectures knowing there are mathematical limits on how far off we might be from the true distribution.
Lalam: Ultimately, this work points toward a future where AI systems can generate content that is not only fluent but also provably compliant with intricate rules, which will definitely make our AI more useful across many professional domains.
University of Toronto
cs.CL, cs.FL
Submitted: 2026-09-30
Updated: 2026-10-06
License: http://creativecommons.org/licenses/by/4.0/
Importance score: 90/100
The gist: Constrained generation aims to sample from language models conditioned on hard constraints, and this paper introduces NFA-LM, a polynomial-time engine for NFA-constrained generation with theoretical
Key concepts
- NFA
- A Non-deterministic Finite Automaton is a mathematical model used to recognize patterns in text, essentially defining the hard constraints for language generation. It is efficiently represented by an unrolled NFA structure that helps manage the complexity of the constraints.
- HMM Phmm(α | x1:ℓ)
- This is a tractable Hidden Markov Model approximation of the true conditional probability distribution Plm(α | x1:ℓ). Instead of calculating the exact, intractable probability, this HMM acts as a simplified model used to guide the generation process efficiently.
- Precomputation
- This stage involves generating and storing HMM-weighted suffix samples. These precomputed results are crucial because they allow the generation phase to quickly estimate how likely certain sequences are under the constraints, making the overall process faster.
- Total Variation (TV) distance
- This is a metric used to measure how close an approximate probability distribution is to the true target distribution. The paper proves that NFA-LM can generate sequences whose distributions are within a bounded TV distance of the actual constrained distribution.
Terminology
Summary
Constrained generation aims to sample from language models conditioned on hard constraints, and this paper introduces NFA-LM, a polynomial-time engine for NFA-constrained generation with theoretical guarantees under mild assumptions.
The gist: NFA-LM is a polynomial-time engine for NFA-constrained generation with theoretical guarantees under mild assumptions.
Problem Context and Motivation
Constrained generation requires sampling from language models (LMs) conditioned on hard constraints, which is difficult because computing the constraint-conditioned probability, Plm(α x1:l), involves marginalizing over future continuations, leading to a counting problem equivalent to the intractable and computationally expensive exact counting of sequences accepted by an NFA. Existing techniques fall into two categories: distribution-agnostic methods that distort the true conditional distribution, and distribution-aware methods (like Langevin-based or MCMC samplers) that lack polynomial-time guarantees for approximating the constrained LM distribution to a prescribed accuracy. The paper addresses the question of whether efficient, distribution-aware constrained generation for NFAs with theoretically bounded error is achievable.
Theoretical Foundation: NFA and HMMs
The framework is built upon the connection between regular expressions and NFAs, which are efficiently represented by NFAs with a linear number of states. The paper introduces an unrolled NFA, Au, which resembles a directed acyclic graph (DAG) with layers of state copies Ql for l = 0 to n. The core idea is to distill a Hidden Markov Model (HMM), Phmm(α x1:l), as a tractable proposal that approximates the true conditional probability Plm(α x1:l). This HMM is specified by an initial distribution vector, an emission matrix, and a transition matrix.
The NFA-LM Framework
NFA-LM operates in three stages: (i) Distillation, (ii) Precomputation, and (iii) Generation. The distillation step involves distilling the HMM Phmm using responses sampled from the base LM Plm. The precomputation stage is central to generating reusable information; it precomputes HMM-weighted suffix samples S r,j (q, b) at rates p j(q, b) such that for every atom z ∈ Z+(q, b), the invariant P[z ∈ S r,j (q, b)] = p j(q, b)w(z, b) holds.
Canonical Runs and Product States
The paper defines canonical runs on the unrolled NFA Au to count every suffix once by fixing a total order ≺ on the states. This leads to the definition of a canonical product run
for an atom z, which is defined inductively based on its successor state. A key concept is the convergence state,
which is reached by the longest common suffix of two canonical runs, denoted v z1,z2. The framework then defines product states v = (q, b) as pairs of NFA states and HMM hidden states, and analyzes transitions between these product states using a weight function ψ(a, b, b′).
Approximation Guarantees and Complexity
The theoretical guarantees are established through a series of lemmas leading to Theorem 1. The algorithm uses an outer repetition structure (nu repetitions) to amplify success probability. Key results include:
-
Lemma 1: Bounding the probability of approximation failure, P[A j], by at most 1/16 for each core repetition j.
-
Lemma 3: Bounding the probability of sample-overflow failure, P[B j ∩ (A j) c], by at most 1/16.
-
Theorem 2: Bounding the Total Variation (TV) distance between the approximate and exact distributions as DTV(Pˆ(· α), P(· α)) ≤ nε, with probability at least 1 − δ.
-
Theorem 3: Establishing the time complexity of Algorithm 1 as O(Σn 2m 3h 2ε-2 log(nmh) log(δ-1)).
Constrained Generation Process
During generation (Algorithm 4), NFA-LM estimates the HMM completion weight Phmm(α x1:l) by combining precomputed results associated with the reachable states R. For every candidate token a ∈ Σ, it samples a with probability proportional to Plm(a x1:l−1)Pˆhmm(α x1:l−1 · a). This process is repeated for n steps, ensuring that every returned sequence satisfies the NFA constraint (Lemma 13). The overall running time combines the precomputation and generation costs.
Improvements for AI systems
Here are the specific improvements that can be made to AI systems based on the NFA-LM paper:
-
Enhance Text Detoxification: The system can reliably generate text that adheres strictly to complex structural constraints (modeled by NFAs/Regex), ensuring outputs meet safety or formatting requirements (e.g.,
must contain concept A and must not contain concept B
). -
Improve Agentic Function Calling: The system can generate structured, constrained code or function calls where the structure is defined by a formal regular expression, leading to higher success rates in agentic workflows involving external APIs or tools.
-
Generate Structured SQL/Code: The system can produce syntactically correct database queries that satisfy complex structural patterns (like specific joins or required keywords) without relying solely on brittle prompt engineering, improving the reliability of LLMs for code generation tasks.
-
Enable Distribution-Aware Constraint Satisfaction: Unlike existing distribution-agnostic methods that distort the probability, NFA-LM can generate outputs that are provably close to the true conditional distribution of a language model given hard constraints (within a theoretical error bound of 10%). This leads to more natural, high-quality text while still respecting complex rules.
-
Achieve Provable Error Bounds: The system provides theoretical guarantees on both approximation error (bounded by ε) and time complexity. This allows developers to trust the output quality and performance characteristics of the constrained generation process under specified parameters (ε, δ).
-
Scale to Complex Constraints: The framework can handle constraints that require a large number of NFA states (up to 50 in experiments), effectively allowing LMs to manage more complex linguistic rules than previously possible with simpler constraint methods.
These improvements enable the AI system to move beyond simple pattern matching or brittle prompt engineering toward a robust, theoretically grounded mechanism for generating high-quality, structurally compliant content.
Abstract
Constrained generation aims to sample from language models (LMs) conditioned on hard constraints. Existing constrained-generation techniques for nondeterministic finite automaton (NFA) constraints either distort the distribution or sacrifice efficiency. Theoretically, this task reduces to counting the length- n sequences accepted by an NFA (#NFA), and the exact #NFA problem is #P-complete. Recent work has shown that #NFA admits a fully polynomial randomized approximation scheme (FPRAS). Inspired by this result, we propose NFA-LM, a polynomial-time engine for NFA-constrained generation with theoretical guarantees under mild assumptions. Experiments show that NFA-LM efficiently generates high-quality outputs with theoretically bounded approximation error.
Sources
- Mitigating Bias in Locally Constrained Decoding via Tractable Proposals
- JSONSchemaBench: A Rigorous Benchmark of Structured Outputs for Language Models
- Synchromesh: Reliable code generation from pre-trained language models
- Evaluating Large Language Models on Controlled Generation Tasks
- Gemma 4 Technical Report
- Qwen3.5-Omni Technical Report
- TRACE Back from the Future: A Probabilistic Reasoning Approach to Controllable Language Generation
Related papers
- Exploring Solution Divergence and Its Effect on Large Language Model Problem Solving
- Ishigaki-IDS-Bench: A Benchmark for Generating Information Delivery Specification from BIM Information Requirements
- Subliminal Steering: Stronger Encoding of Hidden Signals
- MedStruct-S: A Benchmark for Key Discovery, Key-Conditioned QA and Semi-Structured Extraction from OCR Clinical Reports
- The End of Transformers? On Challenging Attention and the Rise of Sub-Quadratic Architectures
- Untangling the Mechanisms of Misleading Context in Medical Question Answering