English

One-Pass to Reason: Token Duplication and Block-Sparse Mask for Efficient Fine-Tuning on Multi-Turn Reasoning

Computation and Language 2025-07-14 v2 Artificial Intelligence Machine Learning

Abstract

Fine-tuning Large Language Models (LLMs) on multi-turn reasoning datasets requires N (number of turns) separate forward passes per conversation due to reasoning token visibility constraints, as reasoning tokens for a turn are discarded in subsequent turns. We propose duplicating response tokens along with a custom attention mask to enable single-pass processing of entire conversations. We prove our method produces identical losses to the N-pass approach while reducing time complexity from O(N3)O\bigl(N^{3}\bigl) to O(N2)O\bigl(N^{2}\bigl) and maintaining the same memory complexity for a transformer based model. Our approach achieves significant training speedup while preserving accuracy. Our implementation is available online (https://github.com/devrev/One-Pass-to-Reason).

Keywords

Cite

@article{arxiv.2504.18246,
  title  = {One-Pass to Reason: Token Duplication and Block-Sparse Mask for Efficient Fine-Tuning on Multi-Turn Reasoning},
  author = {Ritesh Goru and Shanay Mehta and Prateek Jain},
  journal= {arXiv preprint arXiv:2504.18246},
  year   = {2025}
}

Comments

9 pages, 3 figures