63 lines
2.2 KiB
Markdown
63 lines
2.2 KiB
Markdown
---
|
||
title: "Block-Sparse Attention Mask (分块稀疏注意力掩码)"
|
||
created: 2025-06-02
|
||
updated: 2025-06-02
|
||
type: concept
|
||
tags: [attention, training-optimization, multi-turn-reasoning]
|
||
sources: ["[[goru-one-pass-to-reason-2025]]"]
|
||
---
|
||
|
||
# Block-Sparse Attention Mask
|
||
|
||
> [[goru-one-pass-to-reason-2025|One-Pass to Reason]] 中设计的自定义注意力掩码，通过为不同 token 类型指定不同的可见性规则，在单次前向传播中同时满足"生成时可见、上下文时隐藏"的冲突需求。
|
||
|
||
## 可见性规则
|
||
|
||
在多轮推理对话中，token 被分为四种角色：
|
||
|
||
| Token 类型 | 角色 | 可见范围 |
|
||
|-----------|------|---------|
|
||
| hi | 人类消息 | A(H_{<i}) — 只看历史 |
|
||
| ti | 推理 token | A(H_{<i}, hi) — 看历史+当前人类消息 |
|
||
| ri_in | 上下文副本 | A(H_{<i}, hi) — 不看推理 token |
|
||
| ri_out | 生成副本 | A(H_{<i}, hi, ti) — 看全部 |
|
||
|
||
## 结构特点
|
||
|
||
1. **分块稀疏**：整个掩码按对话轮次和 token 类型划分为 block，每个 block 有不同的稀疏模式
|
||
2. **因果性保持**：所有 token 仍遵循因果约束（不偷看未来）
|
||
3. **选择性可见**：核心创新——ri_in 和 ri_out 虽然内容相同，但对 ti 的可见性不同
|
||
|
||
## 与标准掩码的区别
|
||
|
||
- **标准因果掩码**：所有 token 看到所有之前的 token
|
||
- **Block-Sparse Mask**：特定 token 类型跳过特定块（ri_in 跳过 ti）
|
||
|
||
## 高效生成
|
||
|
||
因为掩码需要在 GPU 上为每个 batch 动态生成，论文提出了向量化生成算法：
|
||
|
||
```python
|
||
# 伪代码
|
||
M[i][j] = can_see(role[i], role[j]) AND causal(i, j)
|
||
# role 映射: 0=pad, 1=human, 2=thinking, 3=response_in, 4=response_out
|
||
```
|
||
|
||
用卡诺图（Karnaugh map）化简布尔逻辑，最小化逻辑运算次数。
|
||
|
||
## 实现后端
|
||
|
||
使用 [[flex-attention|PyTorch FlexAttention]]（Dong et al., 2024），因为 FlashAttention-2 不支持传入自定义掩码。
|
||
|
||
## 与 packing 的兼容
|
||
|
||
可与 [[sequence-packing]] 叠加：打包掩码（防跨样本污染）与自定义注意力掩码通过逻辑 AND 结合。
|
||
|
||
## 相关
|
||
|
||
- [[one-pass-fine-tuning]]
|
||
- [[token-duplication]]
|
||
- [[visibility-constraint]]
|
||
- [[flex-attention]]
|
||
- [[goru-one-pass-to-reason-2025|One-Pass to Reason 论文]]
|