Hao Zhang

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.

Designing a unified neural network to efficiently and inherently process sequential data with arbitrary lengths is a central and challenging problem in sequence modeling. The design choices in Transformer, including quadratic complexity and weak length extrapolation, have limited their ability to scale to long sequences. In this work, we propose Gecko, a neural architecture that inherits the design of Mega and Megalodon (exponential moving average with gated attention), and further introduces multiple technical components to improve its capability to capture long range dependencies, including timestep decay normalization, sliding chunk attention mechanism, and adaptive working memory. In a controlled pretraining comparison with Llama2 and Megalodon in the scale of 7 billion parameters and 2 trillion training tokens, Gecko achieves better efficiency and long-context scalability. Gecko reaches a training loss of 1.68, significantly outperforming Llama2-7B (1.75) and Megalodon-7B (1.70), and landing close to Llama2-13B (1.67). Notably, without relying on any context-extension techniques, Gecko exhibits inherent long-context processing and retrieval capabilities, stably handling sequences of up to 4 million tokens and retrieving information from contexts up to $4\times$ longer than its attention window.

Translating natural language into SQL (Text2SQL) is a longstanding challenge at the intersection of natural language understanding and structured data access. While large language models (LLMs) have significantly improved fluency in SQL generation, producing correct and executable SQL, particular...