spot_img
HomeResearch & DevelopmentEnhancing Large Language Models' Reasoning Through Step-by-Step Code Execution

Enhancing Large Language Models’ Reasoning Through Step-by-Step Code Execution

TLDR: Researchers introduce TracePile, a 2.6 million sample dataset that converts code execution into explicit, step-by-step “Chain of Execution” (CoE) rationales. Training large language models (LLMs) with TracePile significantly improves their general reasoning abilities across mathematics, code, logic, and algorithms, outperforming traditional training methods by providing fine-grained, procedural supervision.

Building large language models (LLMs) with robust and general reasoning capabilities is a major goal in artificial intelligence. While LLMs have excelled in understanding and generating language, their ability to perform complex reasoning, such as mathematical problem-solving and code generation, remains a critical area for improvement.

Recent research has increasingly turned to code as a rich source for training, given its inherent logical structure and diverse problem-solving patterns like divide-and-conquer or enumeration. However, the reasoning within raw code is often implicit, hidden by syntax or implementation details, making it difficult for LLMs to learn effectively from it directly.

To address this challenge, a new approach introduces TracePile, a massive dataset comprising 2.6 million samples. TracePile transforms code execution into explicit, step-by-step rationales, which the researchers call Chain of Execution (CoE). This CoE format is designed to be similar to a chain-of-thought, but specifically for code execution. The dataset covers a wide range of domains, including mathematics, classical algorithms, and competitive programming problems. To further enhance logical granularity and code diversity, TracePile is enriched with variable-tracing questions and code rewritings.

The researchers evaluated TracePile using three different training setups: continue-pretraining, instruction tuning after pretraining, and a two-stage finetuning process. These experiments were conducted across four different base models, including LLaMA 3, LLaMA 3.1, Qwen-2.5, and Qwen-2.5 Coder. The models were tested on 20 benchmarks covering mathematical reasoning, coding, logical reasoning, and algorithmic problem-solving.

The results showed consistent and significant improvements across the board. For instance, TracePile boosted LLaMA3.1-8B by an average of 7.1% across nine mathematical datasets after continue-pretraining. Clear gains were also observed on benchmarks like LiveCodeBench, CRUX, and MMLU under the two-stage finetuning setup. These improvements highlight TracePile’s effectiveness in enhancing both domain-specific and general reasoning capabilities.

How TracePile is Built

TracePile gathers data from three main sources: algorithmic competition code (like Codeforces submissions), classical algorithm implementations (inspired by works like CLRS), and mathematical code (from datasets like OpenMath). To ensure quality, competition code is cleaned and validated, while mathematical problems are filtered for difficulty.

To enrich the dataset’s diversity, two key augmentation strategies are employed. First, query diversification involves generating alternative questions for code snippets, moving beyond simple input/output predictions to target internal execution states and variable changes. Second, code diversification involves rewriting original code implementations to introduce syntactic and structural variations while preserving semantics. This helps models generalize beyond rigid code templates.

The core of TracePile is the Chain of Execution generation. Frontier LLMs are prompted to create detailed, step-by-step natural language narrations that follow the execution trace of the code. This process is guided by few-shot prompting and intermediate result verification to ensure logical consistency and detail. The final CoE outputs are converted back to natural language after verification to avoid issues with structured JSON outputs during continued pretraining.

Also Read:

Key Findings and Impact

Ablation studies revealed that all components of TracePile—mathematical, algorithmic, and competition code—contribute significantly to its effectiveness. Removing any single source led to performance degradation across various reasoning domains, indicating strong cross-domain transfer. Diversification strategies, particularly code rewrites and query augmentations, were also crucial for enhancing generalization.

The research also demonstrated that performance scales with the amount of TracePile data, showing a clear upward trend across all reasoning domains as the dataset size increased. This confirms that more CoE supervision directly leads to stronger reasoning capabilities.

Qualitative analysis on challenging reasoning tasks from the BBH benchmark showed that TracePile significantly improves fine-grained reasoning abilities. Models trained with TracePile outperformed base models in tasks requiring precise state tracking, multi-step arithmetic, logical deduction, and symbolic reasoning, suggesting enhanced robustness and a reduction in typical reasoning errors.

In conclusion, TracePile offers a novel and effective way to supervise LLMs with explicit, step-by-step code execution traces. This method goes beyond traditional final-answer supervision, providing a structured reasoning signal that significantly boosts general reasoning capabilities across diverse domains. For more details, you can read the full paper here.

Nikhil Patel
Nikhil Patelhttps://blogs.edgentiq.com
Nikhil Patel is a tech analyst and AI news reporter who brings a practitioner's perspective to every article. With prior experience working at an AI startup, he decodes the business mechanics behind product innovations, funding trends, and partnerships in the GenAI space. Nikhil's insights are sharp, forward-looking, and trusted by insiders and newcomers alike. You can reach him out at: [email protected]

- Advertisement -

spot_img

Gen AI News and Updates

spot_img

- Advertisement -