论文阅读笔记(四)——HyenaDNA: Long-Range Genomic Sequence Modeling at Single Nucleotide Resolution
摘要
基因组(DNA)序列编码了基因调控、蛋白质合成和许多其他细胞特性的大量信息。与自然语言模型类似,研究人员提出了基因组学的基础模型,从未标记的基因组数据中学习可概括的特征,然后可以对下游任务进行微调,如识别调控元件。由于注意力的二次缩放,以前基于transformer的基因组模型使用512到4k标记作为上下文(<人类基因组的0.001%),这大大限制了DNA中远程相互作用的建模。此外,这些方法依赖于标记器或固定的k-mers来聚集有意义的DNA单元,失去单核苷酸分辨率(即DNA“特征”),其中微妙的遗传变异可以通过单核苷酸多态性(snp)完全改变蛋白质功能。最近,基于隐式卷积的大型语言模型Hyena被证明在质量上匹配注意力,同时允许更长的上下文长度和更低的时间复杂度。利用鬣狗新的远程能力,我们提出了HyenaDNA,这是一个在人类参考基因组上进行预训练的基因组基础模型,其上下文长度在单核苷酸水平上高达100万个标记-比以前密集的基于注意力的模型增加了500倍。HyenaDNA按序列长度进行次二次扩展(训练速度比Transformer快160倍),使用单核苷酸标记,并且在每层都具有完整的全局上下文。我们探索了更长的上下文能够实现什么-包括基因组学中首次使用上下文学习来简单适应新任务,而无需更新预训练的模型权重。在长距离的物种分类任务中,HyenaDNA可以在不降采样的情况下将上下文长度增加到1M,从而有效地解决了这一挑战。在核苷酸转换器的微调基准测试中,HyenaDNA使用具有数量级更少参数和预训练数据的模型,在18个数据集中的12个数据集上达到了最先进的水平(SotA)在genomicbenchmark上,HyenaDNA在8个数据集中的7个数据集上平均超过SotA 10个精度点,在增强子识别上平均超过SotA 20个精度点。
资源和链接
Method
Preliminaries and Related Work
Transformers and Attention
支持许多最近的基础模型的是注意力机制。给定长度为l的序列x∈R L×D,一个(单头)缩放自关注层[2,51]是从R L×D到R L×D的映射,该映射执行如下操作:
A
(
x
)
=
σ
(
x
W
q
W
k
⊤
x
⊤
)
,
y
=
A
(
x
)
x
W
v
\mathrm{A}(x)=\sigma\left(x \mathrm{~W}_{q} \mathrm{~W}_{k}^{\top} x^{\top}\right), \quad y=\mathrm{A}(x) x \mathrm{~W}_{v}
A(x)=σ(x Wq Wk⊤x⊤),y=A(x)x Wv
其中D是嵌入维数,Wq, Wk, Wv∈R D×D是可学习的线性映射,σ表示逐行softmax(可选缩放)。注意计算每个标记的所有成对比较,并按顺序长度缩放为0 (l2)。这允许高分辨率的全局上下文,但限制了当前硬件上上下文的大小。
以前降低二次注意力成本的方法使用专门的方法来近似全密集注意力[18]。在稀疏注意中,元素只关注所有其他位置的一个子集。另外,线性注意力方法构造A(u)的近似,可以在次二次时间内评估。然而,这两类方法都以较低的时间复杂度(允许较长的序列)来换取表现力上的损失。
Long Context Strategies in Genomics
为了获得更长的上下文,基因组模型依赖于两种策略:
- tokenization
- dilation and downsampling。
标记化是具有双向转换架构(BERT)的掩码语言建模(MLM)的必要步骤,这是基因组学中的一种常见模型。这些标记器使用固定的k-mers(长度为k的短重叠序列)或基于频率的字节对编码(BPE),试图将DNA聚合成有意义的单位。因此,这些聚合技术创建了大量的新词汇表(与4个核苷酸的自然词汇表相比),这些词汇表的泛化性较差。第二种策略使用扩展卷积和下采样,这两种方法本质上都是对权重之间的元素进行平均或跳过。一个典型的例子是Enformer,它使用扩张和下采样来达到100k核苷酸的上下文长度来预测基因表达轨迹。在标记化、扩展和下采样中常见的是牺牲单核苷酸分辨率以获得更长的上下文。
Large Convolutional Models
长度为L的输入x与(可学习的)滤波器h之间的离散卷积由下式给出:
y
t
=
(
h
∗
x
)
t
=
∑
t
′
=
0
L
−
1
h
t
−
t
′
x
t
′
or equivalently
y
=
T
x
.
y_{t}=(h * x)_{t}=\sum_{t^{\prime}=0}^{L-1} h_{t-t^{\prime}} x_{t^{\prime}} \quad \text { or equivalently } \quad y=\mathrm{T} x .
yt=(h∗x)t=t′=0∑L−1ht−t′xt′ or equivalently y=Tx.
其中T∈R L×L为卷积对应的Toeplitz矩阵。从历史上看,卷积在深度学习和更广泛的信号处理中发挥了重要作用。最近,有研究表明,通过堆叠k个长卷积层,其中k通过函数γθ即k:= γθ(L)参数化,可以在涉及长序列的各种基准测试中获得最先进的性能,例如长距离竞技场(LRA)。文献中提出了不同的γθ:状态空间模型和通过神经场的隐式参数化。在语言方面,隐式卷积语言模型的h族,H3和Hyena使用长卷积和门通来匹配Transformer在O(Llog2 L)时间内的性能,明显低于基于注意力的模型的O(l2)。
HyenaDNA Long-Range Genomic Foundation Models
The HyenaDNA Model
HyenaDNA模型是一种仅解码器的序列到序列架构,由Hyena操作符[37]组成的一堆块定义,然后是前馈神经网络(见图1.3)。
给定输入KaTeX parse error: Undefined control sequence: \inR at position 2: x\̲i̲n̲R̲^{L}(L为序列长度),Hyena算子可定义为:
(
x
1
,
x
2
,
v
)
↦
H
(
x
1
,
x
2
)
v
H
(
x
1
,
x
2
)
=
D
x
2
⊤
h
D
x
1
\begin{aligned} \left(x_{1}, x_{2}, v\right) & \mapsto \mathrm{H}\left(x_{1}, x_{2}\right) v \\ \mathrm{H}\left(x_{1}, x_{2}\right) & =\mathrm{D}_{x_{2}} \top_{h} \mathrm{D}_{x_{1}} \end{aligned}
(x1,x2,v)H(x1,x2)↦H(x1,x2)v=Dx2⊤hDx1
其中x1, x2, v是输入的投影,
T
h
∈
R
L
×
L
T_{h} \in R^{L \times L}
Th∈RL×L是由一个可学习的长卷积滤波器构造的Toeplitz矩阵,作为神经网络的输出,
(
T
h
)
i
j
=
h
i
−
j
(T_h)_{ij}=h_{i-j}
(Th)ij=hi−j(Th)ij = hi−j。卷积滤波值本身是通过一个以时间(位置)索引和可选的位置编码ht = γθ(t)为输入的小神经网络γθ获得的,这使得算子能够处理很长的序列,而不需要参数数量线性增长。进一步,矩阵
D
x
1
,
D
x
2
∈
R
L
×
L
D_{x1},D_{x2} \in R^{L \times L}
Dx1,Dx2∈RL×L以对角线上的x1, x2构造,并计算为元素门控。通过对输入序列进行密集线性层和短卷积得到投影,如图3.1所示。
Proposition 3.1. Hyena算子可以在O(llog2l)时间内求值
Training Long Sequence Models
Tokenization: HyenaDNA序列长度的次二次代价允许该模型直接在单核苷酸水平处理超长序列,而不需要基于频率的聚合标记器。这可以实现短序列和长序列的细粒度分辨率,这对于检测单核苷酸多态性或突变以及模拟基因表达中的远程依赖关系至关重要。
我们使用自然DNA词汇,并将每个核苷酸作为标记。标记包括“A”、“G”、“C”、“T”和“N”(非特定核苷酸)以及用于填充、分隔和未知字符的特殊字符标记。令牌被映射到嵌入维度D。
Sequence length warm-up for ultralong sequences: 随着梯度方差的增大,直接对长序列进行训练会影响训练的稳定性[28]。最初在较短的序列上训练(随后是较长的序列)被用于训练小规模的变形金刚并减少训练时间,而使用序列长度热身来解决多达2k个令牌的稳定性问题。
对于超长序列(200k+),我们开发了一个新的热身计划,逐步增加序列长度,以提高稳定性和减少训练时间。
我们的序列长度调度从L1 = 64开始,然后在每个阶段将窗口加倍,同时保持全局批大小不变。通过这样做,每个连续阶段的迭代将包含更多令牌,确保调度器也可以作为批处理大小预热的一种形式。在图3.2中,我们观察到序列长度调度在序列长度大于450k时尤为重要,在此长度下,训练时间减少40%,最终准确率提高7.5%,这将在后面的4.4.3节中描述。
Downstream Adaptation
长上下文模型的可调提示传统上用于引导FM的输出,方法是在输入前添加额外的上下文。在此方法的基础上,引入了软可调提示,将可学习的令牌(作为权重)直接注入输入,作为模型微调的替代方案。
使用扩展的上下文长度(L),我们能够在预训练后探索适应fm的新范式。给定一个下游任务,提示符xp∈R T,对应的标签为yp,我们在嵌入步骤后加上N≤L−T维数为D的可训练参数θ:
x
←
concat
[
embed
(
x
p
)
,
θ
]
,
x
∈
R
L
×
(
T
+
N
)
x \leftarrow \operatorname{concat}\left[\operatorname{embed}\left(x_{p}\right), \theta\right], \quad x \in \mathbb{R}^{L \times(T+N)}
x←concat[embed(xp),θ],x∈RL×(T+N)
然后由模型处理得到的序列x, θ在包含输入序列标签yp的损失函数上进行优化。关键是,软提示需要利用提示和标签对的一小部分来优化θ。
在软提示过程中,HyenaDNA只优化输入序列中提示符的参数,同时保持所有其他模型参数不变。因此,软提示提供了一种灵活且计算效率高的方法来使基因组FMs适应新的下游任务。
结果