PaperlayerView exact revision on arXiv

Source-linked AI summary

Constraint-Aware Training

Jinwoo Kim

arXiv:2610.02909v1cs.LG

TL;DR

Ordinary training already teaches language models to suppress invalid tokens that constrained decoding will filter, raising whether this analysis can instead be externalized. The paper defines constraint-aware training and formalizes its benefits; its controlled experiments report lower loss than ordinary cross-entropy training. The analysis must be efficient enough not to impose significant inference overhead.

  • Problem

    Ordinary training already suppresses tokens that constrained decoding will filter, motivating whether the analysis can be externalized to avoid teaching the model that same distinction.

    Source: §1 · p. 1

  • Method

    The paper defines a general constraint-aware objective that applies a prefix analysis, renormalizes over allowed tokens, and trains against the target distribution on that support.

    Source: §1 · p. 2, §2.1 · p. 3

  • Results

    Constraint-aware models achieve consistently lower loss than ordinary cross-entropy models, including under constrained decoding, and the advantage persists where theory no longer guarantees it.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results motivate training objectives that incorporate the analyses used during generation.

    Source: §1 · p. 2

    The analysis must be efficient enough not to impose significant overhead during inference, because constraint-aware training assumes it is always deployed.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Prefix-grounded analyses can filter invalid next tokens during constrained decoding, although ordinary training already teaches the model to suppress those same tokens. The paper introduces constraint-aware training to externalize that analysis and reports lower prediction loss than ordinary cross-entropy training in controlled synthetic experiments.

1 Introduction

  • Ordinary next-token training already suppresses invalid tokens that constrained decoding will filter using prefix-grounded analyses.For example, scope analysis can allow an in-scope variable while excluding undefined names.Source: §1 · p. 1
  • Constraint-aware training applies the prefix analysis, renormalizes probability over allowed tokens, and computes cross-entropy against the target.Only tokens admitted by the analysis receive a training signal.Source: §1 · p. 2, Figure 1 · p. 2
  • The paper formalizes externalization and proves two transformer capacity results and a data-efficiency result, then tests whether the benefits carry over to training dynamics.The preliminary test uses controlled small-scale training runs over a synthetic language.Source: §1 · p. 2
  • Constraint-aware models achieve consistently lower loss than ordinary cross-entropy models, including when both are evaluated under constrained decoding.The advantage persists across training even when the theory no longer guarantees a benefit, such as when the model is wide enough.Source: §1 · p. 2

2 Externalizing Analyses Using Constraint-Aware Loss

The section decomposes standard cross-entropy into a penalty for probability assigned to disallowed tokens and a loss for the distribution over allowed tokens. Constraint-aware training drops the former component and fits the latter after masking and renormalization.

2.1 Constraint-Aware Loss Decomposition

  • A prefix analysis defines allowed and disallowed tokens at each prefix, and the paper assumes every training target is allowed by the analysis.The analysis is a computable function of a prefix and candidate token.Source: §2.1 · p. 3
  • Standard cross-entropy decomposes into support loss for probability on disallowed tokens and constraint-aware loss for misprediction within the allowed set.In the gradient decomposition, the support component pushes probability off disallowed tokens while the constraint-aware component redistributes mass among allowed tokens.Source: §2.1 · p. 3

2.2 Constraint-Aware Training

  • Constraint-aware training drops the support gradient and fits the corpus frequencies after masking and renormalization.This trains the model on the relative distribution among allowed tokens rather than on probability assigned to tokens the analysis will filter.Source: §2.2 · p. 4
  • Constraint-aware loss minimizes the weighted average KL divergence between the corpus next-token distribution and the model distribution renormalized over allowed tokens.The fixed corpus-entropy term means minimizing average constraint-aware loss also minimizes this divergence.Source: §2.2 · p. 4
  • Different full-vocabulary predictions can satisfy the same constraint-aware objective when their renormalized distributions on allowed tokens match the corpus.For a corpus with four x targets and one y target, both (0.8, 0.2, 0.0) and (0.4, 0.1, 0.5) satisfy the objective when only x and y are allowed.Source: §2.2 · p. 4

3 The Benefits of Externalization: Width

Externalizing scope analysis reduces the representational burden of suppressing out-of-scope names. In a synthetic width sweep, constraint-aware (CA) training achieves substantially lower excess loss than standard cross-entropy (CE) training.

3 The Benefits of Externalization: Width

  • In the example with any pair of five names allowed, a one-dimensional readout can model the target pair’s relative probabilities under CA, while unfiltered CE requires width at least 4.The synthetic experiment uses 12 in-scope names drawn from 61, with 12 functions specifying their target distributions.Source: §3 · p. 5
  • Theorem 3 requires hidden-state width at least m − 1 for CA and 2m for CE to support arbitrary m-token allowed sets with arbitrary precision.The result concerns a capable readout; finding a suitable readout and trunk through training is a separate practical question.Source: §3 · p. 5
  • At w = 10, CA excess loss was 0.01; it fell to 0.0004 at w = 11 and stayed near 0.0004–0.0005 beyond that.The sharp drop occurs at the width threshold predicted by Theorem 3.Source: §3 · p. 5, Figure 2 · p. 6
  • For CE, excess loss stayed above 0.8 below w = 24 and remained above 0.05 at all tested widths; at w = 32, ℓCE was 0.1812.Although Theorem 3 gives w = 24 for an optimal capable readout, the trained CE model did not learn it; w = 32 was the smallest width displaying capability across sampled queries.Source: §3 · p. 6, Figure 2 · p. 6
  • Across the reported width sweep, ℓCA was much lower than ℓCE, often by orders of magnitude; externalizing scope during training yielded a closer fit to the variable-usage distribution.Source: §3 · p. 6

4 The Benefits of Externalization: Depth

When evaluating formulas determines which names are allowed, CE training must learn that analysis as well as the target distribution. In the synthetic depth experiment, CA loss stayed nearly constant as depth varied, while CE loss improved with depth but remained higher.

4 The Benefits of Externalization: Depth

  • Theorem 4 requires depth L ≥ dA(ϵ/log 2) for expected KL divergence at most ϵ, where dA measures the depth needed to identify allowed sets with bounded average error.The bound connects the depth needed for the analysis to the depth needed for the model’s prediction loss.Source: §4 · p. 6
  • For a fixed error bound, dA can grow with input length for common analyses, so CE-trained transformers must grow in depth to maintain expected loss at most ϵ on longer inputs.The paper notes that formal-language recognition can require transformer depth to grow with input length.Source: §4 · p. 6
  • On about 550,000 held-out queries, CA loss stayed around 2.7 × 10^-5 nats as depth increased, while CE loss fell from 1.1 nats at depth 2 to 0.53 at depth 8.The CE model’s additional layers were used to evaluate F, the analysis; CA outperformed both total CE loss and CE loss measured only on allowed tokens.Source: Figure 3 · p. 7, §4 · p. 7
  • The synthetic corpus pairs arithmetic formulas with target tokens; its analysis evaluates each formula and returns the four allowed names.CA learns the tag-dependent distribution over allowed names, while CE must also learn to evaluate the formula.Source: §4 · pp. 6–7
  • As formula operator count grows, CE loss rises and eventually plateaus around 1.1 nats when the model predicts a uniform distribution instead of identifying the allowed set.The experiment indicates CE learns an input-size-dependent circuit for evaluating F, whereas CA avoids that burden and maintains much lower constant excess loss.Source: §4 · p. 7

5 The Benefits of Externalization: Data Efficiency

The section argues that corpus-to-corpus variation contributes substantially to expected excess loss and tests whether constraint-aware training reduces that variation and the data needed for a given loss. The synthetic experiment reports lower loss and less prediction fluctuation for CA, with its variance component flattening at 7.2M training pairs.

5 The Benefits of Externalization: Data Efficiency

  • The depth bounds may not capture the depth needed to model allowed-token distributions; how these requirements interact is left for future work.This qualification concerns the interaction of inner-distribution depth with Theorem 4’s bounds.Source: §5 · p. 7
  • CA achieves significantly lower excess loss across corpus sizes and less prediction fluctuation, with its variance component flattening by 7.2M training pairs.CE variance continues decreasing from 24M to 36M pairs; the authors conjecture that CA’s ≈7 × 10^-5 floor may reflect training noise or kernel nondeterminism.Source: Figure 4 · p. 8, §5 · pp. 8–9
  • Under minimum-change fitting and conditional-centering assumptions, Theorem 5 attributes higher expected CE loss to greater parameter change and resulting variance.The theorem compares expected KL divergences over corpora and held-out prefixes; it does not imply a loss floor.Source: §5 · p. 8
  • The controlled experiment varies corpus size across 2.4M, 7.2M, 24M, and 36M pairs and trains 16 independent corpora per size with other hyperparameters fixed.Smaller pools are revisited to allocate the same training compute and isolate corpus size and variation.Source: §5 · p. 8
  • The variance component ℓvar is roughly an order of magnitude higher than ℓmean, indicating corpus variation is a significant contributor to expected excess loss.The experiment defines ℓmean as the loss of the mean prediction across 16 models and ℓvar as the average individual loss minus ℓmean.Source: §5 · pp. 8–9

6 Discussion and Related Work

The discussion extends the loss decomposition to nested analyses, allowing separate weights for their contributions, while noting that only the zero-weight case is studied. The appendix derives loss and gradient identities that clarify how constraint-aware training reallocates learning between allowed and rejected tokens.

6 Discussion and Related Work

  • For nested parser, scope, and type analyses, the loss decomposes into the final allowed-set loss plus separately weighted losses for each analysis.Each component measures probability mass removed by an analysis after the preceding analyses; the paper studies the case where all weights are zero.Source: §6 · p. 9
  • The coefficients on analysis-specific losses remain in their parameter gradients, so they can amplify or dampen each analysis signal; other weight settings remain future work.The paper studies only λ_i = 0.Source: §6 · p. 9
  • Constraint-aware training assumes the analysis is always used at inference, so that analysis must be efficient enough not to impose significant runtime overhead.The paper situates this work alongside prior constrained-decoding methods [3, 7–9].Source: §6 · p. 9
  • Prior masked-loss variants include GraFT [1] and Liao et al. [5], which focus on fine-tuning, whereas this paper studies analyses during pretraining.The paper presents its distinction as generalizing the decomposition to any analysis and formalizing its training benefits.Source: §6 · p. 9

A.1 Loss and gradient decomposition

  • The CE excess loss separates masked KL error from support loss, while their logit gradients respectively rebalance allowed-token mass and push probability off disallowed tokens.The CA gradient is the within-support term; combining it with the support gradient recovers the ordinary CE gradient.Source: §A.1 · p. 11

A.2 Proof of Theorem 2

The proof shows that minimizing CA loss fits empirical conditional token frequencies after masking and renormalization, while leaving disallowed-token mass unconstrained. Under conditions that make masked sampling a normalized string distribution, the same objective minimizes divergence from the empirical string distribution.

A.2 Proof of Theorem 2

  • CA loss minimizers match the empirical conditional distribution on allowed tokens at every observed prefix, while raw predictions may allocate remaining mass arbitrarily among disallowed tokens.Exact agreement can require zero masked probability for an allowed token with zero empirical frequency, which finite softmax logits only approach as a limit.Source: §A.2 · pp. 11–12

A.3 Likelihood of complete strings

  • When the analysis retains every target, admits a token at each reachable nonterminal prefix, and forces EOS within a fixed number of steps, minimizing average CA loss also minimizes KL divergence from the empirical string distribution.The result assumes every position, including termination, is scored; scoring only designated answer positions does not define a likelihood for the surrounding program.Source: §A.3 · p. 12

B Proofs for Section 3: Width

  • The width proof distinguishes fitting ratios among allowed tokens from making those tokens collectively more likely than rejected tokens.The theorem’s constructions and lower bounds concern readouts; a particular trunk must still compute the required hidden states.Source: B Proofs for Section 3: Width · p. 12

B.1 Capability and proof of Theorem 3

Theorem 3 characterizes when a fixed readout can fit allowed-token distributions and suppress rejected tokens, giving different width bounds for CA and CE. The proofs and qualifications address fixed-readout loss bounds, subword tokenization, and how the width argument relates to depth.

B.1 Capability and proof of Theorem 3

  • CA capability requires affine independence of the allowed-token readout rows; CE additionally requires a direction that ties allowed logits and places every rejected logit below them.The fixed-readout equivalence characterizes whether arbitrary positive target distributions can be approximated after masking, or before masking for CE.Source: §B.1 · pp. 13–14
  • For allowed sets of size at most m, a CA-capable readout exists at width m − 1, while CE capability requires width 2m when every m-subset of a pool of at least 2m + 1 tokens is allowed.The CE lower bound depends on that additional subset condition; a particular trunk must also realize the needed hidden states.Source: §B.1 · p. 13

B.2 Lower bounds for the loss of a learned readout

  • The fixed-readout lower-bound procedure verifies positive loss bounds using log-sum-exp inequalities and rational interval arithmetic; a verified bound cannot be invalidated by optimization.The bound is specific to the fixed readout, and changing that readout changes the bound.Source: §B.2 · pp. 14–15

B.3 When a variable name spans several vocabulary tokens

  • For multi-token names, the analysis must induce allowed sets at intermediate token prefixes; lexeme-level set sizes and subset conditions do not automatically transfer to vocabulary tokens.Implementations must also handle unfinished lexemes and tokens crossing lexeme boundaries.Source: §B.3 · p. 15

C Proofs for Section 4: Depth

  • The depth theorem relates the depth needed for small CE excess to the depth needed to identify the allowed set, while CA can leave that decision to the analysis.A separate result derives an expected CE excess lower bound from per-input-size decision-accuracy bounds.Source: C Proofs for Section 4: Depth · p. 15

C.1 Analysis accuracy and proof of Theorem 4

The depth analysis relates CE excess loss to a model’s ability to identify the analysis-selected candidate set, while the later proof passages formalize how leaving that decision to the analysis can reduce fitting requirements. These results rely on specified candidate-set and fitting assumptions.

C.1 Analysis accuracy and proof of Theorem 4

  • Expected CE excess at most ϵ implies average decision error at most ϵ/log 2, so model depth must be at least dA(ϵ/log 2).The error is averaged over corpus prefix sizes; it need not meet the bound separately at every size.Source: §C.1 · p. 16

C.2 Leaving the decision to the analysis

  • Cue-only predictions can attain zero CA excess without computing which candidate set the analysis selects, provided the cue-dependent probabilities are computable at that depth.For the uniform mixture over r candidate-set target laws, the support loss is log r.Source: §C.2 · p. 16

C.3 A loss bound from decision hardness at larger input sizes

  • If every depth-L model has decision error at least 1 − 1/r − δ above size R(L), its CE excess is at least (1 − 1/r − δ) log 2 times the corpus probability of those sizes.The lower bound means any CE model meeting tolerance ϵ must make this bound at most ϵ, while the stipulated cue-only model can meet that tolerance under CA.Source: §C.3 · pp. 16–17
  • The decision-hardness condition must hold for every model in the specified class on the actual encoded prefix distribution; trained-model errors do not establish it.The reduction from decision error to loss also depends on disjoint candidate sets and does not directly apply to arbitrary overlapping scope sets.Source: §C.1 · p. 16, §C.3 · p. 17

D.1 The two fitting problems

  • The data-efficiency proof fits full conditional target laws at sampled prefixes; its idealized interpolation assumption requires independent stacked constraints and excludes guarantees for ordinary sampled-token AdamW training.The stated sample-size comparison applies only when n(|V| − 1) ≤ P and the constraints have full row rank.Source: §D.1 · p. 17

D.2 The extra parameter change needed to fit support

  • The minimum-norm CE fit adds a parameter change orthogonal to the CA fit to satisfy rejected-token support constraints, while preserving the fitted training conditionals.At held-out prefixes, that additional change affects relative logits through the allowed-token Jacobian.Source: §D.2 · p. 18

D.3 Full statement and proof of Theorem 5

Theorem 5 compares CA and CE fits in an idealized setting, where CE’s extra support-fitting update preserves training conditionals. Under conditional centering, CA has no greater expected held-out within-mask excess, and this yields a sample-size comparison within the theorem’s valid range.

D.3 Full statement and proof of Theorem 5

  • Under conditional centering, expected held-out within-mask excess is no greater for CA than CE, and is strictly lower when the extra CE update changes held-out predictions with positive probability.The centering assumption is needed to remove the linear term; preserving training conditionals alone does not prevent systematic change at unseen prefixes.Source: §D.3 · pp. 18–19

D.4 How many examples are needed to reach a target loss?

  • At any common valid sample size, CA requires no more examples than CE to reach a fixed expected within-mask excess-loss target.The comparison concerns the specified learner families and valid sample-size range, not a universal sample multiplier or AdamW rate.Source: §D.4 · p. 19

D.5 The corpus-fluctuation loss in Figure 4

  • The Figure 4 corpus-fluctuation loss is the average loss removed by averaging predictions, not the variance of scalar losses.Average individual excess decomposes into the loss of mean predictions plus this component, with both measured in nats.Source: §D.5 · p. 19

E Experimental Setup and Supporting Results

  • The experiments use known target distributions and vary model size, depth, or available training corpus across three studies.The shared decoder uses four attention heads, vocabulary size 134, context length 4,096, and AdamW optimization; parameters and samples are matched within each objective pair.Source: E Experimental Setup and Supporting… · pp. 19–20

E.1 Width: setup and controls for Section 3

The width study tests readout capacity on synthetic scope-conditioned targets, with matched training samples and evaluation queries. Supporting checks compare learned readout fit, a fixed-scope control, and a longer training schedule.

E.1 Width: setup and controls for Section 3

  • The width sweep uses four layers, internal width 64, and readout widths from 8 to 64, evaluating excess loss on 480,000 held-out queries.The target distribution is a softmax over identifiers restricted to the twelve names in scope; the table caption describes six matched two-stage pairs on a uniform-size panel.Source: §E.1 · p. 20, Table 1 · p. 21
  • At readout width 11, the CA readout fits all 3,072 checked targets exactly, while CE has mean excess-loss lower bound 1.0172 and actual excess 1.0786 on that panel.At width 24 the learned CE readout still has a positive mean lower bound of approximately 0.153; width 32 is the first tested CE width fitting every checked target to arbitrary precision.Source: §E.1 · p. 20
  • The fixed-scope control lets both objectives fit accurately at width 12: CE total excess is 0.000229 and CA masked excess is 0.000200.This control changes both the input distribution and the task of identifying the current scope.Source: §E.1 · p. 20
  • With the longer schedule, CE masked excess falls from 0.027137 to 0.015593 but remains above CA’s 0.000472; CE total excess falls from 0.0537 to 0.0290.The comparison changes exposure and learning-rate trajectory together, and each objective has one run under each schedule.Source: §E.1 · p. 20

E.2 Depth: setup and controls for Section 4

In the depth comparison, constraint-aware (CA) training has lower masked excess than cross-entropy (CE) across matched model pairs. CE’s excess is dominated by support loss, while its allowed-set identification improves with depth.

Depth comparison

  • CA learns tag-dependent probabilities at the smallest tested depth, while additional CE layers reduce support loss but leave masked prediction less accurate.Both objectives follow the same curriculum in this depth comparison.Source: §E.2 · p. 21

Setup and controls

  • Both objectives use width 128, depths 2, 4, and 8, and two stages of about 6 × 10^8 positions each, sharing corpora, initialization, and ordering seeds.The first stage uses uniform short sizes; the second uses Pr(s) ∝ 0.85s^-2 for 2 ≤ s ≤ 24. Separate hyperparameter selection for CA and CE was not tested.Source: §E.2 · p. 21

Depth comparison

  • CA’s masked excess is lower in every pair, by factors of 54–91.For depth 2, seed 0, CA is [3.004, 3.013]×10^-5 and CE is [1.631, 1.647]×10^-3; the intervals describe held-out-document variation, not variation across training runs.Source: §E.2 · p. 21
  • CE’s total excess is over 99.6% support loss in each pair, while its allowed-set identification accuracy improves with depth.CA’s decision accuracy remains near 1/5; at deployment, the analysis supplies the correct group.Source: §E.2 · p. 21

Evaluation criterion

  • The reported CE 95%-accuracy endpoints are operator counts 4, 5, and 8 at depths 2, 4, and 8 for seed 0, and 3, 5, and 7 for seed 1.Each endpoint requires all sizes through it to meet the stated nominal Bonferroni document-bootstrap lower-bound criterion.Source: §E.2 · p. 21

Evaluation size distribution

  • Under the geometric size law used for stage two, CE’s total-excess ordering by depth is unchanged; CA’s masked losses remain in the same small range as under uniform sizes.The alternative evaluation used 2,000 held-out documents and 230,537 queries; neither evaluation establishes a lower bound on the best achievable error at a given depth.Source: §E.2 · pp. 21–22

E.3 Data efficiency: setup and controls for Section 5

The data-efficiency study compares CA and CE across independently generated corpora, matched training recipes, and additional tuning and schedule controls. CA’s masked excess remains lower across the reported comparisons, while corpus and recipe sensitivity qualify the interpretation.

Primary study setup

  • The primary study uses sixteen independently generated corpora per size and matched CA/CE pairs sharing corpus, initialization, batch-order seed, schedule, and evaluation panel.Pool sizes are 2.4M, 7.2M, 24M, and 36M training pairs; the largest pool contributes 146,496 distinct documents within the fixed budget.Source: §E.3 · p. 22

Primary data study

  • Across four pool sizes, averaging the sixteen predictions removes 89–93% of CA’s excess and 91–93% of CE’s; the amount removed is 75–146 times larger for CE.CA at 2.4M pairs has lower masked excess and ℓvar than CE at 36M pairs under the fixed recipe.Source: §E.3 · p. 22

Independent selection

  • With independent hyperparameter selection, CA has lower masked excess in all eight evaluation pairs; mean CE-to-CA excess ratios are approximately 193 at 2.4M pairs and 156 at 36M.The comparisons assess sensitivity within the tested recipe grid, not optimal hyperparameters.Source: §E.3 · p. 23

Longer training and larger pools

  • The longer-training cells separate repeating a fixed pool from accessing a larger one, but the missing 72M-pair short-schedule cell prevents separating all corpus-size and schedule effects.The larger-pool cell changes the loader permutation, so it is not a continuation of the original realized training stream.Source: §E.3 · p. 23
  • CA has smaller mean masked excess and ℓvar in all three longer-training cells, but its fluctuation loss falls from 6.64 × 10^-5 to 4.14 × 10^-5 nats with a longer schedule at the same pool size.This indicates the residual in the primary sweep is sensitive to the training recipe; controls support the observed CA advantage while leaving the residual-error source open.Source: §E.3 · pp. 23–24

Corpus variation

  • CE’s loss varies substantially across corpora, while resampling bands for its changes in log fluctuation loss span zero.The bands describe sensitivity to the observed corpora, not a general training-run confidence interval.Source: §E.3 · p. 24
Loading 2610.02909v1…