解析GPT模型generate方法输出的注意力机制结果

示例代码 #

# Importing necessary modules
from transformers import GPT2Tokenizer, GPT2LMHeadModel

# Loading pre-trained GPT-2 tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

# Encoding input text
input_ids = tokenizer.encode("The dog is running", return_tensors='pt')

# Generating model output with attention information
output = model.generate(
    input_ids,
    max_length=6,
    num_return_sequences=1,
    no_repeat_ngram_size=2,
    output_attentions=True,
    return_dict_in_generate=True,
)

# Extracting attention tensors
attn = output.attentions

观察总结 #

• attn变量是一个包含两个元素的元组，对应新增生成的token数量（因6-4=2）。
• 每个元素是由12个张量组成的元组，对应GPT每个block的层数。
• 第一个张量形状为[1, 12, 4, 4]，第二个张量形状为[1, 12, 1, 5]。
• 可视化时，形状为[1, 12, 4, 4]的张量代表掩码注意力。

疑问解答 #

1. 两种形状张量的含义与差异 #

先明确GPT-2注意力张量的维度定义：[batch_size, num_heads, target_seq_len, source_seq_len]，各维度对应：

• batch_size：批量大小，此处为1
• num_heads：注意力头数量，GPT-2基础版固定为12
• target_seq_len：当前需计算注意力的目标序列长度
• source_seq_len：可被关注的源序列长度

形状[1, 12, 4, 4]的张量

这是生成第一个新token时的注意力张量：

• 初始输入序列长度为4（对应"The dog is running"的4个token），生成第一个新token时，模型需要基于这4个输入token完成注意力计算，此时目标序列长度为4（每个输入token都要和所有输入token做注意力交互），源序列长度也为4。
• 由于GPT是自回归模型，生成阶段采用因果掩码，每个位置只能关注自身及之前的位置，最终呈现为4×4的掩码注意力矩阵。

形状[1, 12, 1, 5]的张量

这是生成第二个新token时的注意力张量：

• 此时序列总长度变为5（初始4个token+刚生成的1个新token），生成第二个新token时，目标序列长度为1（仅需为这个新生成的token计算注意力），源序列长度为5（可关注包括新token在内的所有历史序列）。
• 这里的1表示当前仅处理新生成的单个token的注意力计算，5是当前已有的全部序列长度。

核心差异

• 序列基础：前者基于初始4-token序列计算，后者基于更新后的5-token序列计算
• 计算目标：前者为所有输入token计算注意力（用于生成第一个新token的前馈计算），后者仅为新生成的单个token计算注意力
• 掩码范围：前者的4×4矩阵是严格下三角掩码（每个位置只能查看前面的token），后者的1×5矩阵中，新token可关注所有历史token（包括自身）

2. 张量来自解码过程的具体阶段 #

这些张量来自GPT-2自回归解码的每一步token生成的注意力计算阶段：

• 当调用generate方法并设置output_attentions=True时，模型会在每一步生成新token的过程中，记录下每个Transformer decoder block内多头注意力层的输出张量。
• 第一个张量对应第一步解码：基于初始输入序列，通过所有12个decoder block的注意力计算，生成第一个新token，同时输出每个block的注意力矩阵。
• 第二个张量对应第二步解码：将第一步生成的token加入输入序列，再次通过所有12个decoder block的注意力计算，生成第二个新token，同时记录这一步每个block的注意力矩阵。

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