0. TL; DR

本文介绍了一种名为 VQDNA 的新型框架，旨在通过向量量化（Vector Quantization, VQ）技术革新基因组序列的标记化方法。VQDNA 利用可学习的向量子码本作为基因组词汇表，能够自适应地将基因组序列标记化为具有模式感知能力的嵌入向量。此外，文章还提出了层次残差量化（Hierarchical Residual Quantization, HRQ）方法，通过多层次的码本设计，以粗到细的方式丰富基因组词汇表。通过在 32 个基因组数据集上的广泛实验，VQDNA 展现了优越的性能和参数效率，并在 SARS-CoV-2 突变分析中揭示了其潜在的生物学意义。

1. 背景介绍

基因组学作为一门研究生物基因组（即生物体内完整的 DNA 指令集）的学科，为科学家们提供了深入了解生命分子机制的机会。近年来，随着深度学习技术的发展，基因组学领域取得了显著进展，例如基因表达预测、DNA 甲基化预测等。这些研究的核心在于如何有效地表示基因组序列，以便模型能够从中学习到有用的模式。

传统的基因组语言模型通常采用手工设计的标记化策略（如 k-mer 或字节对编码 BPE），但这些方法可能无法充分捕捉基因组数据中的复杂模式。因此，本文提出了 VQDNA 框架，通过学习一个可区分的基因组词汇表来替代手工标记化方法，从而更有效地表示基因组序列。

2. VQDNA 模型

VQDNA 的核心思想是将基因组标记化问题转化为基因组词汇表学习问题。通过使用向量量化（VQ）技术，VQDNA 能够将基因组序列映射到一个离散的嵌入空间中，从而生成具有模式感知能力的基因组嵌入向量。

2.1 向量量化基因组词汇学习

VQDNA 基于 VQ-VAE（向量量化 - 变分自编码器）框架，通过以下步骤实现基因组词汇学习：

编码器：给定输入基因组序列 $X \in \mathbb{R}^{L \times d}$，编码器 $E_\theta(\cdot)$ 将其映射到潜在空间 $Z = E_\theta(X) \in \mathbb{R}^{L \times D}$。
量化：通过一个有限的词汇表 $C = \{(k, e(k))\}_{k=1}^K$，其中每个整数 $k$ 对应一个可学习的嵌入向量 $e(k) \in \mathbb{R}^D$，潜在表示 $Z$ 可以通过以下公式进行量化：$M_i = \arg\min_{k \in [K]} |Z_i - e(k)|_2^2$，其中 $M \in [K]^L$ 表示量化后的整数索引。
解码器：解码器 $G_\phi(\cdot)$ 将量化后的嵌入 $\hat{Z} = e(M)$ 映射回基因组序列空间，生成重构序列 $\hat{X} = G_\phi(\hat{Z})$。

整个框架的优化目标是最小化 VQ-VAE 损失函数：

\[\begin{aligned} L_{\text{VQ}} &= L_{\text{CE}}(X , \hat{X}) + || \text{sg}[Z] - \hat{Z}||_2^2 + \beta || Z - \text{sg}[\hat{Z}]||_2^2 \\ &= L_{\text{rec}} + L_{\text{code}} + \beta L_{\text{commit}} \end{aligned}\]

其中$L_{\text{rec}}$ 是重构损失，用于优化编码器和解码器；$L_{\text{code}}$ 是码本损失，用于更新码本嵌入；$L_{\text{commit}}$ 是承诺损失，用于确保训练的稳定性。

2.2 层次残差量化（Hierarchical Residual Quantization, HRQ）

为了进一步提升 VQ 量化器的性能，本文提出了层次残差量化（HRQ）方法。HRQ 通过设计多层次的码本，以粗到细的方式逐步丰富基因组词汇表。具体来说，HRQ 在编码器的不同层次输出$H^{(n)}$上引入多个码本$C^{(n)}$，每个码本负责捕捉不同粒度的基因组模式。

给定编码器第 $n$ 层的输出 $H^{(n)}$，HRQ 的量化过程如下：

量化：通过第 $n$ 层的码本 $C^{(n)}$ 对 $H^{(n)}$ 进行量化，得到量化后的代码索引 $M^{(n)}$：

\[M^{(n)}_i = \arg\min_{k \in [2^n K]} \|H^{(n)}_i - e^{(n)}(k)\|_2^2\]

残差计算：通过以下公式计算下一层的输入 $H^{(n+1)}$：

\[H^{(n+1)}_i = \begin{cases} 2Z^{(n)}_i - e^{(n)}(M^{(n)}_i), & \text{if } n \geq 2 \\ e^{(1)}(M^{(1)}_i), & \text{if } n = 1 \end{cases}\]

通过这种方式，HRQ 能够在不同层次上逐步捕捉基因组序列中的粗粒度语义和细粒度细节，从而生成更具区分性的基因组嵌入向量。

量化后的嵌入定义为所有层次化嵌入的均值:

\[\hat{Z}_i = \frac{1}{N}\sum_{n=1}^{N} \hat{H}^{(n)}_i = \frac{1}{N}\sum_{n=1}^{N} e^{(n)}(M^{(n)}_i)\]

3. 实验分析

本文通过在 32 个基因组数据集上的广泛实验，验证了 VQDNA 的性能和泛化能力。这些数据集涵盖了多种基因组分析任务，包括表观遗传标记预测、转录因子结合位点预测、SARS-CoV-2 变异分类等。其中VQVAE选用ConvNeXt结构。

实验结果表明，VQDNA 在多个任务上均取得了优于现有基因组语言模型的性能。例如，在 SARS-CoV-2 变异分类任务中，VQDNA（HRQ）的 F1 分数达到了 74.32%，显著高于其他模型。

VQDNA（HRQ）在实现优异性能的同时，仅使用了 110M 参数，远少于其他大规模模型（如 NT-2500M-1000g 的 2537M 参数）。这表明 VQDNA 在参数效率方面具有显著优势，能够在更少的参数下实现更好的性能。

通过对 SARS-CoV-2 突变的分析，VQDNA 展现了其生物学意义。实验中，HRQ 词汇表能够将具有相似突变的病毒株聚集在一起，同时将不同突变谱系的病毒株区分开来。例如，Lambda 变异株虽然与 Delta 变异株有部分相似的属性，但属于不同的谱系，HRQ 词汇表能够准确地将它们区分开来，这表明 VQDNA 能够捕捉到基因组序列中的细粒度模式。

本文还对 VQDNA 的关键设置进行了消融研究。例如，通过调整码本大小和维度，发现码本大小为 512 时，VQDNA 在重构和区分能力之间取得了最佳平衡。此外，实验还表明，25% 的掩码比例在掩码语言建模（MLM）预训练阶段能够帮助 VQDNA 学习到更好的表示。