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

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. D, 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. D 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 M Large language models (LLMs) have become powerful tools for automating software engineering tasks, including code completion and program repair. Despite their fluency and generalization capabilities, LLMs often produce brittle or incorrect code, especially in tasks that require multi-step reasoning or nuanced understanding of program semantics. This limitation stems from the fact that most LLMs generate code in a single forward pass, without explicitly modeling the reasoning process that underlies correct and efficient solutions. To address this, recent work has explored Chain-of-Thought (CoT) prompting [1], which encourages LLMs to generate intermediate reasoning steps before producing a final answer. Originally developed for arithmetic and symbolic reasoning tasks, CoT has been adapted to code-related domains to improve correctness [2]-[6].