Z1: コードを用いた効率的なテスト時スケーリング

Z1: Efficient Test-time Scaling with Code

April 1, 2025

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

cs.AI

要旨

大規模言語モデル（LLM）は、テスト時の計算スケーリングを通じて複雑な問題解決能力を向上させることが可能ですが、これにはしばしば長いコンテキストと多数の推論トークンのコストが伴います。本論文では、コード関連の推論軌跡を用いてLLMを訓練し、性能を維持しながら過剰な思考トークンを削減する効率的なテスト時スケーリング手法を提案します。まず、単純および複雑なコーディング問題とそれらの短いおよび長い解決軌跡をペアにしたキュレーションデータセット、Z1-Code-Reasoning-107Kを作成します。次に、コンテキストを区切るタグ（例：<think>. . . </think>）を削除し、推論トークンを制限することで、過剰思考のオーバーヘッドを軽減する新しいShifted Thinking Windowを提示します。長いおよび短い軌跡データで訓練され、Shifted Thinking Windowを備えた私たちのモデル、Z1-7Bは、問題の複雑さに応じて推論レベルを調整する能力を示し、R1-Distill-Qwen-7Bの性能に匹敵する異なる推論タスクにおいて、平均思考トークンの約30%で効率的なテスト時スケーリングを実現します。特に、コード軌跡のみでファインチューニングされたZ1-7Bは、より広範な推論タスク（GPQA Diamondで47.5%）への一般化能力を示します。効率的な推論の引き出しに関する私たちの分析は、今後の研究にとって貴重な洞察を提供します。

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.