深入理解注意力机制中的评分函数：从Gluon教程出发

深入理解注意力机制中的评分函数：从Gluon教程出发

d2l-zh 项目地址: https://gitcode.com/gh_mirrors/d2l/d2l-zh

注意力机制的核心思想

注意力机制是现代深度学习模型中不可或缺的重要组成部分，特别是在自然语言处理和计算机视觉领域。它的核心思想是让模型能够"有选择地关注"输入数据中最相关的部分，而不是对所有部分一视同仁。这种选择性关注的能力极大地提升了模型处理长序列和复杂模式的能力。

注意力评分函数的基本原理

在注意力机制中，评分函数(Scoring Function)扮演着关键角色，它决定了查询(Query)和键(Key)之间的相关程度。具体来说，评分函数需要完成以下任务：

1. 计算查询向量q和键向量k之间的相关性得分
2. 通过softmax函数将得分转换为概率分布(注意力权重)
3. 使用这些权重对值(Value)向量进行加权求和

数学表达式可以表示为：

$$f(\mathbf{q}, (\mathbf{k}_1, \mathbf{v}_1), \ldots, (\mathbf{k}_m, \mathbf{v}m)) = \sum{i=1}^m \alpha(\mathbf{q}, \mathbf{k}_i) \mathbf{v}_i$$

其中注意力权重计算为：

$$\alpha(\mathbf{q}, \mathbf{k}_i) = \mathrm{softmax}(a(\mathbf{q}, \mathbf{k}_i))$$

掩蔽Softmax操作

在实际应用中，我们经常需要处理变长序列，这时就需要掩蔽Softmax操作(Masked Softmax)。它的主要功能是：

1. 过滤掉无效位置(如填充位置)的影响
2. 确保注意力权重仅分布在有效位置上

实现原理是将无效位置的值替换为一个极小的负数(-1e6)，这样在Softmax计算时，这些位置的输出就会趋近于0。

def masked_softmax(X, valid_lens):
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1)
    else:
        # 将无效位置掩蔽为极小的值
        X = sequence_mask(X, valid_lens, value=-1e6)
        return nn.functional.softmax(X, dim=-1)

两种重要的评分函数

1. 加性注意力(Additive Attention)

加性注意力适用于查询和键维度不同的情况。它通过一个单隐藏层的MLP来计算相关性得分：

$$a(\mathbf q, \mathbf k) = \mathbf w_v^\top \text{tanh}(\mathbf W_q\mathbf q + \mathbf W_k \mathbf k)$$

实现要点：

• 使用tanh作为激活函数
• 不包含偏置项
• 需要学习三个参数矩阵：W_q, W_k和w_v

class AdditiveAttention(nn.Module):
    def __init__(self, key_size, query_size, num_hiddens, dropout):
        super().__init__()
        self.W_k = nn.Linear(key_size, num_hiddens, bias=False)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=False)
        self.w_v = nn.Linear(num_hiddens, 1, bias=False)
        self.dropout = nn.Dropout(dropout)

2. 缩放点积注意力(Scaled Dot-Product Attention)

当查询和键的维度相同时，可以使用计算效率更高的点积注意力。为了避免点积结果随维度增大而方差变大，需要对结果进行缩放：

$$a(\mathbf q, \mathbf k) = \mathbf{q}^\top \mathbf{k} /\sqrt{d}$$

批量计算形式为：

$$\mathrm{softmax}\left(\frac{\mathbf Q \mathbf K^\top }{\sqrt{d}}\right) \mathbf V$$

实现要点：

• 计算查询和键的点积
• 除以√d进行缩放
• 应用softmax和dropout

class DotProductAttention(nn.Module):
    def __init__(self, dropout):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)
        self.attention_weights = masked_softmax(scores, valid_lens)
        return torch.bmm(self.dropout(self.attention_weights), values)

实际应用中的选择建议

1. 维度匹配：当查询和键的维度相同时，优先选择缩放点积注意力，计算效率更高；维度不同时则必须使用加性注意力。

2. 计算效率：点积注意力的计算复杂度是O(n²d)，而加性注意力是O(n²h)，其中h是隐藏层大小。当d < h时，点积更高效。

3. 表达能力：加性注意力由于有额外的非线性变换，理论上表达能力更强，适合更复杂的相关性模式。

4. 实践建议：现代Transformer架构大多采用缩放点积注意力，因为它在GPU上可以高度优化；但在某些特殊场景下，加性注意力可能表现更好。

可视化理解

通过热图可以直观地观察注意力权重的分布。例如，当所有键相同时，注意力权重会均匀分布：

d2l.show_heatmaps(attention_weights, xlabel='Keys', ylabel='Queries')

这种可视化工具对于调试和理解模型行为非常有帮助，特别是在处理复杂序列时。

总结

注意力评分函数是注意力机制的核心组件，决定了模型如何分配注意力资源。加性注意力和缩放点积注意力各有优缺点，适用于不同场景。理解这些评分函数的工作原理，对于设计和优化基于注意力机制的模型至关重要。在实际应用中，还需要考虑计算效率、内存占用和模型表现之间的平衡。

d2l-zh 项目地址: https://gitcode.com/gh_mirrors/d2l/d2l-zh