如何解读BERT模型last layer最后attention head的注意力权重矩阵？

核心特性拆解 #

• 每行求和为1：注意力权重是经过Softmax归一化的分配系数，每行对应一个目标token（记为token_i），行内每个值代表token_i在生成自身语义表征时，分配给其他所有token（包括自身）的注意力占比。归一化到总和为1是为了让权重符合概率分布逻辑，确保注意力的分配是合理的比例关系。
• 矩阵非对称：这是完全正常的现象。token_i对token_j的关注度，和token_j对token_i的关注度没有强制相等的逻辑。以你使用的ProtBERT为例，蛋白质序列中某个氨基酸token可能重点关注相邻的疏水残基，但那个残基可能更关注下游的功能修饰位点，这种单向的功能关联就会导致矩阵不对称。

结合ProtBERT的具体解读建议 #

从你提供的热图来看（颜色越深权重越高）：

• 对角线区域颜色偏深：说明多数token会优先关注自身，这是Transformer模型的普遍特性——自身的基础信息是生成表征的核心。
• 局部非对角线深色块：对应蛋白质序列中空间相邻或功能协同的氨基酸对，ProtBERT预训练在海量蛋白质数据上，注意力头通常会学习到蛋白质结构里的残基接触关系、催化位点的协同作用这类生物特征。

提取代码修正说明 #

你代码里存在变量未定义的小问题，修正后的完整提取逻辑如下：

tokenizer = BertTokenizer.from_pretrained('Rostlab/prot_bert')
inputs = tokenizer(input_text, return_tensors='pt') 
attention_mask = inputs['attention_mask']
outputs = model(inputs['input_ids'], attention_mask=attention_mask)
# 获取所有层的注意力张量，outputs.attentions是元组，每个元素对应一层的注意力
all_attentions = outputs.attentions
# 取最后一层注意力：形状为(批量大小, 注意力头数, 序列长度, 序列长度)
last_layer_attention = all_attentions[-1]
# 提取最后一个注意力头的权重矩阵，转为numpy数组方便后续分析
last_head_attention = last_layer_attention[0, -1, :, :].detach().cpu().numpy()
# 后续可结合tokenizer的token列表，将矩阵行列与具体氨基酸token对应

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