说真的我第一次跑DNABERT-2的时候和大多数人一样只把它当成一个序列分类器序列进去embedding出来接个分类头就完事了。直到我做下游的调控元件预测要解释模型到底看到了什么时才意识到手里攥着一堆注意力矩阵其实是座没被认真挖过的富矿。这篇就把我这段时间折腾DNABERT-2注意力提取的经验整理一遍包括为什么要抠、怎么抠、抠出来怎么用以及一堆只有亲手跑过才会踩到的坑。1. 为什么要抠DNABERT-2的注意力一个被忽略的信息富矿1.1 DNABERT-2和经典BERT的差异决定了你怎么看注意力很多人直接把DNABERT-2当成DNA版的BERT来用这个直觉对但不全对。它在两个关键位置上动过刀一是tokenizer从k-mer改成了单碱基BPE二是位置编码从绝对位置嵌入换成了ALiBiAttention with Linear Biases。第二点对注意力提取影响尤其大。经典BERT把位置信息以position embedding的形式加到输入和隐藏状态上模型在计算self-attention时看到的query、key里已经混入了绝对位置。而DNABERT-2不同它是在attention logits上直接叠加一个随query和key之间相对距离线性变化的偏置然后再做softmax。换句话说你最后拿到的attention_probs已经天然耦合并入了相对位置先验。这意味着什么意味着你在解释attention map的时候看到某个位置集中了高注意力它可能是序列内容引起的也可能是距离引起的——更常见的是两者叠加。不要天真地以为注意力高就是这个碱基重要后面我会专门讲怎么处理这个问题。1.2 注意力矩阵里其实藏着预训练阶段学到的调控语法DNABERT-2在预训练阶段只干了很简单的一件事给定被mask的序列片段预测被遮住的碱基。但为了做好这件事模型必须学会上下文。你想想DNA不像自然语言没有固定的词序但转录因子结合位点、剪接信号、启动子模体这些东西本质上就是局部上下文规律。所以预训练完成之后模型不同层、不同注意力头会自发分化出不同的关注偏好有的头偏好盯AT-rich区域有的头负责建立跨长距离的接触有的头则对特定k-mer模式敏感。这些模式并不写在任何训练标注里但会体现在注意力权重上。这就是为什么值得去抠注意力它是目前我们能拿到的、最能反映模型内在行为的中间产物之一。无论你是做可解释性分析、想给下游分类任务加特征还是想做变异效应排序注意力矩阵都可以直接用。2. 准备工作加载117M参数模型时你先要搞定的四件事2.1 模型加载不是一句from_pretrained就完事DNABERT-2的权重和代码在模型库里是以自定义代码方式发布的标准加载要带trust_remote_codeTrue。我建议先把模型权重下载到本地再通过本地路径加载别每次跑都现场拉远程代码。原因很简单一是网络不稳定会让整个流程卡死在加载阶段二是远程代码版本可能悄悄变动你上周和这周拿到的模型行为可能不一样做注意力对比时会莫名踩坑。from transformers import AutoModel, AutoTokenizer MODEL_PATH /path/to/your/dnabert2-117m # 替换成你实际下载好的目录 tokenizer AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_codeTrue) model AutoModel.from_pretrained(MODEL_PATH, trust_remote_codeTrue) model.eval()提示如果条件允许用torch.compile或者至少torch.inference_mode()包住推理。117M参数不算大但你要扫全基因组片段时每一帧都是时间去积累出来的。2.2 输入预处理DNA序列不是普通文本DNABERT-2的分词器是基于单碱基BPE的词汇表很小但它对输入格式仍然有要求。我踩过的最蠢的坑直接把小写序列丢进去结果某些实现里小写字符会被当作未知字符产生一堆奇怪的token。seq ACGTACGTGATCGATCGATGCTAGCTAGCTAGCTAG seq seq.upper().replace( , ) # 只保留大写字母 inputs tokenizer(seq, return_tensorspt)这里还要注意DNA序列里如果有N未知碱基尽量先做平滑处理或按区间过滤。BPE分词器遇到连续N可能会合并成奇怪的token影响后面注意力对齐。我一般先把含N比例1%的窗口直接滤掉。2.3 序列长度和batch注意力是O(n²)的别贪心DNABERT-2的默认训练/推理截断长度一般是512。用ALiBi理论上可以外推到更长序列但attention_matrix是随长度平方膨胀的512的矩阵已经是512×512翻到1024直接变成四倍显存和四倍计算量。你在准备提取注意力之前先想清楚自己到底需不需要全长的attention map。我实际跑增强子这类1-2kb的调控元件时分成若干窗口比硬扩长更划算。窗口之间留一定overlap后续做碱基级聚合时方便平滑。2.4 前向输出里到底有哪些东西在直接开跑之前先把输出结构摸一遍。批量推理时加两个开关with torch.inference_mode(): outputs model(**inputs, output_attentionsTrue, output_hidden_statesTrue)然后看outputs.attentions是不是真的返回了tuple别想当然。DNABERT-2走的是远程自定义代码不同版本对output_attentions的支持不一致有的版本这个参数根本没传到底层返回的attentions可能是None。先打印再往下走。3. 提取注意力的三条路接口返回、hook注册与手动forward改造3.1 最小侵入方式直接用output_attentions如果前面的output_attentionsTrue生效那每条样本会得到一个tuple长度等于层数。每个元素对应一层的注意力权重形状一般是[batch_size, num_heads, seq_len, seq_len]。for layer_idx, attn in enumerate(outputs.attentions): print(flayer {layer_idx}: {attn.shape})通道是这个套路。attn_cache {} def make_hook(layer_name): def hook(module, args, output): # 不同模块输出结构不同这里是最常见的两种 candidate output[0] if isinstance(output, tuple) else output if candidate is not None and candidate.dim() 4: attn_cache[layer_name] candidate.detach().float() return hook for name, module in model.named_modules(): # 根据实际打印出来的name过滤这里有多种命名风格 if attention.self in name or self_attn in name: module.register_forward_hook(make_hook(name))hook里存tensor时要显式.detach()别直接存原tensor再带去backward。我在跑微调时犯过这个错后来排查发现hook把计算图挂住了显存暴涨。3.3 必要时手动改forward最后的选择output_attentions不可用、hook又因为远程代码模块结构太乱而抓不准时我见过有人直接切到源码里在attention_probs torch.softmax(attn_weights alibi_bias, dim-1)这一行后面加一句返回。我建议先跑两次hook方式。打印出来再过滤比盲改源码安全得多。最保险还是那句先跑一次把模型所有module的名字打出来看清楚你的版本里真正的self-attention模块叫什么再写过滤条件。这套流程只花两分钟能省两小时的debug。3.4 形状与token对齐别把位置坐标搞错拿到attention矩阵之后首先要确认坐标对应关系。在一段[CLS] DNA_token [SEP]的输入里矩阵的第一行第一列对应 [CLS] token最后一列对应 [SEP] token中间部分才是你序列的token。但注意一个关键问题DNABERT-2分词器是BPE中间token的数量不一定等于碱基数。如果你做的是motif级或碱基级分析要么把tokenizer的token映射回字符区间要么干脆在token粒度上做分析并接受一定的模糊性。我一般先做token粒度分析确定候选区间后再用滑动窗口精化到碱基级。直接把attention position当成碱基坐标是新手最容易犯的错误之一。4. 注意力矩阵里到底藏着什么motif、远程依赖和头分工4.1 先做一个粗糙的head扫描我建议拿到注意力矩阵后先别急着算各种花哨统计数据。第一步是直接可视化几个head的attention map或者做一个快速的CLS关注度扫描attn attn_cache[encoder.layer.11.attention.self] # 以某一层为例 # 假设形状 [batch, heads, q_len, k_len] # 取第0个样本[CLS]所在的行对所有head做平均 cls_attention attn[0, :, 0, :].mean(dim0) # 每个token一维我去年跑了一段启动子序列用最后几层的某个head做这种扫描发现模型对TATA box附近的AT-rich区域有非常强的集中注意力而对随机散布的重复序列则比较均匀。这种现象不是偶然的模型在做MLM预训练时要预测被mask的碱基就必须学会关注决定转录起始的关键区域。4.2 头分工用entropy筛掉高熵噪声头每个头都会有自己的性格。怎么量化我一般先算每个head的注意力分布熵对每一层每一个head把 [CLS] 行或者全矩阵做softmax归一化后计算熵。熵低说明注意力集中在少数token上——这是有强偏好的头熵高则说明注意力大范围分散基本是背景噪声头。实操中我通常会丢掉熵最高的20%-30%的头只保留低熵、注意力集中的头来解读。这么做不是因为我主观喜欢锐利的头而是因为高熵头在后续聚合时会提供近乎均匀的常数项稀释真正有信息量的信号。4.3 注意力重要性这个等式必须打问号做注意力分析绕不开一个争议attention weight到底能不能当importance用。我的立场是可以当作线索但不能当成证据。原因在于attention weight是从query和key算出来的它描述的是在softmax打分的分布里模型把多少权重分给某个位置。但真正的预测强度还要看value向量到底携带了什么信息。一个位置可能注意力不高但value里有一个大模长的信号对最终输出影响反而很大。所以我的习惯是注意力用于生成假设哪段区域可能与调控相关然后切换到你自己的任务里去验证比如看这个区域的突变对预测分数的影响或者配合gradient做归因。两条证据交叉验证了再下结论。5. 从注意力到下游任务多头聚合、注意力池化与变异评分5.1 多头聚合成位点评分平均还是按熵加权做完了head筛查接下来就是把多个头的注意力合成为一个位点级别的attention score。最简单的办法是对所有头取平均# 假设attn形状 [layer, heads, seq_len, seq_len] cls_attn attn[:, :, 0, :].mean(dim(0, 1)) # 所有层所有头平均但我实测下来平均不如按熵加权稳。具体做法对每个头算熵转成权重w_h 1 / entropy_h再对CLS关注度做加权平均。这样做相当于自动给专注的头更大的话语权抑制分散的头的噪声。下面这个函数可以直接拿去做变异评分的第一步def get_attention_profile(seq, model, tokenizer): inputs tokenizer(seq, return_tensorspt) with torch.inference_mode(): outputs model(**inputs, output_attentionsTrue) profiles [] for attn in outputs.attentions: # attn: [batch, heads, q_len, k_len] # 取CLS行做mean - [seq_len] profiles.append(attn[0, :, 0, :].mean(dim0)) # 在层维度上再平均 return torch.stack(profiles).mean(dim0)5.2 注意力池化别只盯着[CLS]DNABERT-2和BERT一样有[CLS] token很多人做序列分类时直接取[CLS]向量。但如果你想拿注意力做特征增强另一个思路是把注意力权重当作软对齐权重把所有位置的hidden state做一个加权求和得到的向量再输入分类头。这就是所谓的attention pooling。我在做增强子/启动子二分类的小样本实验时注意力池化比直接拼[CLS]稳定很多。原因倒不复杂DNA调控元件的关键信号往往落在特定区域注意力池化相当于让模型自己学会该重点聚合哪一段。实现上取最后一层hidden state把每个位置的权重定成该位置token对[CLS]的注意力或者做个层间平均然后加权求和。微调时这个权重可以固定也可以当作可学习参数一起训练。5.3 变异效应评分用一个零样本文本距离来排序变异把注意力用于变异排序是我个人觉得最实用的一条路。逻辑很简单一个碱基突变如果改变了模型的注意力分布说明这个位置在模型眼中动了格局反之如果注意力分布几乎不变说明这个位置的变异很可能不影响模型对该序列的理解。具体计算方式跑一遍参考序列得到所有层的注意力矩阵再跑一遍突变序列得到另一个注意力矩阵然后计算两个矩阵之间的某种距离比如L2范数。ref_attn get_attention_batch(ref_seq, model, tokenizer) mut_attn get_attention_batch(mut_seq, model, tokenizer) delta torch.norm(ref_attn - mut_attn)这个delta就可以当作一个无监督的变异影响分数用在全基因组扫描或者作为候选变异排序的初级筛。这个方法的好处是完全不需要标签跑一次前向就能打分我用它做已知致病变异回溯top-20候选里的召回率比我单纯用保守性打分的时候要好。但要注意如果两个序列之间只有单个碱基差异delta往往很小容易被批次噪声淹没。所以我一律用窗口平滑取以突变位点为中心上下游各5-10个位置的平均再看差异。6. 一些只有亲手跑过才知道的坑6.1 方向性DNA双链不是闹着玩的DNA有正负链同一个位点在基因组上从两条链读出来是反向互补的。如果把正链和负链的序列都丢进同一个模型模型对它们的处理是不同的——注意力矩阵里的位置含义直接跟着链方向变化。做变异评分或者motif分析时必须固定一个链方向。我的做法是统一用转录方向或者统一用基因组参考链正链然后对整个数据集保持一致。否则你会发现同一个变异换个方向跑出来的delta分可以差一倍。6.2 长序列的内存和精度fp16下注意力会变暴躁我刚开始为了提高速度用fp16跑全基因组扫描。结果发现fp16下softmax的数值精度变差注意力分布会出现大量接近0或1的极端值可视化图像变成一片噪点。如果要做注意力解释或变异评分我强烈建议用fp32推理。如果显存实在不够先降batch size而不是降精度。牺牲一点速度换来的是可解释的分布。6.3 自注意力对角线的假亮任何transformer的attention map主对角线附近永远是高亮区——因为当前位置对当前位置的注意力天然很高这让热图看起来像一条大裂缝。做可视化前先把对角线附近几个位置的权重抹掉或者做一次背景校正比如把每行的注意力都减去行均值否则你会误以为模型只在看自己。具体做法在热图矩阵中把|i-j| 1的权重置0然后重新按行归一化。这一步能瞬间让真正的长距离交互浮现出来。6.4 别想着存全矩阵有一次我扫了4000条调控元件每个元件512token每层注意力512×51212层×12个头……你算一下显存和磁盘一个样本就是12×12×512×512×4字节大约是150MB4000条就是600GB。这还是不压缩的情况。所以如果要规模化扫描千万别把完整矩阵落地。建议在跑前向的同时直接计算你需要的统计量比如CLS关注度向量、每头的entropy、变异前后的delta分把那些30KB的统计结果存下来就够了。真到了要可视化某个具体样本的阶段再单独跑一次提取那一条样本的完整矩阵。6.5 记录模型和分词器版本远程代码最大的坑是版本漂移。你跑完一批数据的注意力三个月后想复现时发现同一个样本出来的attention map不一样了不用惊讶——可能只是模型仓库里代码更新了或者你换了新的权重文件。我在做任何批次前都会把三个东西写进meta文件model.name_or_path、tokenizer.__class__、当天跑脚本时的git commit号。没有这三个东西之后的任何attention对比分析在方法学上都站不住。最后再分享一个小习惯我处理完注意力矩阵从来不直接看原始权重而是先把每个head的分布做一次标准化再转成log-domain去比较。标准化后的注意力对不同序列长度、不同GC含量没那么敏感跨样本可比性会好很多。你要是刚开始做DNABERT-2的注意力分析建议也把这个当作默认操作。