首页
/ Transformers 中的 BERTology 支持:读取 BERT/GPT 隐藏状态、注意力权重与注意力头的重要性分析

Transformers 中的 BERTology 支持:读取 BERT/GPT 隐藏状态、注意力权重与注意力头的重要性分析

2026-09-05 12:52:31作者:余洋婵Anita

BERTology(“BERT 学”)是研究大规模 Transformer 模型内部工作机制的新兴领域。HuggingFace Transformers 在 BERT/GPT/GPT-2 等模型中内置了访问内部表示的能力——包括全部隐藏状态、每个注意力头的全部注意力权重,以及用于计算头重要性得分并剪枝注意力头的机制——并在研究项目仓库中提供了配套的 bertology.py 示例脚本,用于在 GLUE 上微调的模型上完成信息提取与头剪枝。读完本文,你将掌握这些内置开关的用法、其源码实现位置与限制条件,以及如何复现注意力头重要性分析的基本流程。

什么是 BERTology

近年来,一个名为“BERTology”的研究方向兴起,专门探索 BERT 这类大规模 Transformer 模型的内部运作机制。原文档(docs/source/ar/bertology.md,另有中文版日文版)列举了该领域的几篇代表性论文:

  • BERT Rediscovers the Classical NLP Pipeline(Ian Tenney、Dipanjan Das、Ellie Pavlick,arXiv:1905.05950):证明 BERT 的不同层自发地重新“发现”了经典 NLP 流水线中的不同处理阶段;
  • Are Sixteen Heads Really Better than One?(Paul Michel、Omer Levy、Graham Neubig,arXiv:1905.10650):提出注意力头重要性得分(head importance score)并对冗余头进行剪枝,是 Transformers 相关功能的直接蓝本;
  • What Does BERT Look At? An Analysis of BERT's Attention(Kevin Clark、Urvashi Khandelwal、Omer Levy、Christopher D. Manning,arXiv:1906.04341):从注意力分布角度分析 BERT 到底在“看”什么;
  • CAT-probing(arXiv:2210.04633):面向编程语言预训练模型,度量其对代码结构的注意力。

为支撑这一领域,Transformers 在 BERT/GPT/GPT-2 模型中加入了三类附加能力,主要改编自 Paul Michel 等人的工作:

  1. 访问 BERT/GPT/GPT-2 的所有隐藏状态(每层各一个);
  2. 访问 BERT/GPT/GPT-2 每个注意力头的所有注意力权重
  3. 检索注意力头的输出值与梯度,从而计算头的重要性得分并对头做剪枝(pruning)。

核心机制之一:output_hidden_states 与 output_attentions 开关

在源码中,这两项能力由统一的配置开关控制。在 配置基类 中,output_attentions 等属性默认为 False

self._output_attentions: bool | None = kwargs.pop("output_attentions", False)

模型前向传播时,只要传入 output_hidden_states=Trueoutput_attentions=True(或在构造/加载 config 时设置),输出对象中就会附带相应字段。各输出数据类的文档约定一致(见 输出定义):

  • hidden_statestuple(torch.FloatTensor),可选):当传入 output_hidden_states=Trueconfig.output_hidden_states=True 时返回,包含嵌入层输出及各层隐藏状态;
  • attentionstuple(torch.FloatTensor),可选):当传入 output_attentions=Trueconfig.output_attentions=True 时返回,包含每一层中每个注意力头的注意力权重。

以 BERT 为例,Bert 建模代码 中通过 output 映射将 attentions 归属到 BertSelfAttention 组件、cross_attentions 归属到 BertCrossAttention 组件,并在各 forward 的返回路径上(如 L815 处)把 outputs.attentions 透传到最终模型输出中。

一个典型的调用方式是:

import torch
from transformers import BertForQuestionAnswering, BertTokenizerFast

tokenizer = BertTokenizerFast.from_pretrained("google-bert/bert-base-uncased")
model = BertForQuestionAnswering.from_pretrained("google-bert/bert-base-uncased")

inputs = tokenizer("Where do I find a good sandwich?", return_tensors="pt")
with torch.no_grad():
    outputs = model(
        **inputs,
        output_hidden_states=True,  # 获取全部 13 个 hidden states(嵌入层 + 12 层)
        output_attentions=True,     # 获取全部 12 层 x 12 头的注意力权重
    )

print(outputs.hidden_states)  # 13 个张量,每个形状为 (batch, seq, hidden_size)
print(outputs.attentions)     # 12 个张量,每个形状为 (batch, n_heads, seq, seq)

关键限制:注意力权重要求 eager 注意力实现

读取源码可以发现一条重要约束:output_attentions 与注意力实现方式(attn_implementation)强相关。configuration_utils.py 中的校验逻辑明确拒绝非 eager 实现下输出注意力权重:

def validate_output_attentions(self):
    if self.output_attentions and self._attn_implementation not in ["eager", None]:
        raise ValueError(
            "The `output_attentions` attribute is not supported when using the `attn_implementation` set to ..."
        )

对应的 property setter(L399-L415)在 eager 尚未被选定前设置该属性时也会给出提示。因此做 BERTology 类分析时,若模型以 flash/sdpa 注意力实现加载,需要改用 attn_implementation="eager" 重新加载,否则无法拿到逐头注意力权重。这一点在做剪枝实验前务必确认。

核心机制之二:注意力头的重要性得分与剪枝

第三项能力——“检索头的输出值与梯度以计算头重要性并剪枝”——对应 Michel 等人论文(arXiv:1905.10650)的做法:对某个头,用其输出与其梯度做外积(outer product)再取绝对值求和,得到重要性得分:

head_importance = sum_j | output_j · grad_output_j |

得分低的头被置零(即通过掩蔽使该头不再参与计算),在不明显损害下游任务性能的前提下减少计算量。原文档指出,上述提取与剪枝流程的完整实现位于一个专门的示例脚本 bertology.py 中,它针对在 GLUE 上微调过的模型演示了信息提取与头剪枝。

需要说明的是:该脚本已不在当前主仓库内。examples/research_projects/README.md 明确记载,此目录下的各研究项目(含 bertology/run_bertology.py)已迁移至独立仓库 transformers-research-projects。因此,当前仓库承担的是“提供底层能力”(隐藏状态、注意力权重、以及输出中携带的完整信息),而“如何算重要性得分、如何逐层剪头”的端到端示例应到该研究仓库中查阅。从源码结构看,主仓库中已不再保留 get_head_wrapper 等旧版专用访问器,头级操作(如按头置零)由研究脚本基于模型内部结构完成;而层级的信息提取则由上文的标准开关完成。

从 BERTology 到可复现实验:使用要点小结

  1. 只读能力开关output_hidden_states=Trueoutput_attentions=True 可在 config 或单次 forward 中设置,默认均为 False,不影响常规推理/训练行为;
  2. 注意力实现约束output_attentions 仅在 attn_implementationeager(或未显式指定)时可用(校验逻辑);
  3. 输出结构hidden_states 比层数多一个元素(含嵌入层输出),attentions 中每个张量的第一维之后是注意力头维度(输出文档约定);
  4. 剪枝实验:头重要性得分与剪枝的完整流程以研究脚本形式存在于独立研究仓库,主仓库提供的是被脚本所依赖的内部表示访问能力;
  5. 内存代价:同时开启两个开关会使模型额外保存/返回每层每头的中间张量,长序列时显存开销显著,建议在小批量、短序列场景下做分析性提取。

通过上述开关与配套研究脚本,Transformers 把 BERTology 从论文层面的分析变成了几行代码即可完成的可复现实验:读取任意层的隐藏状态、任意头的注意力分布,并在此基础上执行头级剪枝实验。

延伸阅读

  • BERT 建模源码attentions/cross_attentions 在各 forward 返回路径上的组装过程;
  • 通用输出数据类hidden_statesattentionscross_attentions 字段的存在条件与形状约定;
  • 配置基类output_attentions 的默认值、setter 提示与 validate_output_attentions 校验;
  • 本文档的多语言版本:中文日文西班牙文
登录后查看全文
热门项目推荐
相关项目推荐