Z1: Scalabilità Efficiente al Momento del Test con Codice

Z1: Efficient Test-time Scaling with Code

April 1, 2025

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) possono ottenere una migliore risoluzione di problemi complessi attraverso il ridimensionamento computazionale al momento del test, ma ciò spesso comporta contesti più lunghi e costi elevati in termini di token di ragionamento. In questo articolo, proponiamo un metodo efficiente di ridimensionamento al momento del test che addestra gli LLM su traiettorie di ragionamento legate al codice, facilitando la riduzione dei token di pensiero in eccesso mantenendo le prestazioni. In primo luogo, creiamo Z1-Code-Reasoning-107K, un dataset curato di problemi di codifica semplici e complessi accoppiati con le loro traiettorie di soluzione brevi e lunghe. In secondo luogo, presentiamo una nuova Finestra di Pensiero Spostata per mitigare il sovraccarico di pensiero rimuovendo i tag di delimitazione del contesto (ad esempio, <think>. . . </think>) e limitando i token di ragionamento. Addestrato con dati di traiettorie lunghe e brevi e dotato della Finestra di Pensiero Spostata, il nostro modello, Z1-7B, dimostra la capacità di adattare il suo livello di ragionamento in base alla complessità dei problemi e mostra un ridimensionamento efficiente al momento del test su diverse attività di ragionamento che eguaglia le prestazioni di R1-Distill-Qwen-7B con circa il 30% dei suoi token di pensiero medi. È degno di nota che, addestrato solo con traiettorie di codice, Z1-7B dimostra una generalizzazione a compiti di ragionamento più ampi (47,5% su GPQA Diamond). La nostra analisi sull'elicitazione efficiente del ragionamento fornisce anche spunti preziosi per la ricerca futura.

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.