🤖 AI Summary
In multi-turn reasoning fine-tuning, reasoning tokens from prior turns are excluded from subsequent inputs, preventing single-pass forward propagation and inducing computational redundancy. Method: We propose a response token replication mechanism coupled with a customized attention masking scheme. During training, historical response tokens are explicitly reused as input, while structured attention masks restrict attention computation to valid contextual spans. Contribution/Results: This enables end-to-end, single-pass forward propagation over multi-turn reasoning data for the first time. The approach is fully compatible with standard Transformer architectures—requiring no architectural modifications—while preserving inference capability. It significantly reduces training latency, improves fine-tuning efficiency and scalability, and establishes a new paradigm for efficiently aligning large language models with multi-turn reasoning capabilities.
📝 Abstract
Training Large Language Models ( LLMs) to generate explicit reasoning before they produce an answer has been shown to improve their performance across various tasks such as mathematics and coding. However, fine-tuning LLMs on multi-turn reasoning datasets presents a unique challenge: LLMs must generate reasoning tokens that are excluded from subsequent inputs to the LLM. This discrepancy prevents us from processing an entire conversation in a single forward pass-an optimization readily available when we fine-tune on a multi-turn non-reasoning dataset. This paper proposes a novel approach that overcomes this limitation through response token duplication and a custom attention mask that enforces appropriate visibility constraints. Our approach significantly reduces the training time and allows efficient fine-tuning on multi-turn reasoning datasets.