Transformers 中支持 4D 掩码
随着最近合并的 PR*，Transformers 现在可以接受自定义的
4D 注意力掩码作为 .forward() 方法的参数。为什么是 4D，它带来了哪些机会？
理解 Transformers 中的注意力掩码
在自然语言处理领域，Transformers 中的注意力掩码通常是一个微妙但至关重要的组成部分。我们通常很少关注注意力掩码（双关语？），并传入 None，这是许多应用程序的默认参数。由分词器生成的标准注意力掩码是 2D 张量，形状为 [batch_size, total_sequence_length]，填充为 1。这种格式表示序列中简单的线性注意力机制。
然而，深入研究 Transformer 架构，你会发现这些掩码会经历一个转换（另一个双关语？）。在模型内部，2D 张量掩码会变成 4D 张量，形状为 [batch_size, heads, input_ids_length, total_sequence_length]。这种格式允许更细致的注意力策略，例如因果解码，它使用一个全 1 的下三角矩阵，有时在存在键值（KV）缓存时还会补充一个全 1 的矩形。这种结构确保每个 token 只关注自身及其前面的邻居。
tensor([[[[1, 1, 1, 1, 0, 0, 0, 0],
[1, 1, 1, 1, 1, 0, 0, 0],
[1, 1, 1, 1, 1, 1, 0, 0],
[1, 1, 1, 1, 1, 1, 1, 0],
[1, 1, 1, 1, 1, 1, 1, 1]]]])
定制 4D 掩码的出现
传统上，这些 4D 掩码是内部表示，用户无法直接访问或修改。最近合并的拉取请求允许将自定义 4D 掩码作为参数传递给 Transformers 中的 .forward() 方法，这标志着一个重大进步。此功能允许更复杂的注意力模式，这对于需要非顺序 token 处理或单个输入张量中包含多个序列的任务至关重要。
重要提示：如果向函数提供自定义掩码，很可能需要提供一个形状为 [batch_size, total_sequence_length] 的 position_ids 张量——它将允许位置编码调整自身以适应非顺序的 token 顺序。并非所有模型都接受自定义 position_ids 参数，但许多现代模型（如 Llama）都接受。其他模型可能需要打补丁。
以下是几个使用 4D 掩码改进内存和/或效率的用例
1. 内存高效的束搜索
考虑束搜索，这是一种在语言模型中生成文本的常见策略。传统上，具有共同前缀和不同结尾的序列批次是分开处理的，消耗大量内存。通过 4D 掩码，这些序列可以紧凑地表示在一个序列中，从而显着节省内存。
假设在束搜索的第 1 步我们得到 2 个束
- 猫坐在**那**
- 猫坐在**我**
在第 2 步，我们评估下一个位置的 4 个候选
- 猫坐在垫子**上**
- 猫坐在地板**上**
- 猫坐在我的椅子**上**
- 猫坐在我的桌子**上**
通常，这些将被评估为大小为 4、长度为 5 的批次。以下是如何将它们打包成长度为 9 的 1 个序列
| token | 猫 | 坐 | 在 | 那 | 我 | 垫 | 地板 | 椅子 | 桌子 |
|---|---|---|---|---|---|---|---|---|---|
| position_ids | 0 | 1 | 2 | 3 | 3 | 4 | 4 | 4 | 4 |
| 蒙版 | 1 | 1 | 1 | 1 | . | . | . | . | . |
| 1 | 1 | 1 | . | 1 | . | . | . | . | |
| 1 | 1 | 1 | 1 | . | 1 | . | . | . | |
| 1 | 1 | 1 | 1 | . | . | 1 | . | . | |
| 1 | 1 | 1 | . | 1 | . | . | 1 | . | |
| 1 | 1 | 1 | . | 1 | . | . | . | 1 |
假设共同前缀已经处理并添加到 KV 缓存中，我们需要将以下内容传递给 model.forward()
input_ids作为分词后的序列垫 地板 椅子 桌子position_ids作为形状为(1, 9)的张量，如上所示attention_mask形状为(1, 1, 5, 9)kv_cache长度等于 5 个 token（那些position_id <= 3的 token）。
结果是内存节省了大约 2 倍，随着使用更长、更多的束，节省的内存会更多。
此方法来自 Xupeng Miao 等人的论文 SpecInfer: 加速生成式大型语言模型服务与推测推理和 Token 树验证，作者在其中引入了拓扑感知因果掩码，以便在推测解码中构建 token 候选树时将所有 token 的树注意力计算融合到单个内核中。
2. SFT（有监督微调）训练中的序列打包。
基于此处的讨论：https://github.com/huggingface/trl/issues/805。
当使用不同长度的序列对语言模型进行微调时，标准方法是使用序列结束 (EOS) token 分隔序列，如 T5 论文中介绍的那样。然而，这并不能阻止跨序列注意力。通过使用 4D 掩码，即使序列打包在一起，注意力也可以仅限于单个序列。这种方法显著降低了序列之间交叉污染的风险。
请参见下图了解方法说明
3. 前瞻解码。
Yichao Fu 等人提出的新方法“前瞻解码”，在“使用前瞻解码打破 LLM 推理的顺序依赖”一文中，将 token 生成和验证融合在一个pass中。这项技术，以前需要自定义 forward 函数来操作 token 注意力，现在可以使用 4D 注意力掩码进行简化。在 4D 掩码实现之前，它需要自定义 forward 函数，并对 token 的注意力进行精细操作。现在只需提供正确的掩码即可完成此操作。
注意力掩码的未来
展望未来，注意力掩码的演变可能不会止步于四维。引入第五维，表示层索引，可能会开启新的实验可能性，例如特定于层的注意力模式。另一个值得探索的领域是 head 维度，它允许像 Big Bird 那样进行稀疏注意力。
结论
在 Transformers 中集成自定义 4D 注意力掩码为处理复杂语言任务提供了更大的灵活性和效率。通过启用更复杂的注意力机制，这些掩码为新应用和现有应用的性能改进打开了大门。
- 感谢 Arthur Zucker、UniverseFly、KexinFeng、PhilJd、shentianxiao 对 4D 掩码 PR 的支持、建议和测试。
