cs.LGOct 2, 2026

Constraint-Aware Training

Authors: Jinwoo Kim

Organizations: Department of Computer Science and Engineering University of California-San Diego La Jolla, CA

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

When generating programs with language models, constrained decoding can apply program analyses to exclude tokens that violate syntax, scope, or typing rules. However, there is a duplication: standard training already teaches the model to suppress the tokens rejected by these analyses. This duplication leads to the question: if we will perform some analysis to filter a set tokens out during inference anyways, can we avoid teaching the model the said analysis altogether during training, and does this externalization lead to more efficient models? This paper defines a general constraint-aware objective satisfying this externalization desideratum and formalizes the benefits of externalization into three concrete theorems about model size and data efficiency. We show, through a controlled synthetic experiment, that the theorems survive training dynamics: constraint-aware training yields lower prediction loss at a matched parameter count and data compared to ordinary cross-entropy training, motivating training objectives that incorporate the analyses used during generation.

Figures & tables

Figure 1: Standard (CE) and constraint-aware (CA) loss compared on three different scenarios on each row. In each row, mask MM denotes the set of tokens allowed (marked with T\mathsf{T} ) according to the analysis, tt the one-hot target, pp the model prediction, and pMp_{M} the model prediction renormalized on tokens allowed by MM . (a) With only the gold x allowed, CA loss is zero. (b) When the mask MM allows both x and y , CA penalizes assigning probability to the non-gold y . (c) Changing the mass on the rejected token (vs. (b)) without changing the x:y\texttt{x}:\texttt{y} ratio changes CE but leaves CA unchanged.
Figure 2: Mean losses ℓCA\ell_{\mathrm{CA}} , ℓCESA\ell_{\mathrm{CE}}^{{S_{A}}} , and ℓCE\ell_{\mathrm{CE}} , each reported in excess of the target distribution’s entropy. All values are nats per answer.
Figure 3: Left: excess losses obtained by training on the depth corpus, associated with ℓCE\ell_{\mathrm{CE}} , ℓCESA\ell_{\mathrm{CE}}^{{S_{A}}} (from CE training), and ℓCA\ell_{\mathrm{CA}} (from CA training). All losses are nats per answer. Right: CE excess loss by formula size for 2, 4, and 8 depth configurations.
Figure 4: Excess loss curves against corpus size (left), and loss of the mean prediction across all 16 models ( ℓmean\ell_{\mathrm{mean}} ) and the average variance component of individual losses ( ℓvar\ell_{\mathrm{var}} ) across corpus sizes (right). Units are in nats. CA achieves a significantly lower excess at all corpus sizes. For both CE and CA, ℓvar\ell_{\mathrm{var}} is much higher then ℓmean\ell_{\mathrm{mean}} , indicating that variance from the training corpus is a significant source of loss.
Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

DepthSeedℓCA\ell_{\mathrm{CA}} ( 10−510^{-5} )ℓCESA\ell_{\mathrm{CE}}^{{S_{A}}} ( 10−310^{-3} )ℓCE\ell_{\mathrm{CE}}CE range
203.0081.6391.10364
212.3821.5361.10063
403.2831.9720.83185
412.4512.2360.85525
803.1781.7060.47648
812.4501.4940.58617
Appendix
Table 1: All six matched two-stage pairs, evaluated on the same uniform-size panel. Losses are nats per answer; the two within-mask columns use different scales. The last column is the largest operator count up to which CE meets the 95%-accuracy criterion at every size. No CA model meets this criterion even at the first evaluated size.
Figure 5: Full-size CE curves for the same two-stage seed-0 models. Left: total excess. Right: query-weighted accuracy of identifying the allowed set. The main figure averages the two seeds and shows loss by size only for operator counts 2–9.
Pairs (M)ℓCA\ell_{\mathrm{CA}}ℓCESA\ell_{\mathrm{CE}}^{{S_{A}}}CE ℓA\ell_{\mathrm{A}}ℓCE\ell_{\mathrm{CE}}
2.40.000379±0.0000800.000379\pm 0.0000800.028276±0.0029980.028276\pm 0.0029980.0174900.045766
7.20.000080±0.0000050.000080\pm 0.0000050.011467±0.0055260.011467\pm 0.0055260.0179690.029436
240.000075±0.0000040.000075\pm 0.0000040.007973±0.0039280.007973\pm 0.0039280.0159020.023875
360.000074±0.0000030.000074\pm 0.0000030.007268±0.0027130.007268\pm 0.0027130.0152190.022487
Appendix
Table 2: Primary data study over sixteen independent training corpora at each pool size, with fixed initialization and order seed. Masked excess is mean ±\pm sample standard deviation; other entries are means. CE masked excess plus support loss equals CE total before rounding. Pool sizes are in millions of training pairs; losses are held-out nats per query.
Pairs (M)ObjectivePeak LRReadout decayMasked excessCE total
2.4CA0.01000.000155±0.0000130.000155\pm 0.000013–
2.4CE0.0030.10.029823±0.0053660.029823\pm 0.0053660.049277
36CA0.00100.000060±0.0000040.000060\pm 0.000004–
36CE0.0030.10.009372±0.0024890.009372\pm 0.0024890.027421
Appendix
Table 3: Independent selection for masked deployment. Each objective uses the same six-recipe search and two tuning corpora per size. Losses are means over four new evaluation corpora; the masked-loss column also gives the sample standard deviation. Pool sizes are in millions of training pairs; all losses are excess in nats per query. The CE total column describes the same masked-loss-selected CE model.
CellObjectiveMasked excess ( 10−310^{-3} nats)ℓvar\ell_{\mathrm{var}} ( 10−310^{-3} nats)
ACA0.074±0.0030.074\pm 0.0030.0664
ACE7.268±2.7137.268\pm 2.7136.5800
BCA0.047±0.0020.047\pm 0.0020.0414
BCE4.899±5.3654.899\pm 5.3654.4032
CCA0.046±0.0020.046\pm 0.0020.0410
CCE5.947±6.7085.947\pm 6.7085.3315
Appendix
Table 4: Larger-pool and longer-schedule study. Within-mask excess is mean ±\pm sample standard deviation over 16 corpora. ℓvar\ell_{\mathrm{var}} is the average loss removed by averaging their masked predictions. Both numerical columns use units of 10−310^{-3} nats.
More in this topic

Constrained Generative Modeling

Constrained generative modeling develops methods that produce samples satisfying structural, physical, semantic, or user-defined requirements.

See all 187 papers

Explore similar work

CardsList
Oct 26, 2026

Constrained Decoding: Harmful or Beneficial? Revisiting Its Impact on LLM Performance

Constrained decoding guarantees that Large Language Models (LLMs) produce outputs conforming to predefined formats, but is often assumed to inherently degrade reasoning. We examine this assumption through a taxonomy of three constraint types: format constraints, which enforce output structure; validity constraints, which mask invalid answer; and reasoning constraints, which restrict reasoning trajectories. We theoretically show that appropriate constraints can preserve or improve performance by eliminating invalid paths. Across four model families and eleven open-source models, causal ablations on diverse benchmarks using state-of-the-art constrained decoding strategies reveal that apparent performance degradation largely stems from non-equivalent prompting and output formats rather than from decoder-level masking. Meanwhile, validity constraints improve reliability and extractability by eliminating invalid answers, while reasoning constraints can improve reasoning and, consequently, accuracy when they enforce beneficial trajectories. These findings show that the impact of constrained decoding depends on the type of constraint and the output space it removes. Our code is available.

Ngoc Trinh Hung Nguyen, Alonso Silva, Armen Aghasaryan +1

Télécom Paris, Institut Polytechnique de Paris, France · Nokia Bell Labs

Constrained DecodingConstraint-Guided Reasoning in Language ModelsControllable Text Generation
AskPDF
May 31, 2024

Grammar-Aligned Decoding

Large Language Models (LLMs) struggle with reliably generating highly structured outputs, such as program code, mathematical formulas, or well-formed markup. Constrained decoding approaches mitigate this problem by greedily restricting what tokens an LLM can output at each step to guarantee that the output matches a given constraint. Specifically, in grammar-constrained decoding (GCD), the LLM's output must follow a given grammar. In this paper, we demonstrate that GCD techniques (and in general constrained decoding techniques) can distort the LLM's distribution, leading to outputs that are grammatical but appear with likelihoods that are not proportional to the ones given by the LLM, and so ultimately are low-quality. We call the problem of aligning sampling with a grammar constraint, grammar-aligned decoding (GAD), and propose adaptive sampling with approximate expected futures (ASAp), a decoding algorithm that guarantees the output to be grammatical while provably producing outputs that match the conditional probability of the LLM's distribution conditioned on the given grammar constraint. Our algorithm uses prior sample outputs to soundly overapproximate the future grammaticality of different output prefixes. Our evaluation on code generation and structured NLP tasks shows how ASAp often produces outputs with higher likelihood (according to the LLM's distribution) than existing GCD techniques, while still enforcing the desired grammatical constraints.

Kanghee Park, Jiayu Wang, Taylor Berg-Kirkpatrick +2

University of Wisconsin-Madison · University of California San Diego

Constrained Generative ModelingConstrained DecodingAdaptive SamplingLLM Code GenerationGrammar-Constrained Decoding
AskPDF
Jul 8, 2026

Constrained Decoding for Diffusion Language Models via Efficient Inference over Finite Automata

Constrained decoding is essential for serving LLMs, ensuring that generated outputs follow specific structures such as JSON schema-formatted function calls. Existing systems are designed for autoregressive models and assume left-to-right generation, masking out invalid next tokens at each step. Diffusion language models, however, break this assumption: they sample multiple positions simultaneously from a fully-factorized mean-field distribution at each denoising step. In this paper, we present an exact and tractable algorithm for sampling from the constrained mean-field posterior under any constraint expressible as a finite automaton. Viewing finite automata as graphical models, we obtain tractable representations of the constrained distribution that enable efficient inference. The approach guarantees constraint satisfaction by construction, supports both greedy and sampling-based decoding, and is compatible with parallel and block-wise decoding under arbitrary remasking schedules. Applying depth-reduction techniques from arithmetic circuit theory, we further reduce sampling depth from linear to logarithmic in the sequence length. Empirical evaluations on Dream-7B and LLaDA-8B show substantial accuracy gains across various tasks including function calling (xLAM, BFCL), planning (Sudoku, Countdown), text-to-SQL (Spider), and math reasoning (GSM-Symbolic), with little inference overhead relative to unconstrained decoding. For example, on BFCL-Live, our approach improves Dream-7B's greedy decoding accuracy from 63.9% to 71.5%, and stochastic sampling accuracy from 22.3% to 69.0%, where the unconstrained baseline collapses, with under 5% wall-clock overhead.

Meihua Dang, Stefano Ermon

Stanford University

Constrained Generative ModelingConstrained DecodingDiffusion Language Model InferenceDiffusion Language Model Decoding
AskPDF