TensorFlow2.x实战使用Keras自定义层实现注意力机制

搭建自定义注意力层的基本结构

在TensorFlow 2.x中，使用Keras自定义层需要继承tf.keras.layers.Layer基类。对于注意力机制，我们首先需要初始化层所需的参数。在__init__方法中，我们通常定义注意力计算中需要的权重矩阵的维度，并调用父类的初始化方法。例如，我们可以初始化查询（Query）、键（Key）、值（Value）的投影矩阵所需的维度参数，并设置一个布尔值来控制是否在Softmax之前应用注意力掩码。

构建投影矩阵

接下来，在build方法中，我们根据输入张量的形状动态创建层的权重。对于自注意力机制，我们需要为查询、键和值分别创建三个密集层（Dense层）。build方法接收一个input_shape参数，它通常是输入张量的形状元组。我们使用self.add_weight方法来创建可训练权重，并指定其形状、初始化器以及名称。例如，查询权重矩阵的形状可能是(input_shape[-1], self.units)，其中self.units是在__init__中定义的输出维度。

实现注意力分数计算

注意力机制的核心是计算注意力分数。在自定义层的call方法中，我们实现前向传播逻辑。首先，使用在build方法中创建的权重矩阵，将输入张量分别投影到查询、键和值空间。然后，通过查询和键的点积计算注意力分数。为了提高数值稳定性，通常会将注意力分数除以键向量维度的平方根进行缩放。如果需要处理变长序列并忽略填充部分，可以在此步骤中应用注意力掩码，将填充位置的分数设置为一个极小的负数（如-1e9）。

实现Scaled Dot-Product Attention

缩放点积注意力是Transformer模型中使用的标准注意力机制。在call方法中，计算完查询和键的点积后，应用缩放因子。然后，对缩放后的分数矩阵应用Softmax函数，以获取注意力权重。Softmax函数沿最后一个轴应用，确保每个查询的注意力权重之和为1。最后，将注意力权重与值张量相乘，得到加权的上下文向量作为输出。

处理多头注意力机制

为了实现更强的表示能力，通常会将注意力扩展为多头。在自定义层中，这可以通过创建多个独立的“注意力头”来实现。每个头都有自己的查询、键、值投影矩阵。在build方法中，我们创建多组权重，而不是一组。在call方法中，我们将输入张量分别投影到每个头的查询、键、值空间，然后在每个头上独立计算缩放点积注意力。最终，将所有头的输出拼接起来，并通过一个最终的线性投影层（Dense层）聚合信息。

整合与输出

计算完多头注意力的输出后，我们通常需要添加残差连接和层归一化（Layer Normalization）来稳定训练过程。虽然这些步骤有时在自定义层外部实现，但也可以将它们集成到层内部。在call方法的最后，返回最终的上下文向量张量。如果该层是编码器-解码器注意力层，那么查询可能来自解码器，而键和值来自编码器，需要在调用时接收不同的输入源。

在模型中使用自定义注意力层

定义好自定义注意力层后，可以像使用任何内置Keras层一样在tf.keras.Sequential模型或函数式API中使用它。在函数式API中，将前一层的输出作为查询、键、值（对于自注意力）输入到自定义层中。如果处理的是序列数据，需要确保输入张量的形状为(batch_size, sequence_length, features)。自定义层会自动处理批量维度，并输出一个相同批大小的新序列。

训练与调试要点

使用自定义层构建模型后，编译和训练过程与标准模型无异。可以使用model.compile()指定优化器和损失函数，然后调用model.fit()进行训练。为了确保自定义层正确实现，建议使用小批量数据和简单任务进行测试，并检查输出的形状和值是否符合预期。TensorFlow的急切执行模式使得在调试时可以直接打印张量的值，这对于验证注意力权重的计算是否正确非常有帮助。