Z1: 코드를 활용한 효율적인 테스트 타임 스케일링

Z1: Efficient Test-time Scaling with Code

April 1, 2025

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

cs.AI

초록

대규모 언어 모델(LLMs)은 테스트 시간 계산 확장을 통해 향상된 복잡 문제 해결 능력을 달성할 수 있지만, 이는 종종 더 긴 컨텍스트와 많은 추론 토큰 비용을 수반합니다. 본 논문에서는 코드 관련 추론 궤적을 통해 LLMs를 훈련시켜 과도한 사고 토큰을 줄이면서도 성능을 유지하는 효율적인 테스트 시간 확장 방법을 제안합니다. 먼저, 간단하고 복잡한 코딩 문제와 그에 대한 짧고 긴 해결 궤적을 짝지은 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.