深入理解Annotated-Transformer中的注意力掩码机制

在Transformer架构的实现中，注意力掩码(Attention Mask)是一个关键但容易被忽视的技术细节。本文将以harvardnlp/annotated-transformer项目为例，深入解析解码器中两种不同类型的注意力机制及其掩码应用方式。

解码器层的注意力结构

Transformer解码器包含两种核心注意力模块：

1. 掩码多头注意力(Masked Multi-Head Attention)
2. 普通多头注意力(Multi-Head Attention)

这两种结构在实现上看似相似，但在功能和应用场景上存在重要差异。

掩码多头注意力的双重作用

掩码多头注意力在解码器中承担着两个关键功能：

1. 填充位置掩码(Padding Mask)：防止模型关注输入序列中的填充位置(padded positions)，这些位置通常是为了保持批次处理时序列长度一致而添加的无意义符号。

2. 序列顺序掩码(Sequence Mask)：确保解码器在预测当前位置时，只能关注该位置之前的token，而不能"偷看"未来的信息。这种掩码通常采用上三角矩阵的形式实现，对角线以上的位置被设置为负无穷。

在annotated-transformer的实现中，这种双重掩码通过tgt_mask参数实现：

x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))

普通多头注意力的作用

解码器中的第二个注意力层是标准的跨注意力机制，它连接编码器和解码器：

1. 查询(Query)来自解码器的输出
2. 键(Key)和值(Value)来自编码器的最终输出（memory）

这个注意力层只使用填充位置掩码（通过src_mask参数），不应用序列顺序掩码：

x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))

实现细节对比

特性	掩码多头注意力	普通多头注意力
输入来源	解码器自身	编码器输出
掩码类型	tgt_mask (双重掩码)	src_mask (填充掩码)
是否限制未来信息	是	否
实现位置	解码器第一个子层	解码器第二个子层

实际应用中的考虑

理解这两种注意力机制的区别对于正确实现Transformer模型至关重要：

1. 在训练阶段，必须正确应用序列顺序掩码，否则模型会通过"作弊"获得未来信息，导致评估指标虚高。

2. 在推理阶段，虽然序列顺序掩码仍然需要，但由于是逐步生成输出，实现方式会有所不同。

3. 对于不同长度的输入序列，填充掩码的处理需要特别注意，特别是在批处理场景下。

通过深入分析annotated-transformer的实现，我们可以更清楚地理解Transformer架构中这些关键组件的设计理念和实现细节，为后续的模型修改和优化奠定坚实基础。