2.2 KiB
2.2 KiB
title, created, updated, type, tags, sources
| title | created | updated | type | tags | sources | ||||
|---|---|---|---|---|---|---|---|---|---|
| Block-Sparse Attention Mask (分块稀疏注意力掩码) | 2025-06-02 | 2025-06-02 | concept |
|
|
Block-Sparse Attention Mask
goru-one-pass-to-reason-2025 中设计的自定义注意力掩码，通过为不同 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) — 看全部 |
结构特点
- 分块稀疏：整个掩码按对话轮次和 token 类型划分为 block，每个 block 有不同的稀疏模式
- 因果性保持：所有 token 仍遵循因果约束（不偷看未来）
- 选择性可见：核心创新——ri_in 和 ri_out 虽然内容相同，但对 ti 的可见性不同
与标准掩码的区别
- 标准因果掩码：所有 token 看到所有之前的 token
- Block-Sparse Mask：特定 token 类型跳过特定块（ri_in 跳过 ti）
高效生成
因为掩码需要在 GPU 上为每个 batch 动态生成，论文提出了向量化生成算法：
# 伪代码
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（Dong et al., 2024），因为 FlashAttention-2 不支持传入自定义掩码。
与 packing 的兼容
可与 sequence-packing 叠加：打包掩码（防跨样本污染）与自定义注意力掩码通过逻辑 AND 结合。