Efficient Single-Pass Training for Multi-Turn Reasoning

本文提出了一种通过响应令牌复制和自定义注意力掩码来实现多轮推理对话单次前向传递训练的方法，显著提高了训练效率，同时维护了推理可见性和位置一致性。

Large Language Model, Fine-tuning, Reasoning, Efficiency, Human-AI Interaction, Multimodal Systems

Ritesh Goru, Shanay Mehta, Prateek Jain

DevRev

Generated by grok-3-mini-latest

Background Problem

大型语言模型（LLMs）通过生成显式推理来提升性能已被证明有效，尤其在数学和编码任务中。然而，在多轮推理数据集上微调LLMs时，存在独特挑战：推理令牌在生成响应后被排除，不再作为后续输入，这使得无法像多轮非推理数据集那样使用单次前向传递优化（即将整个对话一次性处理），导致训练效率低下。具体问题包括可见性约束（推理令牌仅在生成当前响应时可见，不能影响后续回合）和位置ID不一致（响应令牌在不同回合中的相对位置不匹配），这些问题在非推理设置中可以通过简单的方法避免，但会显著增加多轮推理训练的计算开销。

Method

核心思想是通过响应令牌复制和自定义注意力掩码来实现多轮推理对话的单次前向传递训练，而不牺牲可见性约束和位置一致性。具体实现步骤包括：

• 响应令牌复制：将每个助理消息的响应部分复制为两个版本——r_in（用于作为输入上下文，不关注推理令牌）和r_out（用于生成响应，关注关联的推理令牌）。
• 自定义注意力掩码：定义特定掩码规则，确保每个令牌类型（人类消息h^i、推理令牌t^i、r_in和r_out）仅关注适当的子序列，例如h^i关注{h^j, r_in^j}{j=1}^{i-1}，t^i关注{h^j, r_in^j}{j=1}^{i-1} ∪ {h_i}，r_in^i关注{h^j, r_in^j}{j=1}^{i-1} ∪ {h_i}，r_out^i关注{h^j, r_in^j}{j=1}^{i-1} ∪ {h_i} ∪ {t_i}。
• 一致位置ID分配：通过公式$\begin{aligned} \text{pos}^f(t_i) &= \text{pos}^f(r_i^{in}) = \text{pos}^l(h_i) + 1 \\ \text{pos}^f(r_i^{out}) &= \text{pos}^l(t_i) + 1 \\ \text{pos}^f(h_{i+1}) &= \text{pos}^l(r_i^{in}) + 1 \end{aligned}$来确保相对位置正确。
• 标签掩码：仅对推理令牌t^i和r_out^i计算损失，以符合推理语言建模损失$\mathcal{L}_c = -\sum_{i=1}^{N} \left( \log(P_{\Phi}(r_i | (h_j, r_j)_{j=1}^{i-1} \oplus h_i \oplus t_i)) + \log(P_{\Phi}(t_i | (h_j, r_j)_{j=1}^{i-1} \oplus h_i)) \right)$。 这种方法仅在数据预处理和注意力机制上进行修改，不改变模型架构。

Experiment

论文未提供实际的实证实验结果，而是通过理论分析评估了方法的有效性。具体实验设置包括时间和内存复杂度的比较：

• 数据集和设置：假设每个消息长度为O(ℓ)，对话有N个回合。朴素N次传递方法的时间复杂度为O(N^3 ℓ^2 d)，而提出方法为O(N^2 ℓ^2 d)，显示出在长对话时的效率优势。内存复杂度均为O(N^2 ℓ^2)，但提出方法由于令牌复制有更高的常量因子。
• 结果和分析：理论结果表明，提出方法在保持正确损失计算的前提下显著降低了计算开销，尤其当N较大时。实验设计合理，因为它直接针对效率问题进行了渐进复杂度分析，并与基准方法（N次传递）比较，符合预期。然而，论文承认缺乏实证验证，并建议未来工作包括实际训练效率的实验确认，这一点合理地突出了方法的潜在优势和局限性。

Further Thoughts

这项工作强调了注意力机制在训练效率中的关键作用，或许可以与其他优化技术如LoRA（Low-Rank Adaptation）结合，进一步减少参数更新开销；此外，在多模态模型中，类似令牌复制策略可能用于处理不同模态的可见性约束，提升跨领域泛化能力；同时，考虑到开源社区如DeepSeek-R1的进展，这方法可推广到强化学习框架中，激励模型在多轮交互中更有效地生成推理。