TreeCoder: Systematic Exploration and Optimisation of Decoding and Constraints for LLM Code Generation

2025-11-27

AI Summary

Large language models (LLMs) frequently generate syntactically or semantically incorrect code when relying solely on natural language prompts, and existing approaches overemphasize prompt engineering while lacking structural and correctness guarantees during decoding. This paper proposes TreeCoder, a constraint-aware tree-search-based decoding framework that—uniquely—unifies syntactic, stylistic, and execution constraints into differentiable, optimizable modules and jointly learns both the decoding policy and constraint hyperparameters. By integrating standard numerical optimization algorithms, TreeCoder dynamically enforces code structural validity and functional correctness throughout generation, eliminating reliance on manual prompt engineering. Evaluated on MBPP and SQL-Spider benchmarks, TreeCoder significantly improves the functional accuracy of open-source LLMs—including CodeLlama, Mistral, and DeepSeek—outperforming unconstrained baselines by substantial margins.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Large language models (LLMs) have shown remarkable ability to generate code, yet their outputs often violate syntactic or semantic constraints when guided only through natural language prompts. We introduce TreeCoder, the most general and flexible framework to date for exploring decoding strategies, constraints, and hyperparameters in LLMs, and use it in code generation to enforce correctness and structure during decoding rather than relying on prompt engineering. TreeCoder represents decoding as a tree search over candidate programs, where both decoding strategies and constraint functions - such as style, syntax, execution - are treated as first-class, optimisable components. This design enables systematic exploration and automatic tuning of decoding configurations using standard optimisation techniques. Experiments on the MBPP (Python) and SQL-Spider benchmarks show that TreeCoder consistently improves accuracy across open-source models such as CodeLlama, Mistral and DeepSeek, often outperforming their unconstrained baselines by considerable margins.

Problem

• Enforces syntactic and semantic correctness in LLM-generated code
• Systematically explores decoding strategies and constraint functions
• Improves code generation accuracy across multiple programming languages

Innovation

• TreeCoder uses tree search over candidate programs
• It treats decoding strategies and constraints as optimizable components
• The framework enables systematic exploration and automatic tuning