bioMoR：面向高效基因组学习的生物引导混合递归框架

# bioMoR：生物学引导的递归混合模型用于高效基因组学习
Source: https://arxiv.org/html/2608.06727
Koushik Howlader1, Tirtho Roy1, Md Tauhidul Islam2, Wei Le1
1美国爱荷华州立大学 2美国斯坦福大学

###### 摘要

用于高维组学分析的Transformer模型会处理数千个基因或通路，尽管其中只有一部分需要深度计算。递归混合（Mixture-of-Recursions, MoR）通过自适应的令牌选择或专家选择路由提高了效率。我们提出了bioMoR，据我们所知，这是第一个将MoR应用于基因水平和通路水平学习的框架。我们的贡献包括确定在MoR骨干中整合结构化生物学知识的三个位置：基于图的信息共享精炼令牌嵌入，结构性偏置引导自注意力关注生物学相关的令牌，以及图感知路由器利用邻域信息确定每个令牌的递归深度。这些技术的核心在于我们的洞见：关于令牌相互作用的额外知识可以有效帮助模型构建嵌入，并选择哪些令牌应当被更深入地学习。在覆盖多种组学数据类型并采用统一五折交叉验证协议进行评估的八个基准上，与最强的生物学无关MoR基线相比，bioMoR将平均宏F1提高了8.2个百分点，平衡准确率提高了7.1个百分点，同时参数减少了75%，FLOPs比非递归Transformer最多减少58%。所选出的标记基因或通路提供了生物学可解释性，而它们各自的令牌递归深度则揭示了计算是如何分配的。

参见图1：bioMoR概述。单细胞和多组学输入被表示为基因或通路令牌。固定的生物学知识被整合在三个位置：(1) 嵌入平滑精炼相关令牌，(2) 生物学引导的路由选择进一步计算的令牌，(3) 注意力偏置在权重共享的递归Transformer块中促进生物学相关令牌之间的交互。

## 1 引言

高通量测序现在可以测量每个样本中的数千个基因、通路或分子特征。Transformer对这些数据具有吸引力，因为注意力可以建模长距离分子依赖关系；最近的单细胞基础模型，包括scBERT (?)、scGPT (?)、scFoundation (?)和Geneformer (?)，在注释和表示学习方面表现出色。然而，标准Transformer对每个令牌都进行相同的深度计算，尽管基因组信号通常是稀疏的：一种细胞类型、肿瘤状态或转移表型通常由一小部分标记基因、调控程序或通路驱动。为了计算效率和准确性，一个有用的基因组模型不应在信息量较弱的特征上浪费计算，同时却对解释表型的分子程序精炼不足。因此，有效的建模应同时回答两个问题：应该保留哪些基因或通路，以及其中哪些令牌应获得更多的模型计算？

**指导性假设**。我们应根据组学数据的语义，以选择性和结构化的方式执行计算和学习。模型应保留具有生物学意义的标记基因/通路令牌，并将更多的递归计算分配给那些生物学邻域对任务最具预测性的令牌。

现有方法只解决了这一假设的一部分。高效注意力方法如Linformer (?)、Performer (?)和Nyströmformer (?)通过低秩或核近似降低序列成本。自适应计算方法如ACT (?)、混合专家 (?)、混合深度 (?)和混合递归 (?)动态路由令牌。然而，这些方法是生物学无关的，例如，它们不知道两个基因是否共表达、共享通路或参与同一生物学过程。相反，生物学感知资源和模型，如CellMarker (?)、PanglaoDB (?)、Reactome (?)、基因组相互作用图 (?)、Pathformer (?)和scBiGNN (?)，提供并考虑了组学数据的有用语义结构，但它们通常用于特征工程、架构约束或事后解释，而不是用于选择令牌以及决定哪些令牌应接受更深的计算。

我们提出**bioMoR**，一个生物学引导的混合递归（Mixture-of-Recursions）框架，用于基因组学习。据我们所知，这是第一个将生物学知识与混合递归架构相结合以驱动基因或通路令牌选择从而实现自适应学习的工作。

如图1所示，bioMoR首先将高维输入压缩为标记基因或Reactome通路令牌。然后它使用一个共享的递归Transformer块，并将每个令牌路由到自适应的递归步数。与生物学无关的MoR (?)的关键区别在于，我们在三个位置向模型中注入相关的生物学知识：它平滑令牌嵌入，将注意力偏向相关令牌，并引导MoR路由器决定令牌特定的递归深度。因此，生物学结构和语义成为自适应计算的控制信号，而不只是静态先验。

这引出了我们评估中的四个问题：bioMoR是否提高了预测性能（RQ1），哪个生物知识注入位置最重要（RQ2），它是否确实使计算更高效（RQ3），以及哪些路由、标记预算和模型宽度设置是最优的（RQ4）。

在八个多样化的基因组学习数据集上，bioMoR将平均宏F1提高了8.2个百分点，平衡准确率提高了7.1个百分点，同时与传统Transformer相比参数减少了75%，FLOPs最多减少58%。它还达到了0.92的线性探测准确率，并成功选择了生物学相关的通路进行更深的计算，这一点得到了先前文献的证实。

总之，我们的工作做出了以下贡献：

- •我们引入了bioMoR，这是第一个用于基因和通路水平组学学习的混合递归（MoR）框架。
- •我们整合了生物学知识，包括基因共表达和通路-通路关系，以利用MoR骨干改进表示学习、注意力学习和自适应路由。
- •我们进行了广泛的评估，表明我们的技术显著提高了F1和平衡准确率，并减少了模型参数以及FLOP成本，并且它们普遍适用于单细胞和多组学基准。

## 2 相关工作

#### 基因组表示学习。

Transformer模型如scBERT (?)、scGPT (?)和Geneformer (?)从大规模单细胞数据集中学习基因表示，并支持细胞类型注释等任务。这些模型是有效的，但通常对每个输入令牌应用相同的模型深度。

#### 组学模型中的生物学知识。

P-NET利用已知的生物学层级进行癌症预测 (?)，而Pathformer (?)和scBiGNN (?)利用通路或基因相互作用信息来指导基因组学习。这些方法展示了生物学知识的价值，但它们没有利用它来决定每个令牌应获得多少计算。

#### 自适应计算。

自适应计算时间 (?)、混合专家 (?)、混合深度 (?)和MoR (?)允许不同令牌获得不同数量的计算。然而，它们的路由决策是生物学无关的。bioMoR通过利用生物学知识来指导令牌表示、注意力和递归深度，将这两个研究方向连接起来。

## 3 方法

首先，在*基因组令牌化*和*混合递归*中，我们提供两个现有组件的背景 (?;?;?)，bioMoR在此基础上构建。其次，在*bioMoR*中，我们解释了我们整合生物学知识以引导MoR学习的新技术。整体工作流程见图1。

### 基因组令牌化

设x∈RN×C表示一个基因组样本，具有N个基因和C个分子通道。单细胞基因表达测量转录本丰度，且C=1 (?)。多组学输入可能还包含突变指标（记录DNA序列改变 (?）以及拷贝数变异（CNV，记录DNA片段的获得或丢失 (?)）。因此，我们将每个基因表示为一个向量xi∈RC，包括表达、突变和CNV通道。

#### 单细胞标记令牌。

一个单细胞包含数千个基因，但只有少数是有信息量的。因此，我们将N个基因转化为M≪N个标记令牌。这使得后续注意力变得便宜（O(M2d)），并让每个令牌可以被读取为一个具名的标记基因。首先，每个基因i被嵌入为一个身份向量加上其表达值，

t_i = e_i + W_v x_i, 其中 e_i∈Rd, W_v∈Rd×C, (1)

遵循单细胞基础模型的基因加值方案 (?;?)。然后我们保留M个可学习查询；查询m通过键k_i=W_k e_i软性地选择基因：

w_mi = exp(q_m^T k_i / (τ√d)) / Σ_{j=1}^N exp(q_m^T k_j / (τ√d)), (2)

并将它们求和为一个令牌 c_m = Σ_{i=1}^N w_mi t_i。在测试时，查询m坍缩到其 top 基因 g_m = argmax_i w_mi，因此每个令牌都命名一个标记基因。

#### 通路令牌。

对于多组学癌症队列，我们将基因分组为Reactome通路，这些通路代表了参与同一生物学过程的经过整理的基因集合 (?)。每个令牌是一个通路，通过对成员基因进行池化构建——对于密集表达和拷贝数通道使用均值池化，对于稀疏突变通道使用求和/负担池化。每个令牌因此是一个具名的通路。

### 混合递归

设T=[c_1,...,c_M]为标记/通路令牌矩阵。MoR应用一个共享Transformer块f_θ最多K次：

H^(0)=T, H^(t+1)=f_θ(H^(t)), t=0,...,K−1。 (3)

共享使得参数数量与K无关，而不像K个独立的Transformer块那样 (?;?;?)。一个路由器决定哪些令牌继续。专家选择在步骤t保留top-⌈ctM⌉个令牌；令牌选择让每个令牌自行选择深度并进行负载均衡 (?)。活跃步数

d_m = Σ_{t=1}^K 1{m 在步骤 t 被保留}， (4)

是令牌特定的递归深度。

### bioMoR

算法1：bioMoR前向传播。
输入：x, 生物学知识B, 令牌预算M, 最大深度K
T ← 从x生成标记基因或通路令牌
A ← 从B生成的行归一化固定知识矩阵
H ← (1−λ)T + λAT；
A ← {1,...,M}
对于 t=1,...,K 当 A≠∅ 时执行：
    H_A ← f_θ(H_A; A)  # 复用同一块
    r_A ← 生物学感知路由器得分
    A_next ← 被路由到下一步的令牌
    不在 A_next 中的令牌停止并保持当前状态
    A ← A_next
返回被保留到递归结束的令牌
y_hat ← 分类器(pool(H))
训练：最小化类别加权交叉熵加辅助损失
输出：y_hat, 标记/通路, 递归深度d_m

#### 生物学知识构建。

遵循Genomap (?)，我们构建一个基因共表达矩阵，总结成对基因关系。其条目是由基因表达矩阵的逆协方差计算出的偏相关，并限制在所选标记上；该矩阵构成单细胞生物学知识B。对于多组学数据，我们使用经过整理的Reactome通路关系，并将其视为无向：当通路i和j相关时B_ij=B_ji=1，否则两个条目均为零 (?)。

设D为B的行和的对角矩阵。我们使用固定的行归一化矩阵

A = D^{−1}B。 (5)

相同的生物学知识被用于嵌入平滑（位置1）、注意力偏置（位置2）和深度路由（位置3）。

#### 位置1：生物学引导的嵌入平滑。

在递归之前，我们使用生物学知识图A平滑令牌，使得生物学相关令牌的信息也被整合到令牌的嵌入中：

T̃ = (1−λ)T + λAT, (6)

其中λ是可学习的。这会在重复递归更新之前将每个标记/通路令牌向其生物学邻域去噪。

#### 位置2：注意力偏置。

在递归块内部，*相同的*生物学知识图A偏置自注意力：

Attn(Q,K,V) = softmax(QK^T/√d + λ_attn 1[A>0])V。 (7)

这鼓励生物学相关的令牌在精炼期间交换信息。

#### 位置3：生物学感知路由。

原始MoR路由器使用当前令牌表示 h_m^(t) 来决定令牌m是否应继续进行下一个递归步骤。bioMoR保留该原始得分，并基于令牌的生物学邻域添加一个可学习的修正：

r̃_m^(t) = (w_r^T h_m^(t)) / τ_r + σ(γ_t) φ_ψ([AH^(t), H^(t)−AH^(t)]_m)。 (8)

第一项是原始MoR路由得分。在第二项中，AH^(t)总结了来自生物学相关令牌的信息，而H^(t)−AH^(t)显示了令牌与其生物学邻域的差异。小网络φ_ψ利用这两个信号调整路由得分，σ(γ_t)控制在递归步骤t添加多少修正。由于φ_ψ被初始化为输出零，修正项在训练开始时为零，bioMoR初始表现得像生物学无关的MoR。在训练过程中，当生物学信息有助于降低任务损失时，修正项变为非零；否则它可以保持接近零。这种设计防止了有噪声的生物学知识在开始时影响路由，同时允许模型在有用时使用它。

#### 训练与效率。

主任务损失是类别加权交叉熵。在训练期间，我们最小化

L = L_CE + α L_suf + β L_div + γ L_cmp + η L_router, (9)

其中辅助项鼓励标记充分性（L_suf）、标记多样性（L_div）、标记压缩/稀疏性（L_cmp）和路由器正则化（L_router）。在递归主干中，参数在所有深度共享；因此总参数规模不随最大深度K增加，这与使用K个独立Transformer块的方法不同。实际的FLOP成本由路由决定：一个令牌如果早停，就不会参与后续递归步骤的注意力计算，因此节省了计算。

## 4 结论

我们介绍了bioMoR，第一个用于基因组学习的生物学引导混合递归框架，它使用固定的生物学关系图来平滑令牌嵌入、偏置注意力，并指导令牌特定的递归深度。在八个组学数据集上的实验表明，与生物学无关的MoR基线相比，bioMoR将平均宏F1提高了8.2个百分点，平衡准确率提高了7.1个百分点，同时参数减少了75%，相对于传统非递归Transformer最多减少了58%的FLOPs。所选的标记基因和递归深度提供了可解释性，并表明生物学知识可以在不牺牲准确性的情况下指导高效的自适应计算。