KTPO: K-Step Test-Time Policy Optimization for Long- Horizon Discovery
Dacheng Li · Fangzhou Zhao · Shiyi Cao · Hao Zhang · Joseph E. Gonzalez · Ion Stoica
We propose KTPO, a novel algorithm that integrates RL and Test time scaling methods for discover problems.
摘要
Large Language Models (LLMs) are emerging as a promising approach for discovery problems in mathematics and systems optimization. Existing search algorithms keep models frozen, yet applying reinforcement learning (RL) over full search trajectories spanning hundreds of refinement steps remains difficult. We present KTPO (K-Step Test-Time Policy Optimization), which addresses this challenge by training on short trajectories of k refinement steps. To retain the benefits of longer search, each trajectory starts from a previously discovered program rather than from scratch. In particular, KTPO selects programs near the model’s current capability frontier, estimated by the model’s improvement rate over existing programs. KTPO with 4B-9B models achieves strong performance across seven mathematical and systems optimization tasks. On mathematical problems, KTPO outperforms or approaches the human record on Circle Packing and the Erd˝os minimum overlap problem. On systems optimization tasks from ADRS, a benchmark suite for real-world systems optimization, KTPO out- performs human record on maximizing SQL query prefix hit rate, and surpasses the prior AI records on minimizing multi-cloud data transfer cost and on balancing expert-parallel load across GPUs. Code is available at https://github.com/DachengLi1/KTPO.