做序列建模的项目最磨人的往往不是模型训练本身而是训练完或拿到现成模型之后下一步不知道怎么用。跑分类任务还好说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这三个概念。我用一个很生活化的方式来解释。假设你在一个巨大的图书馆里查资料。你现在手里有一张便利贴上面写着你想查的关键词这个便利贴就是QueryQ。图书馆里的每本书封面上都有标签这些标签就是KeyK。你要做的是拿着便利贴去比对每一本书的标签看哪本书跟你这张便利贴最匹配匹配结果就是一个分数。最后你真正翻开的那些书的内容就是ValueV。注意力计算过程就是对所有Key算一个相似度分数经过缩放和softmax归一化之后变成一组权重然后用这组权重去加权求和所有的Value。这个加权求和的结果就是当前token在融合了全局信息之后的表示。在DNABERT-2这种双向Transformer里每个token都会作为Query去关注序列里的所有token。所以注意力矩阵的第i行含义就是“第i个token在编码时对序列中每个token分别投注了百分之多少的关注”。行求和是1这就是为什么它叫“注意力权重”。2.2 DNABERT-2的模型配置与注意力张量尺寸DNABERT-2的公开实现基于BERT架构做适配典型的配置是6层Transformer编码器隐藏维度76812个注意力头。注意每个头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_codeTrue) 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_attentionsTrue返回值里就会多出一个attentions字段是一个元组每个元素对应一层的注意力张量。sequence ACGTACGTACGTACGTACGTACGTACGTACGTACGTACGT inputs tokenizer(sequence, return_tensorspt) with torch.no_grad(): outputs model(**inputs, output_attentionsTrue) 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, cmapviridis, vmin0.0, vmaxattn.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表示是真实token0表示是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注意力模式在细节上会有差别。如果要做严谨的结论建议用多个种子或多个公开预训练版本做一个集成稳定性验证。只用单个模型跑一次得到的结果严格来说只能作为生成假设的依据不能作为定论。提注意力权重这件事代码本身不难难的是理解你手上拿到的到底是一幅什么图。多花点时间在不同序列上画热力图、对比层与层的差异、观察哪些头稳定地关注某些模式这些功夫最后都会反馈到你的下游任务准确率上。