CodeIO: Condensing Reasoning Patterns via Code Input-Output Prediction

CodeIO: Condensing Reasoning Patterns via Code Input-Output Prediction

Speakers: Junlong Li, Daya Guo, Dejian Yang, Runxin Xu, Yu Wu, Junxian He

Conference: ICML 2025

YouTube: https://slideslive.com/39044114

Overview

The talk introduces CodeIO, a novel framework designed to enhance the general reasoning capabilities of large language models (LLMs) by leveraging "code in the wild." Presented by a proxy on behalf of the authors, including first author Junlong Li, CodeIO addresses a critical bias in current LLMs: their disproportionate strength in mathematical and code generation reasoning due to the abundance of high-quality, structured datasets in these domains (e.g., Art of Problem Solving, LeetCode). In contrast, reasoning in other crucial areas like logical, symbolic, scientific, and natural language understanding remains underdeveloped due to data scarcity.

CodeIO proposes a unique approach that teaches LLMs to predict both the input and output of code execution, expressed as natural language chain-of-thought. This methodology aims to distill core reasoning patterns from diverse real-world code, making them explicit and transferable across various tasks and domains. By transforming raw code into a structured, executable format and generating corresponding natural language queries and I/O pairs, CodeIO creates a scalable, high-quality dataset that fundamentally shifts LLM training from mere output construction to deeply understanding the problem-solving process.

This work is significant because it offers a pathway to unlock broader and more robust reasoning skills in LLMs, moving beyond their current domain-specific proficiencies. By tapping into the vast, yet often noisy, reservoir of real-world code, CodeIO provides a systematic method to extract universal reasoning primitives and package them in a format conducive to generalizable LLM learning. The framework's emphasis on verifiable correctness and natural language expression ensures that the learned skills are not only accurate but also interpretable and adaptable.

Background

▶ Watch: LLM reasoning bias and diverse data need (0:00)

The rapid advancements in large language models (LLMs) have showcased impressive reasoning abilities, particularly with models like GPT-4, DeepSeek R1, and Gemini 1.5. However, a closer examination reveals a significant bias: these models predominantly excel in mathematical problem-solving and code generation. This specialization is not arbitrary; it stems directly from the availability of extensive, meticulously curated datasets in these fields. For instance, platforms like Art of Problem Solving provide rich resources for math reasoning, while LeetCode and Codeforces offer vast repositories for coding challenges, complete with solutions and test cases.

Conversely, LLMs often struggle with reasoning tasks in domains where such well-structured data is sparse. This includes logical deduction, symbolic manipulation, scientific inquiry, and nuanced language-based reasoning. The lack of diverse, high-quality training data across these areas hinders the development of truly generalized reasoning capabilities, leading to models that, despite their apparent sophistication, possess a narrow cognitive scope. A truly helpful and versatile AI model must learn to reason effectively across a multitude of domains.

To address this data scarcity, researchers have explored various avenues, with "code in the wild" emerging as a promising, yet challenging, resource. Code written for real-world applications reflects diverse problem-solving needs—from string processing to complex calculations—and encapsulates universal reasoning primitives. Programming language keywords themselves embody these primitives: if and return denote conditional logic and planning, while while loops enable state-space search through iteration. The challenge lies in extracting these valuable reasoning signals from the often-messy and domain-specific context of raw code.

Past attempts to leverage code for LLM training have faced limitations:

1. Raw Code Pre-training: Directly feeding large corpora of real-world code to models. While simple, this approach suffers from noise; much of the code (e.g., UI rendering, visualization, configuration) contains irrelevant signals that can distract the model from learning core reasoning skills. The reasoning patterns are implicit and often buried within vast amounts of boilerplate.
2. Structured Code Generation Tasks: Transforming code into tasks where models generate code based on natural language descriptions. This approach distills reasoning patterns more explicitly and ensures syntactic correctness. However, it constrains the model's output to programming language syntax, potentially limiting its ability to express general reasoning in natural language and transfer these skills to non-coding tasks.

CodeIO distinguishes itself from these paradigms by proposing a method that simultaneously distills reasoning patterns and ensures their transferability. Instead of merely generating code or processing raw code, CodeIO trains models to predict both the input and output of code execution, describing this transformation using natural language chain-of-thought. This dual prediction in natural language ensures that the core logic is preserved and condensed, while the reasoning process becomes explicit, interpretable, and highly transferable across different tasks and domains, moving beyond the limitations of syntax-constrained outputs.

Key Findings

▶ Watch: CodeIO data construction pipeline explained (2:58)

The CodeIO framework demonstrates significant advancements in improving LLM reasoning, substantiated by several key findings:

• Superior Performance Across Diverse Tasks: CodeIO models consistently outperform all established baselines across 14 distinct reasoning tasks spanning various domains. This includes models trained solely with standard supervised fine-tuning (SFT), as well as those leveraging alternative pre-training strategies such as large-scale instruction tuning on STEM datasets or continued pre-training on raw code files. This indicates that CodeIO's approach effectively extracts and condenses generalizable reasoning patterns.

• Enhanced Performance with Revision (CodeIO++): The optional revision step, which creates the CodeIO++ dataset, consistently yields even better performance. By identifying and correcting incorrect initial responses through re-execution verification and regenerating improved answers with verifier feedback, CodeIO++ refines the training signal, leading to more accurate and robust reasoning.

• Crucial Role of Bidirectional Reasoning: Ablation studies revealed that training models to predict both the input and output (bidirectional reasoning) significantly outperforms models trained to predict only the input or only the output. This highlights the value of understanding the transformation process from both directions, fostering a more comprehensive grasp of the underlying logic.

• Data Scale Matters: Experiments confirm that CodeIO's performance scales positively with data size. Whether increasing the number of unique raw code files or generating more input-output pairs per code sample, larger datasets consistently lead to better and often significantly improved results. This suggests that the method can continue to benefit from expanding its data collection efforts.

• Beyond Synthesizer Quality: A controlled experiment demonstrated that CodeIO's efficacy is not merely a byproduct of using a superior data synthesizer (e.g., DeepSeek V2.5). While a better synthesizer (like DeepSeek V2.5 in WIDS 2.5) does improve performance compared to older synthesizers, CodeIO still significantly outperforms these models even when both use the same powerful synthesizer. This underscores that the unique reasoning prompts and condensed reasoning patterns inherent in the CodeIO data construction pipeline are the primary drivers of its success.

• Optimality of Single-Round Revision: While CodeIO++ significantly improves performance, attempts to implement multiple rounds of revision (CodeIO+++) showed inconsistent gains. This suggests that one round of targeted revision based on verifiable feedback is sufficient and optimal for maximizing performance without introducing diminishing or negative returns.

• Necessity of Two-Stage Training: The research found that a two-stage training strategy is essential. Models trained first on CodeIO/CodeIO++ data to develop strong reasoning capabilities, followed by continued training on broader supervised fine-tuning data for instruction alignment, consistently outperform models trained in a single stage.

• Generalizability to Reinforcement Learning: CodeIO data proves highly valuable in the context of Reinforcement Learning for Verifiable Reward (RLVR). Training on CodeIO data leads to non-trivial performance gains not only in code-related tasks but also in mathematics, science, logic, and tabular reasoning. This finding reinforces CodeIO's ability to teach generalizable reasoning skills that transfer effectively across different learning paradigms and diverse domains.

Technical Deep Dive

▶ Watch: CodeIO++ data revision for improved correctness (4:50)

The technical core of CodeIO lies in its innovative data construction pipeline and a two-stage training strategy, meticulously designed to distill and transfer reasoning patterns from code into natural language chain-of-thought.

Core Idea

CodeIO's fundamental principle is to teach LLMs to predict both the input and output of a piece of code, articulating the transformation process in a natural language chain-of-thought. This approach aims to move beyond simple code generation or raw code pre-training by forcing the model to "think through" the problem in an interpretable, step-by-step manner, much like a human programmer traces execution.

Data Construction Pipeline

The creation of the CodeIO and CodeIO++ datasets involves a five-step process:

1. Collect Raw Codes: The initial step involves gathering a large corpus of raw code files from high-quality, diverse sources such as GitHub repositories and Python education platforms. This ensures a wide variety of real-world contexts and problem types.

1. Transform Raw Code into a Unified Form: This is a crucial step for extracting core reasoning logic and preparing the code for automated execution and data generation. It involves three key components:

• Prepare Reference Code with Main Function: Raw code is often messy and contains elements irrelevant to core logic (e.g., UI rendering, visualization, external dependencies). This component refactors the code into a clean, self-contained reference code snippet, typically wrapped in a main function or similar entry point. This filters out noise and isolates the specific function or logic intended for reasoning.
• Define an Input Generator: To scale data generation efficiently and ensure diversity, an input generator function is defined for each reference code. This function automatically produces realistic and varied inputs that conform to the expected format of the reference code, enabling the creation of numerous I/O pairs.
• Generate Natural Language Queries: For each transformed code, a concise natural language query is created. This query clearly describes the task the code performs and what its input and output represent, bridging the gap between code logic and human understanding.
• Example: A messy raw code file might be transformed into a clean Python function, accompanied by a query like "Given a list of numbers, return their sum," an input description "A list of integers," an output description "A single integer representing the sum," and an input_generator function that produces [1,2,3], [10,-5], etc.

1. Sample Multiple Input and Output Pairs: With the transformed code and input generator in place, this step involves systematically executing the reference code with various inputs generated by the input generator. The key advantage here is correctness: since the outputs are derived from actual code execution, they are guaranteed to be 100% accurate, providing ground truth for training.

1. Assemble the Training Dataset (CodeIO): The collected I/O pairs and natural language queries are then assembled into a supervised fine-tuning (SFT) dataset. Each example consists of a prompt and a response:

• Prompt Construction: The prompt is designed to present an input/output prediction task. For output prediction, it combines the natural language query, a specific input, and the reference code. For input prediction, it combines the query, a specific output, and the reference code.
• Response Generation: The model is expected to generate the missing element (either the corresponding output or input) along with a detailed full chain-of-thought reasoning process. This entire response, including the reasoning and final answer, is generated using a powerful LLM, specifically DeepSeek V2.5 in this work.
• Example: For output prediction, the prompt might be: "Task: Sum numbers. Code: def sum_list(nums): return sum(nums) Input: [1, 2, 3]." The model's response would be: "Chain of thought: The code iterates through the list and adds each number. 1 + 2 = 3, 3 + 3 = 6. Output: 6."

1. Optionally Revise the Generated Response (CodeIO++): This final step introduces a quality assurance layer. Since the correctness of code execution is verifiable, CodeIO can identify and improve faulty generations.

• Verification:
• For output prediction, the model's predicted output is directly compared to the ground truth output obtained in Step 3.
• For input prediction, the reference code is re-executed with the model's predicted input, and the resulting output is checked against the ground truth output provided in the prompt.
• Revision Process: In practice, approximately 50% of the initial responses require revision. Incorrect responses are fed back to DeepSeek V2.5 along with explicit verifier feedback (e.g., "Your output 5 for input [1,2,3] is incorrect, expected 6.") to prompt the LLM to regenerate a better answer.
• Final Data: Both the original (if correct) and the revised (if needed) versions are concatenated into a single, long-form response, forming the CodeIO++ dataset.
• The final combined dataset of CodeIO and CodeIO++ comprises approximately 3.5 million high-quality instances derived from nearly half a million raw code files.

Training Strategy

CodeIO employs a two-stage training strategy to maximize the development of robust reasoning capabilities and ensure instruction-following ability:

1. Stage 1: Reasoning Capability Development: The model is exclusively trained on the CodeIO and CodeIO++ datasets. This stage focuses on grounding the model in real code logic and developing strong, verifiable reasoning capabilities through the input/output prediction tasks.
2. Stage 2: Instruction Fine-tuning: After Stage 1, training continues on broader supervised fine-tuning (SFT) data. This stage is crucial for aligning the model with general instruction-following and making the learned reasoning skills applicable across a wider array of natural language tasks. This sequential approach ensures that the model first acquires deep reasoning skills and then learns how to effectively apply them based on human instructions.

Experimental Setup & Results

▶ Watch: Two-stage training strategy and data scale (5:39)

To validate the effectiveness of CodeIO, a comprehensive experimental setup was devised, involving multiple model configurations, diverse benchmarks, and rigorous ablation studies.

Models and Evaluation:

The study trained and evaluated four different LLMs, varying in size and architectural specifics (though precise model architectures and parameter counts were not detailed in the transcript, they were noted as "varying size and architecture"). These models were tested on a suite of 14 distinct reasoning tasks spanning various domains, including math, logic, science, and tabular reasoning, in addition to code-related tasks.

Baselines:

To isolate the impact of CodeIO's unique reasoning simulation, two types of baselines were established:

1. Stage 2 SFT Only: Models trained solely with the standard supervised fine-tuning data, without any prior CodeIO-style pre-training. This represents a typical LLM training pipeline.
2. Alternative Stage 1 Pre-training: Models where the CodeIO Stage 1 pre-training was replaced with alternative strategies:

• Large-scale instruction tuning on STEM (Science, Technology, Engineering, Math) or Computer Science datasets. This evaluates if generic domain-specific instruction tuning can achieve similar benefits.
• Continue pre-training on raw code files. This tests if simply feeding raw code, without the structured I/O prediction, can be as effective.

Headline Results:

• CodeIO's Dominance: CodeIO-trained models consistently and significantly outperformed all baselines across the majority of the 14 benchmarks. This clearly demonstrated the superior efficacy of the CodeIO approach in fostering generalizable reasoning.
• CodeIO++ Advantage: The refined CodeIO++ dataset, incorporating the revision step, consistently yielded even better performance than CodeIO, underscoring the value of automated verification and iterative refinement in data generation.

Key Ablation Studies:

Two critical ablation studies provided deeper insights into CodeIO's performance drivers:

1. Impact of Error Correction: Surprisingly, the study found that keeping all responses, even those with initial errors before revision, performed almost as well as efforts like reject sampling or filtering out incorrect generations. This suggests that the model can still learn effectively even from "imperfect" initial chain-of-thought attempts, especially when paired with the ground truth. However, the revision process in CodeIO++ still provided a clear net gain.
2. Role of Prediction Direction: The research analyzed the impact of predicting only the input, only the output, or both. It was found that models trained to predict only one direction (unidirectional) performed worse than the final setup. Predicting both input and output together resulted in the best performance, highlighting the crucial value of bidirectional reasoning for a comprehensive understanding of the code's logic.

Synthesizer Control Experiment:

A natural question was whether CodeIO's strong performance simply stemmed from using a powerful synthesizer like DeepSeek V2.5 for generating the chain-of-thought responses. To address this, a controlled experiment was conducted:

• The Web Instruction Data (WIDS), a common SFT dataset, was regenerated using the same DeepSeek V2.5 model, creating WIDS 2.5.
• While WIDS 2.5 performed better than the original WIDS (synthesized by models like Qwen or Mistral), it still significantly underperformed CodeIO. This crucial finding demonstrates that while a better synthesizer improves data quality, the fundamental advantage of CodeIO lies in its better reasoning prompts and condensed reasoning patterns, not just the underlying model used for data synthesis.

Scaling and Training Dynamics:

• Data Scaling: Experiments showed that CodeIO scales well with data size. Increasing the number of raw code files or the number of input/output pairs generated per code sample consistently led to improved performance. This indicates a clear path for further gains by expanding the dataset.
• Comparison to Code Generation/Raw Code Pre-training: Using the same raw code files, CodeIO demonstrated much stronger performance compared to traditional code generation supervision or raw code continue pre-training, reinforcing the value of its structured I/O prediction approach.
• Multi-round Revision (CodeIO+++): Attempts to allow for a second round of revision (CodeIO+++) for persistently incorrect responses yielded inconsistent results, often showing no further gains or even slight declines. Consequently, the final CodeIO++ pipeline retains only one round of revision as optimal.
• Two-Stage Training Necessity: The two-stage training strategy (CodeIO first, then broader SFT) was confirmed as necessary, consistently outperforming single-stage training, highlighting the importance of dedicated reasoning capability development before general instruction tuning.

Reinforcement Learning for Verifiable Reward (RLVR) Performance:

Finally, CodeIO's utility extends to RL settings. Leveraging its verifiable nature, CodeIO data was used in an RLVR context. The results showed non-trivial gains not only on code-related tasks but also across math, science, logic, and tabular reasoning. This striking finding underscores CodeIO's ability to teach truly generalizable reasoning skills that transfer effectively even in reinforcement learning paradigms, mirroring its success in supervised fine-tuning.

Practical Implications

▶ Watch: Ablation studies: impact of errors and prediction direction (7:14)

CodeIO presents several significant practical implications for the development and deployment of advanced LLMs, impacting practitioners, infrastructure teams, and model builders.

For Practitioners and Model Builders

• Diverse Reasoning Data Generation: CodeIO offers a robust and scalable methodology for creating high-quality, diverse reasoning data, particularly for domains where curated datasets are scarce. This is a game-changer for model builders looking to move beyond the current bias towards math and code generation, enabling LLMs to develop more generalized intelligence.
• Improved Generalization: By distilling reasoning patterns from "code in the wild" and expressing them in natural language chain-of-thought, CodeIO allows LLMs to learn skills that are highly transferable across a wide array of tasks and domains. Practitioners can leverage this to build models that are more versatile and less prone to "catastrophic forgetting" when fine-tuned on new tasks.
• Verifiable Learning: The inherent verifiability of code execution provides a strong mechanism for quality control during data generation. This means the training data is grounded in objective correctness, leading to more reliable and trustworthy reasoning capabilities in the trained models.
• Targeted Skill Development: The two-stage training strategy (CodeIO first, then broader SFT) provides a clear blueprint for developing strong core reasoning abilities before integrating them with general instruction-following. This allows for more targeted and efficient skill acquisition.

For Infrastructure Teams and Deployers

• Automated Data Pipeline: The CodeIO framework describes an automated pipeline for transforming raw code, generating inputs, executing code, and verifying outputs. Infrastructure teams can adapt this pipeline to continually generate fresh, high-quality reasoning data, potentially from internal codebases or specific domain-specific code relevant to their applications.
• Specialized Model Development: Organizations with proprietary code or domain-specific logic can use the CodeIO methodology to build specialized LLMs that excel in reasoning within their unique operational contexts. This can lead to more accurate code assistants, automated debugging tools, or domain-specific problem solvers.
• RL-Ready Data: The demonstrated efficacy of CodeIO data in RLVR settings opens avenues for deploying LLMs in environments where verifiable rewards are crucial, such as automated system control, complex decision-making, or highly constrained problem-solving tasks.

Tradeoffs and Limitations

While powerful, CodeIO is not without its tradeoffs and limitations:

• Computational Cost of Data Generation: The data generation pipeline, involving code transformation, input generation, execution, chain-of-thought synthesis by a powerful LLM (DeepSeek V2.5), and an optional revision loop, is computationally intensive. Scaling this process to millions of examples requires significant computational resources.
• Reliance on Strong Synthesizer LLM: The quality of the natural language chain-of-thought reasoning is heavily dependent on the capabilities of the underlying LLM used for synthesis (e.g., DeepSeek V2.5). While CodeIO's method transcends the synthesizer's raw power, a weaker synthesizer would likely degrade the quality of the generated reasoning.
• Code Amenability: The approach relies on "code in the wild" that can be meaningfully transformed into executable, input-output predictable units. Some types of code, such as those heavily involved in UI rendering, complex system interactions, or highly stateful processes with non-deterministic outputs, might be more challenging to fit into the CodeIO paradigm.
• Reasoning Type: While CodeIO promotes generalizable reasoning, the nature of this reasoning is intrinsically tied to the logic embedded in programming constructs. It might be less effective for forms of reasoning that are purely abstract, intuitive, or heavily reliant on real-world common sense not easily derivable from code execution.
• Complexity of Pipeline Management: Setting up and maintaining the multi-stage data generation and training pipeline requires expertise in both ML engineering and software development.

Despite these considerations, CodeIO offers a compelling path forward for developing more intelligent, versatile, and robust LLMs by systematically leveraging the rich, yet underexplored, reasoning patterns embedded in real-world code.

Key Takeaways

• Novel Data Generation: CodeIO introduces a unique framework for generating high-quality, diverse reasoning data by training LLMs to predict both the input and output of code execution in natural language chain-of-thought.
• Enhanced General Reasoning: The method significantly improves LLM reasoning capabilities across 14 diverse tasks, moving beyond the traditional bias towards math and code generation and demonstrating strong generalization.
• Bidirectional Prediction is Key: Ablation studies reveal that training models to predict both input and output (bidirectional reasoning) is crucial for optimal performance, leading to a more comprehensive understanding of logic.
• Beyond Synthesizer Power: CodeIO's effectiveness stems primarily from its novel data construction pipeline and structured reasoning prompts, rather than solely relying on the power of the underlying LLM used for chain-of-thought synthesis.
• Verifiable Revision for Quality: The CodeIO++ dataset, incorporating a single round of automated, verifiable revision, consistently yields better performance, highlighting the value of iterative refinement in data generation.
• Versatile Learning Paradigm: CodeIO data proves effective not only in supervised fine-tuning but also in Reinforcement Learning for Verifiable Reward (RLVR), showing non-trivial gains across multiple domains and underscoring its broad applicability.

About the Speaker(s)

The CodeIO paper was authored by Junlong Li, Daya Guo, Dejian Yang, Runxin Xu, Yu Wu, and Junxian He. The presentation at ICML 2025 was delivered by a proxy on behalf of the authors, as none of them were able to attend in person due to visa constraints. Junlong Li is credited as the first author and the primary creator of the presentation slides. No specific affiliations or titles for the authors were mentioned within the provided transcript.