Z1 : Mise à l'échelle efficace en temps de test avec code

Z1: Efficient Test-time Scaling with Code

April 1, 2025

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

cs.AI

Résumé

Les modèles de langage de grande taille (LLMs) peuvent améliorer leur capacité à résoudre des problèmes complexes grâce à une mise à l'échelle du calcul au moment du test, mais cela implique souvent des contextes plus longs et des coûts élevés en tokens de raisonnement. Dans cet article, nous proposons une méthode efficace de mise à l'échelle au moment du test qui entraîne les LLMs sur des trajectoires de raisonnement liées au code, facilitant ainsi la réduction des tokens de pensée excessifs tout en maintenant les performances. Premièrement, nous créons Z1-Code-Reasoning-107K, un ensemble de données soigneusement sélectionné de problèmes de codage simples et complexes accompagnés de leurs trajectoires de solution courtes et longues. Deuxièmement, nous présentons une nouvelle fenêtre de pensée décalée (Shifted Thinking Window) pour atténuer la surcharge de réflexion en supprimant les balises de délimitation de contexte (par exemple, <think>. . . </think>) et en limitant les tokens de raisonnement. Entraîné avec des données de trajectoires longues et courtes et équipé de la fenêtre de pensée décalée, notre modèle, Z1-7B, démontre la capacité à ajuster son niveau de raisonnement en fonction de la complexité des problèmes et présente une mise à l'échelle efficace au moment du test pour différentes tâches de raisonnement, égalant les performances de R1-Distill-Qwen-7B avec environ 30 % de ses tokens de pensée moyens. Notamment, affiné uniquement avec des trajectoires de code, Z1-7B montre une généralisation à des tâches de raisonnement plus larges (47,5 % sur GPQA Diamond). Notre analyse de l'élicitation efficace du raisonnement fournit également des insights précieux pour les recherches futures.

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.