English

Reinforcement Learning-Guided Chain-of-Draft for Token-Efficient Code Generation

Software Engineering 2025-10-01 v1 Artificial Intelligence

Abstract

LLMs demonstrate surface-level fluency in code generation but struggle with structured reasoning tasks requiring correctness and semantic alignment. While Chain-of-Thought (CoT) prompting enhances reasoning through intermediate steps, it suffers from verbosity and inefficiency. Chain-of-Draft (CoD) prompting offers more concise reasoning, but the stochastic nature of LLMs produces varying solution quality, making optimal selection challenging. We propose \multicod, a reinforcement learning framework that learns to select the most promising candidate from CoD-generated solutions. Our approach uses strategy-guided prompting to encourage diverse reasoning styles and models solution selection as a contextual bandit problem. The framework optimizes interpretable features including code complexity, reasoning structure, and strategic metadata through a reward function balancing correctness, efficiency, and clarity. Experiments on MBPP, BigCodeBench, SWE-bench Verified, and Defects4J show \multicod~outperforms and in some cases, on par with standard prompting, CoT, and CoD baselines while achieving cost and token efficiency from the user's perspective through a multi-candidate design that charges only for the selected output, reducing user billing by over 50\% and improving LLM response quality, making \multicod~more sustainable and scalable for real-world deployment. Our code is available: https://anonymous.4open.science/r/MultiCoD.

Keywords

Cite

@article{arxiv.2509.25243,
  title  = {Reinforcement Learning-Guided Chain-of-Draft for Token-Efficient Code Generation},
  author = {Xunzhu Tang and Iyiola Emmanuel Olatunji and Tiezhu Sun and Jacques Klein and Tegawende F. Bissyande},
  journal= {arXiv preprint arXiv:2509.25243},
  year   = {2025}
}
R2 v1 2026-07-01T06:05:37.537Z