Constraint-Aware Training

Jinwoo Kim

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

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

Panel of Figure 1 — Figure 1: Standard (CE) and constraint-aware (CA) loss compared on three different scenarios on each row. In each row, mask $M$ denotes the set of tokens allowed (marked with $\mathsf{T}$ ) according to the analysis, $t$ the one-hot target, $p$ the model prediction, and $p_{M}$ the model prediction renormalized on tokens allowed by $M$ . (a) With only the gold x allowed, CA loss is zero. (b) When the mask $M$ 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 $\texttt{x}:\texttt{y}$ ratio changes CE but leaves CA unchanged.

Panel of Figure 2 — Figure 2: Mean losses $\ell_{\mathrm{CA}}$ , $\ell_{\mathrm{CE}}^{{S_{A}}}$ , and $\ell_{\mathrm{CE}}$ , each reported in excess of the target distribution’s entropy. All values are nats per answer.

Panel of Figure 3 — Figure 3: Left: excess losses obtained by training on the depth corpus, associated with $\ell_{\mathrm{CE}}$ , $\ell_{\mathrm{CE}}^{{S_{A}}}$ (from CE training), and $\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.

Panel of Figure 4 — Figure 4: Excess loss curves against corpus size (left), and loss of the mean prediction across all 16 models ( $\ell_{\mathrm{mean}}$ ) and the average variance component of individual losses ( $\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, $\ell_{\mathrm{var}}$ is much higher then $\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

Depth	Seed	$\ell_{\mathrm{CA}}$ ( $10^{-5}$ )	$\ell_{\mathrm{CE}}^{{S_{A}}}$ ( $10^{-3}$ )	$\ell_{\mathrm{CE}}$	CE range
2	0	3.008	1.639	1.1036	4
2	1	2.382	1.536	1.1006	3
4	0	3.283	1.972	0.8318	5
4	1	2.451	2.236	0.8552	5
8	0	3.178	1.706	0.4764	8
8	1	2.450	1.494	0.5861	7

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.

Panel of Figure 5 — 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)	$\ell_{\mathrm{CA}}$	$\ell_{\mathrm{CE}}^{{S_{A}}}$	CE $\ell_{\mathrm{A}}$	$\ell_{\mathrm{CE}}$
2.4	$0.000379\pm 0.000080$	$0.028276\pm 0.002998$	0.017490	0.045766
7.2	$0.000080\pm 0.000005$	$0.011467\pm 0.005526$	0.017969	0.029436
24	$0.000075\pm 0.000004$	$0.007973\pm 0.003928$	0.015902	0.023875
36	$0.000074\pm 0.000003$	$0.007268\pm 0.002713$	0.015219	0.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)	Objective	Peak LR	Readout decay	Masked excess	CE total
2.4	CA	0.010	0	$0.000155\pm 0.000013$	–
2.4	CE	0.003	0.1	$0.029823\pm 0.005366$	0.049277
36	CA	0.001	0	$0.000060\pm 0.000004$	–
36	CE	0.003	0.1	$0.009372\pm 0.002489$	0.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.

Cell	Objective	Masked excess ( $10^{-3}$ nats)	$\ell_{\mathrm{var}}$ ( $10^{-3}$ nats)
A	CA	$0.074\pm 0.003$	0.0664
A	CE	$7.268\pm 2.713$	6.5800
B	CA	$0.047\pm 0.002$	0.0414
B	CE	$4.899\pm 5.365$	4.4032
C	CA	$0.046\pm 0.002$	0.0410
C	CE	$5.947\pm 6.708$	5.3315

Appendix

Table 4: Larger-pool and longer-schedule study. Within-mask excess is mean

\pm

sample standard deviation over 16 corpora.

\ell_{\mathrm{var}}

is the average loss removed by averaging their masked predictions. Both numerical columns use units of

10^{-3}

nats.

Constrained Generative Modeling

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

See all 187 papers

Constrained Decoding: Harmful or Beneficial? Revisiting Its Impact on LLM Performance
Oct 26, 2026Ngoc Trinh Hung Nguyen, Alonso Silva, Armen Aghasaryan +1Constrained Decoding Constraint-Guided Reasoning in Language Models
Télécom Paris, Institut Polytechnique de Paris, France · Nokia Bell Labs
Grammar-Aligned Decoding
May 31, 2024Kanghee Park, Jiayu Wang, Taylor Berg-Kirkpatrick +2Constrained Generative Modeling Constrained Decoding
University of Wisconsin-Madison · University of California San Diego
Constrained Decoding for Diffusion Language Models via Efficient Inference over Finite Automata
Jul 8, 2026Meihua Dang, Stefano ErmonConstrained Generative Modeling Constrained Decoding
Stanford University

Oct 26, 2026J/K move · Enter open · S save

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

Ngoc Trinh Hung Nguyen, Alonso Silva, Armen Aghasaryan, Mehwish Alam

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

Constrained Decoding Constraint-Guided Reasoning in Language Models Controllable Text Generation

Open paper AskPDF arXiv

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.

Abstract

Constrained Generative Modeling

Explore similar work

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

Grammar-Aligned Decoding

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