$\textbf{PLUM}$: Improving Code LMs with Execution-Guided On-Policy Preference Learning Driven By Synthetic Test Cases

TL;DR

PLUM introduces an execution-guided, on-policy preference learning framework that uses automatically generated test cases to train code language models without reward models. By sampling solutions on-policy, evaluating them with generated test suites, and applying preference learning (DPO or KTO), PLUM achieves consistent improvements across multiple benchmarks (HumanEval, MBPP, LiveCodeBench) and diverse model families. The approach demonstrates notable gains, especially in challenging tasks, and enables online, iterative alignment without external labeling. This work highlights the practicality and generality of test-case-based feedback for scaling functional correctness in code LMs and reduces distributional drift by staying in-distribution during learning.

Abstract

Preference learning provides a promising solution to address the limitations of supervised fine-tuning (SFT) for code language models, where the model is not explicitly trained to differentiate between correct and incorrect code. Recent findings demonstrate that on-policy data is the key to successful preference learning, where the preference data is collected using the same policy LM being trained. Inspired by this, we propose PLUM, an on-policy $\textbf{P}$reference $\textbf{L}$earning framework A$\textbf{u}$gmented with test cases for code L$\textbf{M}$ s. The framework operates in three key stages: (1) automatic generation of test cases from natural language instructions, (2) creation of a preference data by evaluating candidate code solutions sampled from the policy, which can then be used to (3) train the policy LM. PLUM levitates the need to train reward models, allowing for large scale on-policy and online preference data collation. PLUM is evaluated on both standard benchmarks (HumanEval, MBPP) and more challenging ones (LiveCodeBench), delivering substantial improvements over original SFT'ed models and other execution-feedback-driven approaches. We show PLUM's benefits are consistent across various widely-used code LMs even they have been well-trained with SFT. For example, PLUM increases pass rates by up to 4.8% on average on standard benchmarks and 11.8% on LiveCodeBench, demonstrating its effectiveness and generalizability. We also demonstrate the benefits of on-policy and online preference learning by comprehensive experimentation.

Figures (3)

• Figure 1: Overview of PLUM. It involves three steps: (1) Generating the test cases; (2) Sampling solutions from the policy and evaluating them against the test cases to collect the preference data for (3) performing preference learning.
• Figure 2: Ablations on preference training signals reveal that using un-runnable code as negative instances does not consistently lead to performance improvements. The baseline results refer to the SFT-model without PLUM.
• Figure 3: Distribution of policy model correctness ratio on OSS-Instruct dataset.