TensorFlow的tfa.layers.MultiHeadAttention是否属于自回归结构？

• 首先明确结论：tfa.layers.MultiHeadAttention本身不内置自回归对应的因果掩码机制，你需要自行实现自回归所需的掩码逻辑。
• 该层没有causal=true/false这类内置参数是设计使然：和Keras核心库自带的tf.keras.layers.MultiHeadAttention不同，TensorFlow Addons中的这版多头注意力层没有封装快捷的因果掩码开关，所有掩码逻辑完全交由调用方通过attention_mask参数传入实现。
• 自回归场景下的使用方法非常简单，你只需要自行生成下三角因果掩码，在调用层时作为参数传入即可，参考代码示例：

import tensorflow as tf
import tensorflow_addons as tfa

# 自定义参数：batch大小、序列长度
batch_size = 2
seq_len = 10

# 生成下三角因果掩码：仅允许当前位置及之前的token参与注意力计算
causal_mask = tf.linalg.band_part(tf.ones((seq_len, seq_len), dtype=tf.bool), -1, 0)
# 扩展匹配batch维度
causal_mask = tf.repeat(causal_mask[tf.newaxis, ...], batch_size, axis=0)

# 初始化注意力层
mha_layer = tfa.layers.MultiHeadAttention(num_heads=8, head_size=64)
# 传入输入张量的同时传入因果掩码，即可实现自回归注意力效果
attn_output = mha_layer(query=your_query_tensor, 
                        value=your_value_tensor, 
                        key=your_key_tensor, 
                        attention_mask=causal_mask)

• 额外注意事项：
  • 纯自回归生成任务（如GPT类结构）中，query、key、value序列长度一致，直接使用上述下三角掩码即可
  • 编码器-解码器结构（如Transformer翻译模型）中，仅需要给解码器的自注意力模块添加因果掩码，交叉注意力模块不需要添加该掩码

内容的提问来源于stack exchange，提问作者Sami