Technical Summary: STEP-KTODER

Problem Statement

While process supervision—providing feedback on intermediate reasoning steps—has significantly improved mathematical reasoning in Large Language Models (LLMs), its application to code generation remains underexplored. In mathematical domains, steps are naturally defined as chains of thought. In code generation, however, there is no standard notion of a "step." Supervision targets can vary widely (lines, reasoning traces, or program states), creating ambiguity in what to label and optimize.

Current preference optimization methods like Direct Preference Optimization (DPO) and Kahneman–Tversky Optimization (KTO) typically operate at the outcome level, treating an entire program as a single unit of feedback. This coarse signal fails to distinguish between correct and incorrect components within a partially successful program. For instance, a program that is mostly correct but fails on a narrow edge case is labeled entirely as "undesirable," obscuring which specific functions contributed to success or failure. This lack of granularity limits the ability to reinforce locally correct code while penalizing specific faulty components.

Methodology: STEP-KTODER

The authors propose STEP-KTODER, a framework for code preference optimization that defines "steps" as module-level functions in decomposed multi-function programs. The method combines function-level process supervision with outcome-level feedback using an execution-based labeling pipeline.

1. Data Construction Pipeline

The framework relies on an automatic data construction process consisting of four stages:

• Reference Decomposition: A strong instruction-tuned code model rewrites ground-truth solutions into decomposed, multi-function programs while preserving the original interface and behavior. Only rewrites that pass the original test suite are retained.
• Function-Level Unit Test Generation: For each decomposed function, the model generates targeted unit tests covering the function's interface and edge cases. These tests are validated by executing them against the reference implementation; only tests that pass and contain nontrivial assertions are kept.
• On-Policy Candidate Generation: The target model generates candidate completions by filling in function skeletons (where function bodies are replaced with pass) derived from the decomposed references.
• Labeling:
  • Outcome Label ($o$): Determined by whether the full candidate program passes the dataset-provided test suite.
  • Step Labels ($z$): Determined by executing each candidate function against its generated unit tests. A label of $1$ indicates local correctness, $0$ indicates failure, and $\emptyset$ indicates no valid test was available.

2. Objective Function

STEP-KTODER extends KTO by introducing a joint objective that combines outcome-level and step-level supervision.

• Outcome-Level KTO: Uses the standard KTO loss ($L_{out}$) based on the binary desirability of the full program.
• Step-Level KTO: Defines a step-wise loss ($L_{step}$) over the tokens of individual functions. For each function $s_i$, a local log-ratio is computed, and a value function is applied based on the step label $z_i$.
• Joint Objective: The final loss is $L_{STEP-KTODER} = L_{out} + \lambda_{step} L_{step}$.

Crucially, the framework employs conflict-preserving labels. It retains training examples where the outcome label ($o$) and step labels ($z$) disagree (e.g., a program that passes the global test suite but contains a locally failing function). These conflicts provide unique training signals that outcome-only methods cannot capture, allowing the model to learn that specific components can be correct even if the whole fails, or vice versa.

3. Reliability of Labels

The paper emphasizes that execution-based labels are essential. The authors tested replacing execution-based step labels with "LLM-as-a-judge" annotations (using GPT-5.4 mini). While the LLM showed moderate overall agreement (73.2%), it exhibited a systematic bias: it over-predicted failures, correctly identifying 94.1% of failing functions but only 52.2% of passing ones. This bias corrupted the positive step labels required for training, leading to degraded performance, particularly on harder benchmarks.

Key Contributions

1. STEP-KTODER Framework: A code-specific instantiation of stepwise KTO that defines steps as independently testable functions, bridging the gap between process supervision and code generation.
2. Automatic Data Pipeline: A pipeline that decomposes solutions, generates function-level unit tests, and derives execution-based step labels, enabling fine-grained supervision without manual annotation.
3. Empirical Validation: Experiments demonstrating that execution-based labels are superior to LLM-based judgments for step supervision and that preserving local/global label conflicts is critical for performance.

Experimental Results

The method was evaluated on HumanEval(+), MBPP(+), BigCodeBench, and LiveCodeBench using instruction-tuned models (Qwen2.5-Coder and DeepSeek-Coder).

• Performance Gains: STEP-KTODER consistently outperformed outcome-only KTO and DPO, with the most significant improvements observed on harder benchmarks.
  • On Qwen2.5-Coder-1.5B, it improved over KTO by +11.5% on BigCodeBench Hard and +7.7% on LiveCodeBench.
  • On Qwen2.5-Coder-3B, gains were +11.3% on BigCodeBench Hard and +9.5% on LiveCodeBench.
  • On DeepSeek-Coder-6.7B, it achieved a +4.9% gain over KTO on BigCodeBench Hard.
• Comparison with Target-DPO: While Target-DPO also targets localized failures, STEP-KTODER achieved superior results on the hardest benchmarks (e.g., +18.5% improvement over Target-DPO on BigCodeBench Hard for the 3B model) using significantly fewer training samples (approx. 11,400 vs. 59,000 pairs).
• Ablation Studies:
  • Conflict Preservation: Masking conflicts (removing examples where $o \neq z$) reduced performance to KTO levels, confirming that these mismatches provide essential training signals.
  • Label Source: Replacing execution-based labels with LLM-as-a-judge annotations caused substantial performance drops (e.g., -25.8% on BigCodeBench Hard), validating the necessity of execution-based feedback.

Significance and Claims

The paper claims that STEP-KTODER offers a practical path toward process supervision for code generation by leveraging the natural modularity of programs. By defining steps as functions and supervising them via execution, the method provides localized feedback that reinforces correct components and penalizes specific faults, a capability absent in outcome-only supervision.

The authors highlight that the effectiveness of the method relies heavily on two factors:

1. Execution-based reliability: Automated unit tests provide the necessary ground truth for step labels, whereas LLM judges introduce systematic biases that degrade learning.
2. Conflict preservation: The ability to learn from examples where local correctness and global success diverge is a key differentiator from standard preference optimization.

The work suggests that function-level execution feedback is a viable and effective mechanism for aligning code models, particularly for complex tasks where coarse-grained feedback is insufficient. The authors note limitations regarding the requirement for solutions to be decomposable into testable functions and the engineering cost of generating on-policy candidates for new models.