Think Outside the Code: Brainstorming Boosts Large Language Models in Code Generation

TL;DR

The paper targets competition-level code generation by enabling large language models to reason at a higher level before coding. It introduces Brainstorm, a three-stage framework that first generates diverse algorithmic thoughts, then selects the most promising one with a neural ranker, and finally implements code conditioned on the chosen thought in a zero-shot manner. Empirical results on CodeContests and APPS show substantial improvements over baselines, with pass@k metrics increasing by substantial margins and performance on LeetCode approaching human levels. The work demonstrates that explicit intermediate reasoning via diverse thoughts can unlock stronger algorithmic capabilities in LLMs, offering a scalable alternative to heavy RL sampling and extensive fine-tuning for competition-style problems.

Abstract

Code generation aims to automatically generate source code from high-level task specifications, which can significantly increase productivity of software engineering. Recently, approaches based on large language models (LLMs) have shown remarkable code generation abilities on simple tasks. However, generate code for more complex tasks, such as competition-level problems, remains challenging. In this paper, we introduce Brainstorm framework for code generation. It leverages a brainstorming step that generates and selects diverse thoughts on the problem to facilitate algorithmic reasoning, where the thoughts are possible blueprint of solving the problem. We demonstrate that Brainstorm significantly enhances the ability of LLMs to solve competition-level programming problems, resulting in a more than 50% increase in the pass@$k$ metrics for ChatGPT on the CodeContests benchmark, achieving state-of-the-art performance. Furthermore, our experiments conducted on LeetCode contests show that our framework boosts the ability of ChatGPT to a level comparable to that of human programmers.

Figures (5)

• Figure 1: An illustrative example for the Brainstorm framework.
• Figure 2: Illustrative example for brainstorming step. Diverse thoughts are generated, the correct ones will be selected and be used for the following code generation.
• Figure 3: The overview for training thoughts ranker model.
• Figure 4: Zero-shot prompt for implementing code.
• Figure 5: Performance of Brainstorm on different tags and ratings.