Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

AdapTrack: Constrained Decoding without Distorting LLM's Output Intent

Yongmin Li, Jia Li, Ge Li, Zhi Jin

TL;DR

AdapTrack addresses the distribution distortion problem in constrained decoding for LM-based code generation by introducing adaptive backtracking. The method maintains prefix-validity estimates and uses rejection sampling to backtrack paths that undermine the model's output intent, ensuring constraint satisfaction without distorting the LM distribution $P(s|c)$. Theoretical guarantees show AdapTrack samples from the constrained distribution, and extensive experiments across synthetic and real-world API datasets, general code benchmarks, and DSL domains demonstrate substantial improvements over traditional constrained decoding. The work suggests significant practical impact for reliable, intent-aligned code generation under constraints, with ongoing work to extend API coverage and efficiency.

Abstract

Language model-based code generation and completion tools have been widely adopted, but they may sometimes produce code that does not meet necessary constraints, such as syntactic correctness or API existence. Constrained decoding techniques are developed to help the model generate code adhering to the constraints by greedily eliminating generation options that violate constraints at each step of the generation process. However, there is a severe limitation of constrained decoding, that it distorts the model's output intent, forcing it to produce code that may satisfy the constraint but does not match the development intent and is therefore incorrect. In response to this challenge, we propose AdapTrack. By incorporating backtracking into the generation process, AdapTrack avoids distorting the output intent of the model, thereby producing results that are not only constraint-compliant but also more semantically aligned with model's output intent. On our synthetic API completion dataset, AdapTrack can achieve up to 360.87% improvement compared to constrained decoding; on the real-world API completion dataset we collect that exhibits similar issues, AdapTrack can achieve up to 38.93% improvement over constrained decoding; in general code genration benchmarks, compared to constrained decoding, AdapTrack can achieve up to 7.84% improvement on HumanEval, and up to 6.42% improvement on MBPP. This indicates that, simply by better adhering to the model's output intent, AdapTrack can achieve significant improvements. We provide a theoretical proof that the distribution produced by AdapTrack aligns with the model's distribution given the generated tokens, thereby ensuring that the model's output intent is not distorted. Experiments on DSL problems show that, compared to existing methods, our approach can provide generation results that are more consistent with the language model's distribution.

AdapTrack: Constrained Decoding without Distorting LLM's Output Intent

TL;DR

AdapTrack addresses the distribution distortion problem in constrained decoding for LM-based code generation by introducing adaptive backtracking. The method maintains prefix-validity estimates and uses rejection sampling to backtrack paths that undermine the model's output intent, ensuring constraint satisfaction without distorting the LM distribution . Theoretical guarantees show AdapTrack samples from the constrained distribution, and extensive experiments across synthetic and real-world API datasets, general code benchmarks, and DSL domains demonstrate substantial improvements over traditional constrained decoding. The work suggests significant practical impact for reliable, intent-aligned code generation under constraints, with ongoing work to extend API coverage and efficiency.

Abstract

Language model-based code generation and completion tools have been widely adopted, but they may sometimes produce code that does not meet necessary constraints, such as syntactic correctness or API existence. Constrained decoding techniques are developed to help the model generate code adhering to the constraints by greedily eliminating generation options that violate constraints at each step of the generation process. However, there is a severe limitation of constrained decoding, that it distorts the model's output intent, forcing it to produce code that may satisfy the constraint but does not match the development intent and is therefore incorrect. In response to this challenge, we propose AdapTrack. By incorporating backtracking into the generation process, AdapTrack avoids distorting the output intent of the model, thereby producing results that are not only constraint-compliant but also more semantically aligned with model's output intent. On our synthetic API completion dataset, AdapTrack can achieve up to 360.87% improvement compared to constrained decoding; on the real-world API completion dataset we collect that exhibits similar issues, AdapTrack can achieve up to 38.93% improvement over constrained decoding; in general code genration benchmarks, compared to constrained decoding, AdapTrack can achieve up to 7.84% improvement on HumanEval, and up to 6.42% improvement on MBPP. This indicates that, simply by better adhering to the model's output intent, AdapTrack can achieve significant improvements. We provide a theoretical proof that the distribution produced by AdapTrack aligns with the model's distribution given the generated tokens, thereby ensuring that the model's output intent is not distorted. Experiments on DSL problems show that, compared to existing methods, our approach can provide generation results that are more consistent with the language model's distribution.
Paper Structure (52 sections, 2 theorems, 13 equations, 7 figures, 5 tables, 2 algorithms)

This paper contains 52 sections, 2 theorems, 13 equations, 7 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Motivating example
  3. Background
  4. Goal
  5. Constrained decoding and its problem
  6. AdapTrack
  7. The proof of AdapTrack
  8. Evaluation
  9. RQ1
  10. RQ2
  11. RQ3
  12. RQ4
  13. RQ1: Performance of AdapTrack in predicting correct and valid APIs
  14. Dataset
  15. Metrics
  16. ...and 37 more sections

Key Result

theorem 1

Algorithm alg:mine only returns a valid complete sequence, and the probability of generating different sequences that have been explored during decoding satisfies equation eq:fair.

Figures (7)

  • Figure 1: An example of incorrect decoding results in constrained decoding.
  • Figure 2: Motivating example: how backtracking works on the example in Figure \ref{['fig:ex1']}.
  • Figure 3: An example code in the TFv1 dataset.
  • Figure 4: The change of KL divergence ($\downarrow$) between every 500 samples and the oracle distribution on different datasets.
  • Figure 5: The change of EM@$k$ in different temperature.
  • ...and 2 more figures

Theorems & Definitions (2)

  • theorem 1
  • theorem 2