Z1: Escalado eficiente en tiempo de prueba con código

Z1: Efficient Test-time Scaling with Code

April 1, 2025

Autores: Zhaojian Yu, Yinghao Wu, Yilun Zhao, Arman Cohan, Xiao-Ping Zhang

cs.AI

Resumen

Los Modelos de Lenguaje a Gran Escala (LLMs, por sus siglas en inglés) pueden lograr una mejora en la resolución de problemas complejos mediante el escalado computacional en tiempo de prueba, aunque esto suele implicar contextos más largos y un mayor costo en tokens de razonamiento. En este artículo, proponemos un método eficiente de escalado en tiempo de prueba que entrena a los LLMs en trayectorias de razonamiento relacionadas con código, facilitando la reducción de tokens de pensamiento excesivos mientras se mantiene el rendimiento. Primero, creamos Z1-Code-Reasoning-107K, un conjunto de datos curado que incluye problemas de codificación simples y complejos junto con sus trayectorias de solución cortas y largas. Segundo, presentamos una Ventana de Pensamiento Desplazado (Shifted Thinking Window) para mitigar la sobrecarga de pensamiento excesivo al eliminar etiquetas delimitadoras de contexto (por ejemplo, <think>. . . </think>) y limitar los tokens de razonamiento. Entrenado con datos de trayectorias largas y cortas y equipado con la Ventana de Pensamiento Desplazado, nuestro modelo, Z1-7B, demuestra la capacidad de ajustar su nivel de razonamiento según la complejidad de los problemas y exhibe un escalado eficiente en tiempo de prueba en diversas tareas de razonamiento, igualando el rendimiento de R1-Distill-Qwen-7B con aproximadamente el 30% de sus tokens de pensamiento promedio. Cabe destacar que, ajustado únicamente con trayectorias de código, Z1-7B muestra generalización a tareas de razonamiento más amplias (47.5% en GPQA Diamond). Nuestro análisis sobre la elicitación eficiente de razonamiento también proporciona valiosas ideas para futuras investigaciones.

English

Large Language Models (LLMs) can achieve enhanced complex problem-solving through test-time computing scaling, yet this often entails longer contexts and numerous reasoning token costs. In this paper, we propose an efficient test-time scaling method that trains LLMs on code-related reasoning trajectories, facilitating their reduction of excess thinking tokens while maintaining performance. First, we create Z1-Code-Reasoning-107K, a curated dataset of simple and complex coding problems paired with their short and long solution trajectories. Second, we present a novel Shifted Thinking Window to mitigate overthinking overhead by removing context-delimiting tags (e.g., <think>. . . </think>) and capping reasoning tokens. Trained with long and short trajectory data and equipped with Shifted Thinking Window, our model, Z1-7B, demonstrates the ability to adjust its reasoning level as the complexity of problems and exhibits efficient test-time scaling across different reasoning tasks that matches R1-Distill-Qwen-7B performance with about 30% of its average thinking tokens. Notably, fine-tuned with only code trajectories, Z1-7B demonstrates generalization to broader reasoning tasks (47.5% on GPQA Diamond). Our analysis of efficient reasoning elicitation also provides valuable insights for future research.