Causal Attention
Transformer的Decoder中最显著的结构是Casual Attention。
通过本篇文章，你将学会
Casual Attention的机制原理
Casual Attention在TensorFlow中的实现原理
如何快速地保存并打印TensorFlow中模型已经训练好的参数
如何实现Transformer的Decoder的前向传播
1、训练阶段
图1
Decoder和RNN一样，是串行地生成Token的。Casual Attention是带掩码的Attention——在标准的注意力分数矩阵上乘了一个下三角因果掩码矩阵，使得最终的输出具备了时序因果性——t时刻Token的预测过程不会受到t时刻之后的Key,Value以及Query的影响。因为Casual Attention这种独特掩码机制，使得Decoder在训练时可以进行并行训练，而非RNN那样只能串行训练。
2、推理阶段
图2
Decoder在推理时是依时序串行生成Token的，这是所有Casual Attention 的作用结果。
每个Casual Atention 在推理时和训练时是一样的，t时刻上的query只会与t时刻及t时刻之前的key计算点积，t时刻上的value只会与t时刻及t时刻之前的value进行加权相加。
其中有一个关键点值得注意。如果decoder在推理时采取和训练时一样的运算流程，那么在每个时间步上，t时刻之前的value,key,query都会被重复计算。
为了提高计算效率，就有了K-V缓存技术。
3、Key-Value缓存
Key-Value缓存是运用在Decoder推理时的技术，它需要缓存每个时间步上所有Causal Attention的Key和Value。如此就可以让Decoder在下一时间步无需再计算之前的Keys和Values了，只需要得到当前的Key和Value，再将当前的Query与所有Key计算注意力分数，接着合并所有Value就行了，合并后的Value再经过预测头，就可以生产新的Token了。
Key-Value缓存能够让Decoder中所有Causal Attention在每个一时间步只需计算上一时间步预测出来的Token的Key，Value和Query。如图2，t=2时刻，2之前的Keys和Values均已缓存，只需要计算Token的Key，Value和Query;t=3时刻，只需要计算Token
的Key，Value和Query。
每个时刻上的query是不需要缓存的，这是因为每个Casual Attention缓存了t时刻之前的value后，t时刻之前的query在t时刻就没用了，且t+1时刻的预测token只与t时刻的query以及t+1时刻之前所有value和key有关。query是有时效性的，无需缓存。
4、Casual Attention在Tensorflow中的实现原理
用Tensorflow训练一个Decoder并保存其模型结构和参数
import tensorflow as tf
# 词汇表
vocabulary_table = {1: "今", 2: "天", 3: "气", 4: "好", 5: "真"}
class Model(tf.keras.Model):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.embedding = tf.keras.layers.Embedding(input_dim=6, output_dim=64, name="Embedding")
self.casual_attention = tf.keras.layers.MultiHeadAttention(num_heads=2, key_dim=64, name="Casual_Attention")
self.dense = tf.keras.layers.Dense(6, activation="softmax", name="output_dense")
def call(self, inputs, training=None, mask=None):
x = self.embedding(inputs)
x = self.casual_attention(value=x,
key=x,
query=x,
use_causal_mask=True)
x = self.dense(x)
return x
class ExportModel(tf.Module):
def __init__(self, model):
self.model = model
@tf.function(input_signature=[tf.TensorSpec(shape=[1, 5], dtype=tf.float32)])
def __call__(self, inputs):
# training=False表明模型工作在推理模式
result = self.model(inputs, training=False)
return result
model = Model(name="Decoder")
# 今天天气真好
tokens = tf.constant([[1., 2., 2., 3., 5., 4.]])
Decoder_input = tokens[:, :-1]
print(Decoder_input)
Decoder_output = tokens[:, 1:]
print(Decoder_output)
model.compile(optimizer="adam", loss=tf.keras.losses.sparse_categorical_crossentropy)
model.fit(
x=Decoder_input,
y=Decoder_output,
epochs=164
)
# 只保存模型参数, 模型结构需要手动构建(前向传播)
model.save_weights("Decoder_weights.h5", save_format="h5")
# 模型结构与参数全部保存
model = ExportModel(model)
tf.saved_model.save(model, "Decoder")读取已训练好的权重参数，并手动实现Decoder的前向传播(这里的实现过程并没有使用K-V缓存)。
注意：在计算完query向量和key向量的点积后，一定要除以向量维数的平方根！
这是因为在Transformer模型中，注意力机制是核心组件。它通过query向量和key向量的点积来计算注意力分数。但当向量维度很高时，点积结果会变得非常大，这可能导致以下问题：
数值不稳定：大值在Softmax函数中会被放大，导致注意力过于集中，分布不均。
训练困难：梯度可能爆炸，影响模型稳定性。
为了解决这些问题，Transformer模型中引入了缩放注意力分数的技巧。具体来说，就是在点积后除以向量维度的平方根。这样做的好处有：
控制数值大小：将点积结果缩小到一个适中的范围，保持数值稳定。
优化Softmax表现：避免生成过于极端的概率分布，使注意力更平滑，模型学习更均匀。
训练更稳定：梯度不会因为指数函数放大而爆炸，收敛速度更快。
import tensorflow as tf
import h5py
import numpy as np
# tokens(Batch_size, num_tokens)
def forward(tokens):
# model.save_weights()仅保存模型权重, 保存格式为.h5文件, 需要用h5py库进行读取
# 读取模型参数文件
weights = h5py.File("Decoder_weights.h5")
# 打印Embedding层参数
# (6, 64), 词汇表大小为5, 再加上一个空白词(Padding 0), 总共6个Embedding.
embeddings_table = weights["Embedding/Decoder/Embedding/embeddings:0"][:]
# 打印Causal Attention层参数
# kernel : (input_dim=64, num_heads=2, key_dim=64)
# bias : (num_heads=2, key_dim=64)
key_dense_kernel = weights["Casual_Attention/Decoder/Casual_Attention/key/kernel:0"][:]
key_dense_bias = weights["Casual_Attention/Decoder/Casual_Attention/key/bias:0"][:]
# kernel : (input_dim=64, num_heads=2, query_dim=64)
# bias : (num_heads=2, query_dim=64)
query_dense_kernel = weights["Casual_Attention/Decoder/Casual_Attention/query/kernel:0"][:]
query_dense_bias = weights["Casual_Attention/Decoder/Casual_Attention/query/bias:0"][:]
# kernel : (input_dim=64, num_heads=2, value_dim=64)
# bias : (num_heads=2, value_dim=64)
value_dense_kernel = weights["Casual_Attention/Decoder/Casual_Attention/value/kernel:0"][:]
value_dense_bias = weights["Casual_Attention/Decoder/Casual_Attention/value/bias:0"][:]
# kernel : (num_heads=2, value_dim=64, output_dim=64)
# bias : (value_dim=64,)
attention_output_kernel = weights["Casual_Attention/Decoder/Casual_Attention/attention_output/kernel:0"][:]
attention_output_bias = weights["Casual_Attention/Decoder/Casual_Attention/attention_output/bias:0"][:]
# 打印最后一层密集层的参数
# (64, 6)
output_dense_kernel = weights["output_dense/Decoder/output_dense/kernel:0"][:]
# (6,)
output_dense_bias = weights["output_dense/Decoder/output_dense/bias:0"][:]
B = tf.shape(tokens)[0]
embeddings_table = tf.tile(embeddings_table[tf.newaxis], [B, 1, 1])
# (Batch_size, num_tokens, embeddings_dim)
embeddings = tf.gather_nd(embeddings_table, tokens[:, :, tf.newaxis], batch_dims=1)
# 将embeddings分别映射成key, query, value
key = tf.einsum("abc,cde->abde", embeddings, key_dense_kernel) + key_dense_bias
query = tf.einsum("abc,cde->abde", embeddings, query_dense_kernel) + query_dense_bias
value = tf.einsum("abc,cde->abde", embeddings, value_dense_kernel) + value_dense_bias
# 计算注意力分数,
# 计算query和key的数量积时, 一定要除以一个dim--query最后一个维度的长度
# 因为query和key的维度很高时, 它两的数量积往往很大, 几个很大的数经过SoftMax后会产生饱和状态
dim = tf.cast(tf.shape(query)[-1], dtype=tf.float32)
# 生成一个下三角的掩码矩阵, 下三角是0, 上三角是一个负无穷数(-1e11)
query_length = tf.shape(query)[1]
key_length = tf.shape(key)[1]
attention_mask = (1 - np.tri(N=query_length, M=key_length)) * -1e11
# query和key的数量积矩阵加上attention_mask后,
# 上三角全变为负无穷, 负无穷数的指数接近于0, 使得上三角的数量积在softmax中不起作用
scores = tf.math.softmax(tf.einsum("aecd, abcd -> acbe", key, query / tf.sqrt(dim)) + attention_mask, axis=-1)
# 利用注意力分数将value进行加权相加.
stacked_value = tf.einsum("acbe,aecd->abcd", scores, value)
# Causal Attention输出映射
attention_output = tf.einsum("abcd, cde -> abe", stacked_value, attention_output_kernel) + attention_output_bias
# 预测映射
prediction = tf.math.softmax(tf.einsum("abc, cd -> abd", attention_output, output_dense_kernel) + output_dense_bias, axis=-1)
return prediction
if __name__ == '__main__':
# [1., 2., 2., 3., 5., 4.]
tokens = tf.constant([[1, 2, 2, 3, 5]], dtype=tf.int32)
Decoder = tf.saved_model.load("Decoder")
print(forward(tokens)[:, -1])
print(Decoder(tf.cast(tokens, dtype=tf.float32))[:, -1])
最后检验上述的前向传播是否实现成功，
if __name__ == '__main__':
# [1., 2., 2., 3., 5., 4.]
tokens = tf.constant([[1, 2, 2, 3, 5]], dtype=tf.int32)
Decoder = tf.saved_model.load("Decoder")
print(forward(tokens)[:, -1])
print(Decoder(tf.cast(tokens, dtype=tf.float32))[:, -1])最终打印结果为：
tf.Tensor(
[[1.4004541e-09 2.9884575e-10 1.6915389e-29 9.1038288e-13 9.9929798e-01
7.0196530e-04]], shape=(1, 6), dtype=float32)
tf.Tensor(
[[1.4004541e-09 2.9884575e-10 1.6915389e-29 9.1038288e-13 9.9929798e-01
7.0196530e-04]], shape=(1, 6), dtype=float32)
两者结果一致，且在索引4处的概率值最大，表明下一个token的预测结果为4，符合真实值。
