论文阅读113：A label disambiguation-based multimodal massive multiple instance learning approach (AAAI)

1 要点

题目：基于标签消歧的多模态大规模多示例学习及免疫库分类

代码：https://github.com/Josie-xufan/LaDM3IL

方法：

1. 利用实例级多示例学习 (MIL) 框架来减少大规模MIL场景下的计算开销，并处理免疫库分类问题下的大存储容量挑战；
2. 模型包括：
   • 特征提取器： 使用门控注意力机制和张量融合的多模态融合模块，整合氨基酸序列和VDJ基因片段的信息；
   • 标签消歧模块，用于降低错误监督的影响；
   • 聚合模块：整合各个受体的预测及其相应频率来生成免疫库水平的预测；
3. 实验：
   • 数据集：巨细胞病毒 (Cytomegalovirus, CMV) 和癌症数据集；
   • 地址：https://clients.adaptivebiotech.com/pub/Emerson-2017-NatGen

背景：

注：该方法与免疫库分类这一背景很相关，可以重点关注其如何利用多模态MIL处理问题这一机制上

2 方法

2.1 问题定义

一个适应性免疫受体库 (adaptive immune receptor repertoires, AIRR) 包含大量的自适应免疫受体 (AIR)，也就是通常所说的包与实例。令 { I R 1 , I R 2 , … , I R N } \{ IR_1,IR_2,\dots,IR_N \} {
IR1​,IR2​,…,IRN​}表示$N$个包 (AIRR)，每个包包含$M$个实例$I{R}_i^1,I{R}_i^2,\dots ,I{R}_i^M$。注意不同包中的实例数量差异悬殊。每个包的标签记为 Y i ∈ { 0 , … , C } Y_i\in\{0,\dots,C\} Yi​∈{
0,…,C}，其中$C$为类别数。

与传统MIL不同之处在于，每个AIR对应一个频率值，对于包而言，记为 { f r e i 1 , f r e i 2 , … , f r e i M } \{fre_i^1,fre_i^2,\dots,fre_i^M\} {
frei1​,frei2​,…,freiM​}，表示对某种抗原的免疫反应的强度。

本文的目的是建立一个映射函数$Y_i=F(I{R}_i)$，即获得每个包$I{R}_i$的免疫状态。与传统的实例级方法类似，实例的标签被初始化为包的标签，不同之处在于，这些伪标签将通过标签消歧模块逐步更新。

2.2 模型架构

LaDM${}^3$IL的整体架构如图2，接下来将依次介绍特征提取、标签消歧、聚合模块，并在最后说明损失函数。

图2：LaDM${}^3$IL总体架构，其包含特征提取器、标签消歧模块，以及一个汇聚模块。a) 特征提取器：一个预训练SC-AIR-BERT作为序列编码器；一个可训练嵌入层作为基因编码器，以转换VDJ基因片段。门控注意力机制和一个张量融合模块用于整合学习的基因与氨基酸 (AA) 序列特征；b) 标签去混淆 (以二分类为例)：首先获取原型，其用于表示每个类的特征嵌入。在训练期间，top-K个实例被选择用于更新原型中相应类的嵌入，并根据实例特征嵌入与原型特征嵌入之间的相似性来调整每个实例的标签；以及c) 聚合模块：通过多层感知机获得每个实例的预测后，聚合模块通过将这些预测与其相应的频率相乘来整合这些预测，然后对结果进行归一化以生成包级别的预测。

2.2.1 特征提取

为了获取每个实例AIR的全面表示，需要整合来自AA序列和VDJ基因片段的信息：

1. 基因编码器：利用一个可训练嵌入层将VDJ基因片段转换为数值表示$h_g$，其包含两个独立的片段，分别编码了来自V基因片段和J基因片段的信息，维度则分别为16和8。注意D基因被排除，因为大部分实例种不包含；
2. 预训练序列编码SC-AIR-BERT：生成实例所对应的AA序列的表示$h_s$，其维度为512；
3. 门控注意力机制：输出两个模态的信息$o_g$和$o_s$：
   $$o_g=\sigma (z_g)\cdot h_g^{\prime }\text{\hspace{1em}(1)}$$其中$h_g^{\prime }=\text{ReLU}(W_g{h}_g+b_g)$是一个线性变换、$z_g=h_g^T{W}_{gs}{h}_s+b_{gs}$是一个双线性变换，以及$\sigma$是sigmoid函数。同理可以计算得到：
   $$o_s=\sigma (z_s)\cdot h_s^{\prime }\text{\hspace{1em}(2)}$$其中$h_s^{\prime }=\text{ReLU}(W_s{h}_s+b_s)$，$z_s=h_s^T{W}_{sg}{h}_g+b_{sg}$；
4. 张量融合模块：整合$o_g$和$o_s$为最终的表示$h$：
   $$h=\text{ReLU}(W_{\text{fusion}}\cdot (o_g\otimes o_s)+b_{\text{fusion}}),\text{\hspace{1em}(3)}$$其中$\otimes$表示Kronecker积。

2.2.2 标签消歧

标签消歧模块用于处理实例的不准确监督问题：

1. 每个实例通过特征提取器后，可以获得一个实例级预测：
   $$p_i^j=\text{softmax}(F{C}_{\text{receptor}}(h_i^j)),\text{\hspace{1em}(4)}$$其中$F{C}_{\text{receptor}}$表示分类器；
2. 在每一个训练轮次$e$，将从每个类别种选取实例级预测大于给定阈值$\theta$的实例：
   k e c − r e c e p t o r = { h i k , e , c ∣ p i k , e , c > θ , c ∈ { 0 , … , C } , k ∈ { 0 , K } , i ∈ { 0 , N } } . (5) \tag{5} \begin{aligned} kec-receptor=&\{ h_i^{k,e,c} | p_i^{k,e,c}>\theta,c\in\{0,\dots,C\},\\ &k\in\{0,K\},i\in\{0,N\}\}. \end{aligned} kec−receptor=​{
   hik,e,c​∣pik,e,c​>θ,c∈{
   0,…,C},k∈{
   0,K},i∈{
   0,N}}.​(5)
3. 基于动量法生成原型，即原型中每个类别$c$将在下一轮次更新：
   E prototype = Normalize ( λ ⋅ E prototype c , e + ( 1 − λ ) ⋅ h i k , e , c ) , h i k , e , c ∈ k e c − r e c e p t o r , c ∈ { 0 , C } , (6) \tag{6} \begin{aligned} E_\text{prototype}=&\text{Normalize}(\lambda\cdot E_\text{prototype}^{c,e}+(1-\lambda)\cdot h_i^{k,e,c}),\\ &h_i^{k,e,c}\in kec-receptor,c\in\{0,C\}, \end{aligned} Eprototype​=​Normalize(λ⋅Eprototypec,e​+(1−λ)⋅hik,e,c​),hik,e,c​∈kec−receptor,c∈{
   0,C},​(6)其中$\lambda \in [0,1]$是动量系数；
4. 每个实例的标签$Y_i^j$被初始化为包的标签$Y_i$，并通过实例与原型相似性的Onehot编码来更新：
   $$\begin{array}{rl}& Y_i^{j,e+1}=\gamma \cdot Y_i^{j,e}+(1-\gamma )\cdot \text{Onehot}(si{m}_i^{j,e}),\\ & \gamma =\frac{e(Epoc{h}_{end}-Epoc{h}_{start})}{Epoch}+Epoc{h}_{start},\end{array}\text{\hspace{1em}(7)}$$其中
   $$si{m}_i^{j,e}=\underset{c}{\mathrm{argmax}}(E_{\text{prototype}}^e\cdot (h_i^{j,e}{)}^T).\text{\hspace{1em}(8)}$$

2.2.3 聚合模块

为了生成每个免疫库，也就是包的预测$p_i$，我们结合实例预测$p_i^j$及其频率$fr{e}_i^j$：
$$p_i=\sum _{j=1}^M(p_i^j\cdot fr{e}_i^j).\text{\hspace{1em}(9)}$$最后通过最大最小标准化获取包预测。

2.2.4 损失函数

训练阶段分为热身阶段和标签消歧阶段：

1. 热身阶段：原型更新和标签消歧被暂停，只需计算实例预测和初始标签之间的交叉熵损失：
   $$L_{\text{receptor}}=-\frac{1}{NM}\sum _{i=1}^N\sum _{j=1}^M\sum _{c=0}^C{Y}_i^{j,c}\cdot \log (p_i^{j,c}).\text{\hspace{1em}(10)}$$
2. 标签消歧阶段：使用标签消歧损失：
   $$L_{\text{disambiguation}}=-\frac{1}{NM}\sum _{i=1}^N\sum _{j=1}^M\sum _{c=0}^C{Y}_i^{j,c,e}\cdot \log (p_i^{j,c,e}).\text{\hspace{1em}(11)}$$

3 实验

3.1 数据集

1. CMV：包含785个受体库，平均每个243960个受体。通过排除缺失信息，最终选择684个受体库，其中正312，负372。该数据集同时完成受体库分类和相关受体鉴定任务；
2. Cancer：训练集30000个肿瘤受体和40000个对照受体，测试集10000个肿瘤受体和19851个对照受体。该数据集重点关注肿瘤受体的鉴定；

3.2 实现细节

1. 对于CMV数据集，依据Widrich的方法来划分数据集；
2. 对于Cancer数据集，依据Beshnova的方法来划分数据集；
3. 对于自己的方法，CMV的训练集比例从10%逐渐增加到60%，Cancer则从20%逐渐增加到60%；
4. 优化器选用Adam，基因编码器的学习率设置为$1e^{-3}$、序列编码器的学习率设置为$1e^{-4}$，以及主模型的学习率设置为$1e^{-3}$；
5. 训练阶段基于网格搜索和验证集来确定超参数，最终如下：
   
6. 评价指标选用AUC和ACC；

3.3 性能对比

1. 肿瘤相关受体鉴定：
   
2. 序列长度对cancer数据集分类的影响：
   

3.4 训练集大小的影响

3.5 消融实验

3.6 包级标签监督下的受体鉴定