Z1: Escalonamento Eficiente em Tempo de Teste com 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

Resumo

Modelos de Linguagem de Grande Escala (LLMs) podem alcançar uma resolução de problemas complexos aprimorada por meio da escalabilidade computacional em tempo de teste, embora isso frequentemente envolva contextos mais longos e custos elevados de tokens de raciocínio. Neste artigo, propomos um método eficiente de escalabilidade em tempo de teste que treina LLMs em trajetórias de raciocínio relacionadas a código, facilitando a redução de tokens de pensamento excessivos enquanto mantém o desempenho. Primeiro, criamos o Z1-Code-Reasoning-107K, um conjunto de dados curado de problemas de codificação simples e complexos, pareados com suas trajetórias de solução curtas e longas. Segundo, apresentamos uma Janela de Pensamento Deslocada (Shifted Thinking Window) para mitigar a sobrecarga de pensamento excessivo, removendo tags delimitadoras de contexto (por exemplo, <think>. . . </think>) e limitando os tokens de raciocínio. Treinado com dados de trajetórias longas e curtas e equipado com a Janela de Pensamento Deslocada, nosso modelo, Z1-7B, demonstra a capacidade de ajustar seu nível de raciocínio conforme a complexidade dos problemas e exibe uma escalabilidade eficiente em tempo de teste em diferentes tarefas de raciocínio, correspondendo ao desempenho do R1-Distill-Qwen-7B com cerca de 30% de seus tokens médios de pensamento. Notavelmente, ajustado apenas com trajetórias de código, o Z1-7B demonstra generalização para tarefas de raciocínio mais amplas (47,5% no GPQA Diamond). Nossa análise da elicitação eficiente de raciocínio também fornece insights valiosos para pesquisas futuras.

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.