做序列建模的项目，最磨人的往往不是模型训练本身，而是训练完或拿到现成模型之后，下一步不知道怎么用。跑分类任务还好说，Dataloader一套、一个predict接出来就完事。但如果想做一点跟可解释性、跟序列结构分析相关的活儿，比如看看模型到底关注了基因组的哪些区域，注意力提取就是绕不开的一步。
DNABERT-2这类基于Transformer架构的DNA语言模型，本质上就是把一段碱基序列当成一篇“文本”。既然它是Transformer，那么每一层的注意力权重就记录了模型在解码时对序列中其他位置的“关注程度”。把这层信息挖出来，就等于拿到了模型的“读片报告”——哪个区段被高亮了、哪段序列彼此之间产生了互动，都能看得清清楚楚。
这篇文章我不打算讲那种教科书式的原理复述，而是直接把我在实际项目里提取DNABERT-2注意力的完整流程、每一步的代码、踩过的坑全部摆出来。内容分五块：整体思路、原理细节、端到端实操、排障实录、以及一些只有自己跑过才知道的经验心得。读完你至少能独立把一段DNA序列的注意力矩阵拉到本地、可视化成热力图，并且知道怎么把它用在结合位点预测这类下游任务上。
1. 项目定位与整体思路
1.1 这个任务到底在解决什么问题
DNABERT-2是一个预训练DNA语言模型，输入是一段核酸序列，输出通常是整条序列的向量表示，或者每个token位置的向量表示。这种模型在启动子预测、剪接位点识别、基因表达调控预测等任务上表现都不错。但多数使用场景只停留在“输入一段序列，输出一个标签或回归值”，模型内部怎么工作，几乎是黑箱。
注意力提取要解决的就是这个“黑箱”问题。Transformer每一层都有一个自注意力子层，它计算的是序列中任意两个token之间的关联权重。如果把每一层的这个权重矩阵拉出来，就能看到模型在处理这条序列时，不同位置的token之间建立了哪些依赖关系。
实际价值主要体现在三块：
- 可解释性分析：定位模型认为重要的序列区域，比如转录因子结合位点附近是否出现注意力热点。
- 弱监督的motif发现：不用额外标注，直接根据注意力权重找到模型关注的高频模式，再做序列比对和motif分析。
- 模型调试：如果模型在训练数据上学到了非预期的模式，比如只关注序列长度而不关注碱基内容，注意力热力图能把这种“偷懒”行为暴露出来。
所以这个任务核心其实就是三件事：把权重从模型里拿出来、把维度对齐到原始序列坐标、再按需求做聚合与分析。
1.2 为什么不直接看Embedding或预测头的输出
你可能会有疑问：想看模型关注哪里，用最后一层的hidden state不行吗？比如拿CLS向量做个梯度分析（Grad-CAM）之类的。
方向没问题，但注意力权重有它独特的优势。
第一，注意力权重不需要任何额外的梯度计算，一次前向传播就能拿到，成本极低。Grad-CAM类方法要回传梯度，而且对梯度饱和比较敏感，在离散token的序列模型上表现往往很一般。
第二，注意力矩阵天然是pairwise的。Embedding只能告诉你“哪个token重要”，注意力还能告诉你“哪两个token之间建立了强关联”。在DNA场景中，启动子和增强子之间往往存在远程互作，这种pairwise的关系是单点向量给不了的。
第三，逐层的注意力演化信息非常有价值。低层和高层关注的东西不一样，这个在NLP里已经被反复验证过，DNA序列上也有类似现象。
所以我把注意力直接作为分析对象，而不是拿它当中间产物做个热力图就完事。后面你自然会看到，多头的不同分工、层间的信号演化，这些都是注意力独有的信息维度。
1.3 技术方案选型：只依赖Transformers生态
提取注意力的整体流程如下：
- 加载预训练模型和配套的tokenizer。
- 把DNA序列token化，变成模型输入。
- 跑一次前向传播，开启
output_attentions选项，或者用hook把中间层输出截获。 - 拿到
layers × batch × heads × seq_len × seq_len五维张量。 - 去除特殊token和padding部分，对齐原始序列坐标。
- 按层、按头聚合，做可视化或统计。
这套流程完全基于开源深度学习的标准工具链：PyTorch负责模型运行与张量操作，Transformers生态负责加载预训练模型，再加一个matplotlib/seaborn做可视化。不需要自己从零实现任何注意力机制，底层代码都是现成的，我们只是把它暴露出来而已。
2. 注意力机制原理与关键细节
2.1 用“查资料”来类比Q/K/V
注意力机制很难绕开Q、K、V这三个概念。我用一个很生活化的方式来解释。
假设你在一个巨大的图书馆里查资料。你现在手里有一张便利贴，上面写着你想查的关键词，这个便利贴就是Query（Q）。图书馆里的每本书封面上都有标签，这些标签就是Key（K）。你要做的，是拿着便利贴去比对每一本书的标签，看哪本书跟你这张便利贴最匹配，匹配结果就是一个分数。最后你真正翻开的那些书的内容，就是Value（V）。
注意力计算过程就是对所有Key算一个相似度分数，经过缩放和softmax归一化之后，变成一组权重，然后用这组权重去加权求和所有的Value。这个加权求和的结果，就是当前token在融合了全局信息之后的表示。
在DNABERT-2这种双向Transformer里，每个token都会作为Query去关注序列里的所有token。所以注意力矩阵的第i行，含义就是“第i个token在编码时，对序列中每个token分别投注了百分之多少的关注”。行求和是1，这就是为什么它叫“注意力权重”。
2.2 DNABERT-2的模型配置与注意力张量尺寸
DNABERT-2的公开实现基于BERT架构做适配，典型的配置是6层Transformer编码器，隐藏维度768，12个注意力头。注意每个头query/key的维度是768除以12，也就是64。你不需要记这些数字，加载模型后打印config就能确认，但理解维度关系有助于后面排查问题。
假设我们只处理一条序列，batch size为1。token化之后序列长度是L（包含[CLS]和[SEP]）。那么注意力张量的shape是：
[6层, 1个batch, 12个头, L, L]每一层、每一个头，都有一张L×L的权重矩阵，第i行第j列表示位置i对位置j的注意力权重。
这里有个容易混淆的坑：Transformers生态在输入模型时也会传一个attention_mask参数，但那是用来屏蔽padding位置的二进制掩码，不是注意力权重。前者是输入，后者是模型内部计算出来的结果，千万别混为一谈。
2.3 层和头之间的语义分化
我自己跑了几个不同组织的序列之后，最大的感受是：注意力头之间不是一样的，它们有明显分工。
浅层（前1-2层）的注意力主要集中在相邻token之间。这个好理解，DNA序列存在很强的局部结构依赖，比如密码子的偏好性、局部GC含量，这些信息浅层就可以捕获。越往后，注意力的关注范围会逐渐变长，能够看到更远的互动关系。深层的一些头会出现跨长距离的强注意力，比如关注序列首尾，或者同时关注两个相隔很远的保守区段。
另外，经常会出现“纯特殊token头”，某些头几乎只把注意力放在[CLS]和[SEP]上。这种头基本不携带序列本身的语境信息，分析时建议适当过滤掉。
所以在做注意力分析时，千万不要把所有头平均一下就完事。不同头和不同层的信息差异本身就是分析对象，盲求平均值会把这些模式全部抹掉。
3. 实操：端到端提取注意力权重
3.1 环境准备与模型加载
先说环境。核心依赖只有三样，PyTorch、Transformers生态的库、以及数据科学常用的几个包。我的版本大致是：PyTorch 2.x、Transformers 4.3x以上、numpy、matplotlib。
如果你的显卡显存不够，CPU也能跑，就是慢一点。分析注意力时一般force CPU推理也不会太夸张，因为前向传播只跑一次，不像训练那样反复迭代。但如果要批量处理上千条序列，那还是建议上GPU。
模板代码如下：
import torch from transformers import AutoModel, AutoTokenizer model_name = "dnabert2_model_identifier" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModel.from_pretrained(model_name, trust_remote_code=True) model.eval()加载后先做一步自检，确认模型结构，特别是编码器层怎么命名。大部分BERT系模型的编码器层路径是model.encoder.layer，但有些自定义实现名字不同，打印一下不会有坏处。
print(model.config.num_hidden_layers) print(model.config.num_attention_heads) print(model.config.hidden_size)如果打印结果是6、12、768，那跟下面的代码就完全对得上。
3.2 方案一：用output_attentions参数直接取值
这是最省事的方案。Transformers生态的大部分模型都支持在forward时传入output_attentions=True，返回值里就会多出一个attentions字段，是一个元组，每个元素对应一层的注意力张量。
sequence = "ACGTACGTACGTACGTACGTACGTACGTACGTACGTACGT" inputs = tokenizer(sequence, return_tensors="pt") with torch.no_grad(): outputs = model(**inputs, output_attentions=True) attentions = outputs.attentions print(len(attentions)) print(attentions[0].shape)输出大概是：
6 torch.Size([1, 12, 42, 42])这个42就是token长度，包含了[CLS]和[SEP]以及序列本身的40个碱基token。
这个方案的好处是简洁，不碰模型内部结构。但它有一个限制：如果你加载的不是底层模型，而是下游任务模型（比如分类头包装过的模型），某些封装类可能没有暴露output_attentions选项，或者即使传了也拿不到，这时就得用方案二。
3.3 方案二：用forward hook插桩提取
如果拿不到返回值，或者说你想在模型前向传播的中间环节抓取某个特定层的注意力，用hook是最灵活的方案。
所谓的hook，就是在某个PyTorch模块的前向传播结束后，插入一个我们自己写的函数，这个函数能拿到模块的输入和输出。注意力子层的输出通常包含两个部分：第一个是注意力加权后的context向量，第二个就是注意力权重本身。
attention_cache = {} def make_hook(layer_idx): def hook_fn(module, input, output): # output大概率是tuple，第二个元素是注意力权重 if isinstance(output, tuple): attn = output[1] else: attn = output attention_cache[layer_idx] = attn.detach().cpu() return hook_fn for idx, layer in enumerate(model.encoder.layer): layer.attention.self.register_forward_hook(make_hook(idx)) with torch.no_grad(): model(**inputs)拿到的attention_cache就是一个以层索引为key的字典，每个value的shape与方案一完全一致。
这个方案最大的价值是灵活。你可以只挂你感兴趣的某几层，比如只提取第2层和第4层，减少计算和内存开销。也可以在hook里直接对权重做裁剪或归一化，不用把全部原始张量带回内存。
3.4 token坐标向碱基坐标的映射
DNABERT-2用的不是单碱基token化，而是基于BPE的token化，这点非常关键，也是最容易翻车的地方。一条300bp的序列，拆出来的token数可能只有100-200个，每个token大概对应几个碱基。所以注意力矩阵的一行不是对应一个碱基，而是对应一个token片段。
如果直接拿注意力热力图叠加在序列上解释，很容易错位。必须做一次token到碱基的坐标映射。
方法不复杂：逐token解码，统计这个token解码后包含多少个有效碱基字符，然后累加。
def token_to_base_positions(sequence, input_ids): base_positions = [] pos = 0 # 跳过CLS和SEP for tid in input_ids[1:-1]: piece = tokenizer.decode([tid]) length = len([c for c in piece if c in "ACGT"]) if length == 0: # 特殊字符，比如##等 continue base_positions.append((pos, pos + length - 1)) pos += length return base_positions有了这个映射，就能把某个token的注意力聚合到它覆盖的碱基区间上，画出来才跟真实序列对得上。
在跨层分析时，我会先把注意力按层做一次平均，再映射到碱基层面。但如果要保留多头的细节，那就先不聚合，等分析目标定下来再决定怎么压缩。
3.5 可视化与初步分析
拿到注意力矩阵之后，第一件事永远是画热力图。别急着统计，先用眼睛看。热力图能直观暴露很多问题，比如某一行几乎全亮，那这个token就是在“偷看全局”，不太正常。
import numpy as np import matplotlib.pyplot as plt layer_idx = 3 head_idx = 2 attn = attentions[layer_idx][0, head_idx].numpy() fig, ax = plt.subplots(figsize=(10, 8)) im = ax.imshow(attn, cmap="viridis", vmin=0.0, vmax=attn.max()) ax.set_xlabel("Key position") ax.set_ylabel("Query position") plt.colorbar(im) plt.show()画完之后注意观察几个点：对角线亮不亮、有没有明显离对角线的热点、有没有整行整列为零。对角线亮说明模型很依赖每个token自身携带的信息，这在BERT系模型里很常见，属于正常现象，但如果一个头对角线上下的边带非常宽，说明它更关注局部邻域而不是本身，那对应的生物学含义可能也不同。
3.6 下游应用：用注意力做结合位点弱监督发掘
我们举一个最常用的下游场景：转录因子结合位点预测。传统做法是标注大量ChIP-seq peak，训练分类模型。但如果手上只有几条阳性序列和阴性序列，怎么做初步的motif发现？
我的做法是：提取阳性序列的注意力，做跨序列平均，然后看哪些碱基区域的平均注意力权重显著高于背景。把这些区域对应的序列片段拿去跑一个motif比对，往往能直接命中已知的TF motif。
这个思路本质上是把注意力当成弱标注信息来用。模型不需要你告诉它结合位点在哪，它自己通过预训练就学到了一部分调控模式的表示，注意力热点区域可以作为候选位点。理论上做这一步的预期是：候选区域富集了真实motif。我自己的结果也基本验证了这一点，热点区域跟已知motif位置高度重合，但偶尔也会出现模型关注GC-rich区域而不是真实结合位点的情况，这个噪声来源需要正视。
4. 常见问题与排障实录
4.1 attention_mask和注意力权重，别搞混
这个坑我已经不知道见了多少人踩过。模型的输入参数attention_mask是一个0/1向量，1表示是真实token，0表示是padding。很多新手看到代码里有“attention”字样就去当成注意力权重取出来分析，拿到一堆0和1还以为模型出bug了。
记住：model(**inputs)里的attention_mask是输入的一部分，用来告诉模型哪些位置不用关注。真正要提取的注意力权重，是前向传播之后返回的attentions字段，或者是hook拿到的中间输出。这两个东西没有任何直接关系。
4.2 维度顺序对不上，代码可能崩在诡异的地方
不同版本、不同实现，返回的注意力张量维度顺序可能不一样。有些模型返回的是[batch, layers, heads, seq_len, seq_len]，有些是[layers, batch, heads, seq_len, seq_len]。如果你按固定索引去取，很容易越界或者取错层。
我的习惯是拿到张量后第一件事就print shape，然后手动确认层面和头面：
attn = outputs.attentions[0] # 第一层 assert len(attn.shape) == 4 # batch, heads, seq_len, seq_len永远不要凭记忆写死索引。每个模型的封装风格有点差异，多打印一次花不了两秒钟，能省掉后面大量的定位时间。
4.3 显存和存储溢出
注意力矩阵看着不大，但实际开销很可观。以6层、12头、每条序列512个token来算，单条序列的注意力权重总数是：
6 × 12 × 512 × 512 = 18,874,368个float
存成float32，大概72MB。如果batch size是32，一趟就是2.3GB。你还要算上原始激活值内存，GPU很容易就被吃满。
之前我批量提取几千条序列的注意力时，直接把所有注意力张量全部堆在内存里，结果程序跑到一半内存爆了。后来改成了边提取边释放，只保留聚合后的结果或top区域坐标。具体来说，如果你只需要热点区域，就别把完整矩阵囤在内存里，直接用一个滑窗统计每个位置的累计注意力贡献，然后只存这个累计向量。空间从72MB降到2KB，效果完全不变。
4.4 hook拿不到输出，或者拿到的为None
用hook方案时，如果输出是None，大概率是挂错了模块。BERT系传统结构的注意力输出在layer.attention.self下，但对于一些重构过的模型，注意力可能会放在layer.attention本身，或者改名成layer.self_attn之类的。
定位方式很简单，把模型结构完整打印出来，找一找哪个子模块的地址包含attention关键字就行。
另一种情况是，hook挂上了，但输出里没有tuple，只有一个张量。这时要看具体实现，有的优化版注意力模块只返回context向量，注意力权重被内部丢弃了。遇到这种情况，要么换用方案一直接让模型返回值带attentions，要么就得手动改模型的前向代码，在注意力计算处抽出权重。动手改代码之前，先确认一下模型源码，不要一上来就硬干。
4.5 训练模式下提取，梯度把显存吃穿
提取注意力一般在推理场景下做，模型要切到eval()模式，外部套torch.no_grad()。原因很简单：注意力权重虽然是前向传播的中间结果，但如果你开着梯度，PyTorch会把所有中间激活都记录下来用于反向传播。本来一条序列跑得飞快，开着梯度能慢好几倍，显存占用也会翻番。
如果你是边训练边分析注意力，也建议把分析操作放到评估阶段，用with torch.no_grad():包裹起来，并在提取之后立刻detach().cpu()。
5. 实操中的经验与心得
5.1 注意力模式与生物学信号的对齐
跑过几次真实序列分析后，我发现注意力信号确实不是均匀分布的。启动子区域附近，浅层注意力的相邻token依赖会更强；而到了深层，某些头会把注意力投到与启动子相隔几百bp的上游区域。这种远程关注在基因调控里其实是有生物学意义的，比如远距离增强子与启动子的互作。
但不要过于乐观地直接把注意力热点等价于功能元件。注意力本质上反映的是模型在语言层面学到的统计规律，它有没有被下游任务利用，取决于模型怎么组合这些头。所以注意力热点可以用来筛选候选区域，但不能作为唯一的决定依据。最好配合序列保守性或ChIP-seq已有的实验结果做交叉验证。
5.2 一个弱监督场景的真实操作
我这边接过一个任务，没有现成的结合位点标注，只给了几十条已知有调控功能的片段和几十条随机负对照。
我的处理流程是：
- 每条序列做BPE token化，跑前向拿到最后两层的注意力。
- 把注意力按列求和，得到每个token被关注的总强度。
- 用坐标映射回碱基位置，得到一条注意力波形曲线。
- 对阳性序列的波形做平均，找出显著高的峰。
- 取峰附近的一段序列与数据库motif做比对。
这个流程跑下来，找出来的候选motif与已知数据库结果重合率达到了可用的水平。最重要的经验是：不要只用某一层某一头的注意力，把最后两层的注意力先求平均再分析，稳定性会明显提升。单个头的方差很大，直接看很容易被噪声带偏。
5.3 一定要留个心眼的地方
最后说两个细节。
第一，BPE的token边界不一定落在碱基的自然单位边界上，一个token可能覆盖半个密码子，也可能跨过一个motif的边界。所以在做motif回收时，可以直接取注意力加权后的基序区间，不要生硬地按token边界切序列。
第二，不同随机种子初始化下的DNABERT-2，注意力模式在细节上会有差别。如果要做严谨的结论，建议用多个种子或多个公开预训练版本做一个集成稳定性验证。只用单个模型跑一次得到的结果，严格来说只能作为生成假设的依据，不能作为定论。
提注意力权重这件事，代码本身不难，难的是理解你手上拿到的到底是一幅什么图。多花点时间在不同序列上画热力图、对比层与层的差异、观察哪些头稳定地关注某些模式，这些功夫最后都会反馈到你的下游任务准确率上。
