Llama2 70B注意力权重维度异常及非下三角矩阵问题咨询

我正在使用HuggingFace的Llama2 70B模型，通过Transformers库设置output_attentions=True来获取注意力权重。预期输出维度应为(层数, batch size, 注意力头数, 输入尺寸, 输入尺寸)，但实际得到的维度是(7, 80, 1, 64, 输入尺寸, 输入尺寸)。请问额外的维度7代表什么？另外，因果模型的注意力矩阵应为下三角矩阵，但输出并非如此，这是为什么？

使用代码 #

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "meta-llama/Llama-2-70b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(model_name)

input_ids = tokenizer([input_prompt], return_tensors="pt")
outputs= model.generate(
    input_ids['input_ids'],
    do_sample=False,
    temperature=None,
    top_p=None,
    top_k=None,
    max_new_tokens=7,
    sequence_bias=sequence_bias,
    return_dict_in_generate=True,
    output_attentions=True
)
attention = outputs.attentions

一、额外维度7的含义 #

这个7对应你设置的max_new_tokens=7，代表生成每个新token时的推理步数。

在generate方法中开启output_attentions=True，返回的注意力权重会保留生成每一个新token过程中，所有层的注意力数据。你预期的维度是单步推理（仅输入prompt时）的输出结构，那是直接调用model(**input_ids)得到的结果；而generate是多步迭代生成，所以会把每一步的注意力都存储下来，第一个维度自然就是生成新token的总步数。

二、注意力矩阵不是下三角的原因 #

这是因为Llama2采用旋转位置编码（RoPE），且Transformers库返回的是未应用因果掩码的原始注意力分数。

因果掩码是在计算注意力最终输出（即对value的加权求和）时才会被施加，而output_attentions=True返回的是掩码前的原始权重矩阵，所以看不到下三角结构。如果需要查看掩码后的结果，可以手动对注意力矩阵施加下三角掩码：

import torch

# 假设attention_step是某一步的注意力张量，维度为(层数, batch_size, 头数, seq_len, seq_len)
seq_len = attention_step.size(-1)
causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=attention_step.device))
masked_attention = attention_step * causal_mask

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