让大模型一次处理多轮推理对话，训练速度翻倍提升
One-Pass to Reason: Token Duplication and Block-Sparse Mask for Efficient Fine-Tuning on Multi-Turn Reasoning
- 复制回复令牌并用特殊掩码，实现单次遍历完成多轮推理
- 训练时间从O(N³)降至O(N²)，速度提升显著且准确率不变
- 适合需要高效微调多轮对话的AI研究者和工程师
在多轮推理数据集上微调大语言模型时，由于每轮推理令牌在后续轮次中被丢弃，需对每条对话执行N（轮数）次前向传播。本文提出复制回复令牌并设计定制注意力掩码，使整个对话可单次遍历处理。理论证明该方法与传统N次遍历方式损失完全一致，同时将时间复杂度从O(N³)降低至O(N²)，内存消耗保持不变。实验表明该方法在保持精度的前提下实现显著训练加速。代码已开源：https://github.com/devrev/One-Pass-to-Reason。
原文摘要 · Abstract (English)
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\bigl(N^{3}\bigl)$ to $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).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。