Let's Revise Step-by-Step: A Unified Local Search Framework for Code Generation with LLMs

TL;DR

ReLoc is a unified local search framework for code generation with LLMs that improves efficiency and scalability by step-by-step code revision, outperforming existing methods across various tasks.

Contribution

The paper introduces ReLoc, a novel local search framework for code generation that combines multiple decision rules and a specialized reward model to enhance performance.

Findings

ReLoc outperforms tree-search and improvement-based methods in diverse code tasks.

ReLoc achieves higher code quality with fewer tokens and faster convergence.

The framework is flexible, supporting algorithms like Hill Climbing and Genetic Algorithm.

Abstract

Large Language Models (LLMs) with inference-time scaling techniques show promise for code generation, yet face notable efficiency and scalability challenges. Construction-based tree-search methods suffer from rapid growth in tree size, high token consumption, and lack of anytime property. In contrast, improvement-based methods offer better performance but often struggle with uninformative reward signals and inefficient search strategies. In this work, we propose \textbf{ReLoc}, a unified local search framework which effectively performs step-by-step code revision. Specifically, ReLoc explores a series of local revisions through four key algorithmic components: initial code drafting, neighborhood code generation, candidate evaluation, and incumbent code updating, each of which can be instantiated with specific decision rules to realize different local search algorithms such as Hill Climbing (HC) or Genetic Algorithm (GA). Furthermore, we develop a specialized revision reward model that evaluates code quality based on revision distance to produce fine-grained preferences that guide the local search toward more promising candidates. Finally, our extensive experimental results demonstrate that our approach achieves superior performance across diverse code generation tasks, significantly outperforming both construction-based tree search as well as the state-of-the-art improvement-based code generation methods.