---
title: "One-Pass to Reason: 多轮推理的高效单遍微调"
authors: "Ritesh Goru, Shanay Mehta, Prateek Jain (DevRev)"
venue: "ICML 2025 Workshop — Efficient Systems for Foundational Models"
arxiv: "2504.18246"
code: "https://github.com/devrev/One-Pass-to-Reason"
dataset: "https://huggingface.co/datasets/devrev-research/MathChatSync-reasoning"
year: 2025
type: paper
tags: [efficient-fine-tuning, multi-turn-reasoning, attention-mask]
---

# One-Pass to Reason

> **核心思想**：通过 token 复制 + 分块稀疏注意力掩码，将多轮推理对话的 N 遍训练压缩为单遍，时间复杂度从 O(N³) 降至 O(N²)。

## 问题背景

推理模型（如 DeepSeek-R1）遵循行业惯例：生成推理 token → 输出回复 → 在后续轮次中**丢弃推理 token**。这导致多轮对话微调时，每个对话需要 N 次独立前向传播（N = 对话轮数）。

两个核心约束：
1. **[[visibility-constraint|可见性约束]]**：推理 token 在生成时必须可见，但在后续轮次中必须隐藏
2. **[[position-id-discrepancy|位置 ID 偏差]]**：回复 token 在生成时紧跟推理 token，但在上下文中的位置紧接着人类消息

## 方法

### Token 复制 ([[token-duplication]])

将每个助手回复的 response token 复制为两份：
- **ri_in**（上下文副本）：不关注推理 token，作为后续轮次的纯上下文
- **ri_out**（生成副本）：关注推理 token，参与 loss 计算

### 分块稀疏注意力掩码 ([[block-sparse-attention]])

定义每种 token 类型（hi, ti, ri_in, ri_out）的可见性规则：
- `hi → A(H<i)` — 人类消息只看历史
- `ti → A(H<i, hi)` — 推理 token 看历史+当前人类消息
- `ri_in → A(H<i, hi)` — 上下文副本不看推理
- `ri_out → A(H<i, hi, ti)` — 生成副本看全部包括推理

### 位置 ID 策略

```python
s_ti = s_ri_in = e_hi + 1   # 推理和上下文副本从人类消息后开始
s_ri_out = e_ti + 1          # 生成副本从推理后开始
s_h_{i+1} = e_ri_in + 1     # 下一轮人类消息从上下文副本后开始
```

### 理论保证（[[one-pass-fine-tuning|Theorem 2.1]]）

**1-Pass 与 N-Pass 的 loss 完全等价**：
$$L_{\text{N-Pass}}(c) = L_{\text{1-Pass}}(c)$$

证明分三部分：位置编码等价 → 注意力模式等价 → loss 函数等价。

## 复杂度分析

| 方法 | 时间复杂度 | 空间复杂度 |
|------|-----------|-----------|
| N-Pass | O(N³ℓ²d) | O(Nℓ²) |
| **1-Pass** | **O(N²ℓ²d)** | O(Nℓ²) |

响应 token 复制带来约 33% 的额外内存开销（因为 ri 被存了两份），但渐进复杂度相同。

## 实验结果

在 Qwen-3 (4B / 8B / 32B) 上使用 QLoRA + 8×H100：

**训练加速**（Flex-Pack-1-Pass vs FA2-Pack-N-Pass）：
- 4B: **1.05×**
- 8B: **1.21×**
- 32B: **1.22×**

vs FlexAttention N-Pass：**1.44×–1.54×**

**深度扩展**：对话越长加速越明显（验证了 O(N²) vs O(N³) 理论优势）

**K-Pass 中间方案** ([[k-pass-training]])：
- K=1：最快，+33% 内存
- K=2：1.30×–1.37× 加速，+20% 内存
- K>4：收益递减

## 数据集

**[[mathchatsync-reasoning|MathChatSync Reasoning]]**：首个公开的多轮推理数据集，基于 MathChatSync，用 GPT-4.1-mini 为每个助手回复生成推理 token。

## 实现细节

- 基于 LLaMA-Factory ([[llama-factory]])
- 使用 [[flex-attention|PyTorch FlexAttention]]（FlashAttention-2 不支持自定义掩码）
- 掩码生成在 GPU 上向量化执行，用卡诺图化简布尔逻辑
- 支持序列打包 ([[sequence-packing]]) 叠加自定义掩码

## 关键洞察

1. **从 O(N³) 到 O(N²)** 的复杂度降低意味着：对话越长，单遍训练的优势越大
2. Token 复制的本质是**用空间换时间**：多存一份 response 换来一个数量级的加速
3. K-Pass 提供了一个优雅的连续统：从完全节省内存（N-Pass）到完全节省时间（1-Pass）

## 相关概念

- [[deepseek-r1]] — 典型推理模型
- [[qlora]] — 实验所用的高效微调方法
- [[flash-attention]] — 快速注意力实现
- [[llama-factory]] — 微调框架
- [[multi-turn-reasoning]] — 多轮推理训练问题域