跳转至

GeneZip: Region-Aware Compression for Long-Context DNA Modeling

会议: NeurIPS 2026
arXiv: 2602.17739
代码: https://github.com/AndyJZhao/GeneZip
领域: 模型压缩
关键词: 基因组长上下文建模, 动态分块路由, 区域感知压缩, DNA基础模型, 重复序列感知

一句话总结

针对长上下文 DNA 建模计算开销大且传统均匀降采样破坏基因功能特征的缺陷,GeneZip 提出基于基因结构注释的区域感知比率损失(RAR)与有界动态路由机制,训练时引导模型将 token 预算向高信息密度的启动子与功能区倾斜,在推理期脱离任何生物注释实现高达 137.6 BPT 的高保真压缩,并在 DNALongBench 多个长程预测任务上斩获第一。

研究背景与动机

长上下文 DNA 基础模型在从序列预测功能活性、变异效应评分以及三维基因组构象(如 Hi-C 接触图)等任务中扮演着核心角色。许多关键的生物调控相互作用与拓扑关联结构跨越数十万至数百万碱基对(100 kb 至 1 Mb 以上)。然而,直接对百万级长度的单碱基或固定 \(k\)-mer 序列进行表征与全局 token 混合(Token Mixing),计算复杂度与显存占用呈几何级数爆炸,使得极长上下文下的深度表征学习极难扩展。

为了规避计算壁垒,现有方法主要依赖两条技术路线:一是借助 HyenaDNA、Caduceus 或 Mamba2 等高效亚二次方长序列状态空间骨干,或者大幅削减模型容量;二是通过类似 U-Net 的多层卷积与池化操作实施固定步长的均匀降采样(Uniform Downsampling)。然而,这两种策略本质上都是按人工预设的神经网络架构硬性分配 token 分辨率,完全脱离了真实的基因组序列生物学特征。DNA 序列内在的信息密度极不均匀:仅占人类基因组约 1% 至 2% 的编码外显子区承受着极强的进化保守约束;启动子、5'/3' UTR、内含子与远端基因间区各自承载着差异巨大的调控信号;更为重要的是,哺乳动物基因组中充斥着海量的转座元件(TE)衍生分散重复序列与串联重复序列(TR),形成了巨大的序列级信息冗余。

如果采用均匀步长压缩,低信息密度的重复序列与长内含子占用了大量宝贵的 token 预算,而核心转录与调控元件的分辨率却被严重稀释。近期受 H-Net 启发的自适应 token 学习或启发式分词尝试打破固定词表,但普遍缺乏与生物功能区域先验相匹配的显式预算约束,极易在超长序列上发生路由塌缩或显存突刺。核心 idea:在层次化动态分块路由中引入区域感知比率(RAR)目标与有界路由约束,在预训练阶段利用静态基因结构注释引导路由网络为不同功能区学习差异化的碱基压缩率(BPT),而在推理阶段完全脱离任何注释直接对原始 DNA 执行兼具功能保真度与重复冗余自适应的极致非均匀压缩。

方法详解

整体框架

GeneZip 将长程 DNA 建模清晰解构为编码器压缩(Encoder)、核心序列交互混合(Token Mixing)与下游解码(Decoder)三个解耦阶段。输入为长度为 \(L\) 碱基对的原始未分词 DNA 序列 \(x_{1:L}\)。首先通过分层动态路由编码器将其非均匀压缩为长度为 \(T \ll L\) 的潜在 token 序列 \(z_{1:T}\),使有效压缩比 \(\text{BPT} = L/T\) 达到百倍以上;随后由深度 Transformer 或 Mamba2 序列骨干执行高阶长程上下文信息交互;最后由轻量级解码器恢复出碱基级特征或直接预测下游任务目标。

整个编码压缩器由 \(S\) 级层次化动态路由阶段(Stage)构成。在每一级中,网络利用序列表征在相邻位置之间的余弦不相似度自适应预测分块边界概率,并配合有界路由约束筛选出局部代表性 token 形成下采样序列。在预训练期,网络联合最小化单核苷酸下一词预测损失(NTP)与多级区域感知比率损失(RAR);完成训练后,GeneZip 仅凭自学习的序列上下文差异自主决定边界位置,推理时无需启动子、外显子等任何先验标记,即可实现端到端的高效快速压缩。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["原始长序列 DNA 输入<br/>长度 L: 12.8K ~ 1M bp"] --> B["层次化动态分块路由<br/>E(s) 提取特征与余弦相异度边界预测"]
    B --> C["有界路由约束投影<br/>K_min / K_max 阶段预算裁剪"]
    C --> D["区域感知比率目标监督<br/>结合静态基因结构锚定功能区 BPT 预算"]
    D --> E["压缩潜在 Token 序列<br/>长度 T = L / BPT, T << L"]
    E --> F["长程 Token Mixing 骨干网络<br/>多层 Mamba2 / Transformer 交互"]
    F --> G["解码器与任务输出<br/>单碱基 NTP 重建 / DNALongBench 下游预测"]

关键设计

1. 层次化动态分块路由:基于上下文表征差异的自适应边界诱导

为了克服固定步长下采样无法适应局部生物信息密度的缺陷,GeneZip 借鉴 H-Net 架构在多尺度表征上自适应划分子序列分块。在第 \(s\) 级路由阶段(序列长度为 \(L_s\)),编码网络提取隐藏状态 \(\hat{h}^{(s)}_{1:L_s}\),并通过线性映射分别计算相邻位置投影向量 \(q_t^{(s)}\) 与 \(k_{t-1}^{(s)}\)。两者的余弦不相似度直接决定了当前位置成为分块边界的概率: $\(p_t^{(s)} = \frac{1}{2}\left(1 - \frac{(q_t^{(s)})^\top k_{t-1}^{(s)}}{\|q_t^{(s)}\| \|k_{t-1}^{(s)}\|}\right), \quad t=2,\dots,L_s\)$ 根据定义设定序列起始点恒为边界 \(p_1^{(s)}=1\)。当连续碱基所处的局部语境发生突变(例如从高变异外显子过渡到高度保守的剪接位点或重复序列边界)时,相邻向量相似度下降,边界概率 \(p_t^{(s)}\) 自发激增。通过将概率以阈值离散化并提取边界位置的隐层向量代表整个 chunk,网络实现了将变长基因组片段聚合为单一 token 的动态降采样。在反向传播中,GeneZip 配合 Straight-Through 梯度估计器与平滑插值路径,实现端到端联合反向传播。

2. 区域感知比率目标(RAR):以基因结构为先验的差异化预算监督

单纯依赖无监督动态分块极易让模型将算力均匀分散在占据基因组绝大比例的无调控区间。为此,GeneZip 创新性地提出了区域感知比率(Region-Aware Ratio, RAR)损失函数,以静态生物学结构为软性约束,引导路由网络学习生物学信息先验。具体而言,针对基因组预定义的 7 类功能区域(启动子、CDS、UTR、外显子、内含子、近端基因间区 NIG、远端基因间区 DIG),给定基准目标压缩率 \(N\) 及各区域乘数 \(\mu_r\) 设定目标压缩率 \(N_r^\star = N \cdot \mu_r\)。在 \(S\) 级路由中几何分解为各级目标保留率 \(\tau_r^{\star(s)} = (N_r^\star)^{-1/S}\)。

在计算损失时,RAR 将投影前预期的保留概率 \(G_r^{(s)}\) 与离散投影后实际保留率 \(F_r^{(s)}\) 进行跨阶段区域对齐,其区域级损失函数形式为: $\(\mathcal{L}_{\mathrm{RAR}}^{(r,s)} = \frac{N_r^{\star(s)}}{N_r^{\star(s)}-1}\Big((N_r^{\star(s)}-1)\,F_r^{(s)}\,G_r^{(s)} + (1-F_r^{(s)})(1-G_r^{(s)})\Big)\)$ 该函数在 \(F_r^{(s)}=G_r^{(s)}=\tau_r^{\star(s)}\) 时取得全局最小值 1。一旦某一功能区域被过度裁剪或保留过多 token,梯度即会自适应调节该区域内各碱基位置的边界预测概率。通过按各区域碱基占比加权求和,GeneZip 成功将生物调控先验注入到了无需人工标签的纯序列推理模型中。

3. 有界路由机制:杜绝显存突刺与容量塌缩的全局预算投影

在对极长序列实施超高压缩率(如 BPT > 100)时,动态路由在优化初期极具不稳定性:第一级压缩阶段可能偶然在整段序列上激发出极密集的边界,导致局部保留 token 数量激增进而引发 GPU 显存溢出(OOM);而在深层第二级阶段,未加约束的路由器极易在初期发生崩溃退化,极端情况下甚至坍缩至全序列只选 1 个 token,导致表征容量彻底丧失与训练梯度爆炸。

GeneZip 引入了有界路由(Bounded Routing)投影算子 \(\Pi(\tilde{b}^{(s)}; p^{(s)})\)。设当前阶段理论预算为 \(K^{(s)} = \tau^{(s)} L_s\),算法在其周围设定下界 \(K_{\min}^{(s)} = \rho_{\min}^{(s)} K^{(s)}\) 与上界 \(K_{\max}^{(s)} = \rho_{\max}^{(s)} K^{(s)}\)(例如对所有模型统一固定底线保留 \(K_{\min}^{(s)} \ge 8\),防止容量归零;并在 636M 大模型训练时施加硬性截断以防 OOM)。当离散掩码中选中的边界数超出区间时,投影算子以预测置信度 \(p_t^{(s)}\) 为标尺,以最少翻转次数动态激活最高概率的静止位置或抑制最低概率的活跃位置,从而在保持梯度连续传递的同时为长上下文预训练构筑了坚固的数值稳定性护栏。

损失函数 / 训练策略

GeneZip 的整体预训练目标由碱基级自回归交叉熵损失 \(\mathcal{L}_{\mathrm{NTP}}\) 与区域比率正则化损失 \(\mathcal{L}_{\mathrm{RAR}}\) 构成: $\(\mathcal{L} = \mathcal{L}_{\mathrm{NTP}} + \alpha\,\mathcal{L}_{\mathrm{RAR}}\)$ 其中超参数 \(\alpha = 0.03\) 平衡了自回归重构精度与区域预算依从性。预训练分为两阶段:第一阶段在 12.8K 长度的 GENCODE 人类基因组切片上采用 AdamW 优化器(学习率 \(1\times 10^{-3}\),warmup 500 步)训练 100B 碱基;第二阶段基于第一阶段权重继续在 128K 超长切片上预训练 100B 碱基。整个预训练流程严格排除了 DNALongBench 评测区间以及人类 8、9、10 号染色体,从根本上防止下游评测的数据泄露。

实验关键数据

主实验:验证集自回归语言建模与区域 PPL

主实验系统对比了 GeneZip 与主流固定步长、非压缩及通用层次化分块基线在 12.8K 序列上下文下的语言建模困惑度(PPL),并报告了在各个基因功能区下的细粒度重构表现。

模型 BPT (碱基/Token) 全局 PPL ↓ 启动子 PPL ↓ CDS PPL ↓ UTR PPL ↓ 外显子 PPL ↓ 内含子 PPL ↓ 近端基因间 NIG ↓
细粒度基线/性能上限
Mamba2-84M 1.0 2.4135 2.0494 2.4136 2.5608 2.6346 2.5076 2.4826
H-Net-3BPT-70M 3.0 2.4115 2.0694 2.4815 2.4306 2.8113 2.5462 2.4357
均匀 / 层次化压缩基线
U-Net-70M 128.0 2.8414 3.1335 3.4665 3.2389 3.0586 2.8415 2.6474
H-Net-128BPT-70M 122.2 2.7930 2.8940 3.4657 3.0241 3.0277 2.8270 2.6375
区域感知压缩 (GeneZip)
GeneZip-70M 137.6 2.7259 2.6765 3.4124 2.9084 2.9847 2.7761 2.6278
GeneZip-70M-128K 169.2 2.6557 2.7520 3.0866 3.1236 3.0401 2.7350 2.5422
GeneZip-636M 137.0 2.6620 2.6094 2.8272 2.9311 3.2052 2.7724 2.5741
GeneZip-636M-128K 161.1 2.6484 2.8121 2.9647 3.1037 2.8712 2.6486 2.4941

消融实验:DNALongBench 全任务下游迁移与区域比率消融

为了检验不同生物先验分配(\(\mu_r\) 分配向量依次对应:启动子、CDS、UTR、外显子、内含子、近端基因间区 NIG、远端基因间区 DIG)对特定下游任务的有效性,在严控数据泄露的标准下进行了全方位消融评测。

配置 / 先验策略 区域权重乘数 \(\mu_r\) 生物学预算倾斜重点 eQTL (AUROC) ↑ ETGP (AUPRC) ↑ CMP (SCC/Corr) ↑ TISP (PCC) ↑
均匀 RAR 基线 (1, 1, 1, 1, 1, 1, 1) 无区域生物学先验偏好 0.8183 0.2297 0.1182 / 0.1607 0.2349
启动子-远端调控先验 (1, 16, 8, 8, 2, 2, 4) 保留启动子、内含子及近端基因间区 0.7967 0.4330 0.1324 / 0.1788 0.2785
顺式调控先验 (cis-reg) (1, 16, 2, 4, 2, 2, 4) 保留顺式非编码调控区与启动子/UTR 0.8200 0.3554 0.1282 / 0.1751 0.2764
基因间区先验 (intergenic) (32, 16, 8, 4, 2, 4, 1) 侧重非编码远端序列,粗化启动子 0.7863 0.4464 0.1280 / 0.1735 0.2238
转录平衡先验 (默认) (1, 1, 2, 2, 8, 8, 16) 侧重启动子与转录本编码区 0.7558 0.3334 0.1332 / 0.1791 0.2741

关键发现

  • 高压缩比下的精度保真:在 70M 参数同等规模下,GeneZip-70M 在取得最高压缩比(137.6 BPT)的同时,全局 PPL(2.7259)与全区域分项 PPL 全面超越同等压缩规模的 U-Net(2.8414)与 H-Net(2.7930)。当上下文扩展至 128K 且参数提升至 636M 时,PPL 进一步平稳优化至 2.6484。
  • 任务特定的生物学先验自洽性:消融实验清楚揭示不存在放之四海而皆准的单一压缩范式。以启动子/转录起始为核心的 TISP 任务在赋予启动子最高分辨率(\(\mu=1\))的配置下达到最佳 PCC(0.2785),而粗化启动子的配置骤降至 0.2238;涉及长程增强子与启动子联动的 ETGP 则在保留基因间区调控语境时 AUPRC 激增近一倍(从 0.2297 跃升至 0.4464);而染色质构象预测(CMP)整合了 1 Mb 范围内的复杂拓扑环化信号,更契合转录平衡的全面先验。
  • 天然的重复序列冗余感知:后验分析(RepeatMasker / TRF)表明,尽管训练时未曾输入任何转座子或串联重复标签,GeneZip 依然自发将更高的 BPT 施加于高度规律易预测的 SINE、LINE 及 TR 序列。人类 SINE 重复区局部 BPT 压缩比高达非重复区的 2.67 倍(H-Net 仅为 1.24 倍),证明该机制成功利用重复冗余换取了核心调控区的高分辨率。
  • 推理与微调效率飞跃:在 DNALongBench eQTL 任务上,GeneZip 仅耗时 50 分钟即完成完整微调,相比未压缩的双向基础大模型 JanusDNA(2520 分钟)实现了高达 50.4 倍的端到端墙钟加速,同时在跨组织平均 AUROC 上达成 0.820 的顶尖表现。

亮点与洞察

  • 训练有标、推理脱标的生物学先验内化:仅在训练期将人类基因组 annotation 作为正则化约束引入 RAR 损失,推理时直接基于纯未标注序列做无标签边界探测,既利用了演化生物学积累的先验知识,又保留了对未见物种、合成 DNA 与从头基因组的泛化能力。
  • 有界路由构筑极端压缩下的数值防线:在百倍压缩比的超深路由网络中,创新性设置上下限软硬截断投影,完美解决了动态分块在大模型极长序列训练中容易引发的 OOM 显存突刺与表征坍缩痛点。
  • 下游任务驱动的可调压缩接口:打破了“单一分词器固定分辨率”的思维定势,提供了一套可解释、可插拔的 \(\mu_r\) 调节接口,使长程生物模型能够按具体的下游生物机理定制 token 预算分配。

局限与展望

  • 预训练数据单一性:现有预训练模型主要在人类 GENCODE 语料库上训练,尽管测试表明其路由策略能够良好零样本迁移至小鼠染色体,但尚未在原生多物种泛基因组(Multi-species pan-genome)上开展大规模联合预训练。
  • 基于表征余弦差异的启发式假设:边界预测依然建立在相邻 token 隐层向量的余弦不相似度之上,在极度缺乏语义信息的低复杂度区域或长均聚物片段中,该信号可能夹杂表征噪声。
  • 静态先验向动态功能信号的拓展:目前的区域监督依赖静态结构标签,未来可进一步融入多模态功能调控信号(如 ATAC-seq 可及性、ChIP-seq 组蛋白修饰或 3D 染色质构象接触),实现更智能的数据驱动动态压缩。

相关工作与启发

  • vs H-Net (ICLR 2026):H-Net 开创了基于表征相异度的端到端层次化动态分块,但采用全局均匀比率损失,在基因组上无法感知启动子与重复序列的信息密度差异,且缺乏有界路由约束容易引发训练震荡;GeneZip 在其基础上增加了 RAR 目标与有界投影,实现了生物学功能可控与高效稳定的极致压缩。
  • vs JanusDNA & HyenaDNA:现有主流 DNA 基础模型要么使用近乎单碱基的极低压缩率强行利用亚二次方算子建模,导致长程任务微调开销巨大(JanusDNA 微调需 42 小时),要么采用机械的卷积跨步;GeneZip 作为即插即用的压缩前端,不仅将微调时间大幅缩短至 50 分钟(加速 50.4 倍),而且保留了核心功能区的高保真度。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将基因结构先验转化为动态路由压缩的可控先验,实现了无标签推理的区域自适应 DNA 压缩。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 128K 预训练、PPL 细粒度评测、RepeatMasker 冗余验证及 DNALongBench 4 项下游长程任务迁移与完整消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 动机严密,数学公式推导扎实,生物学实验分析与机制论证极其详实。
  • 价值: ⭐⭐⭐⭐⭐ 破解了长上下文 DNA 基础模型走向百万级规模的算力与分辨率分配核心瓶颈,极具实用工程与学术启发价值。