3D Masked Autoencoders是显微镜下体积和多模态细胞表示的鲁棒学习器

# 3D 掩码自编码器：显微镜体积与多模态细胞表征的鲁棒学习器
来源：https://arxiv.org/html/2606.23964
11机构：AI for Health 研究所 & Helmholtz AI，计算健康中心，Helmholtz Munich – 德国环境健康研究中心，德国纽赫贝格
22机构：慕尼黑路德维希-马克西米利安大学医院第三内科，德国慕尼黑
33机构：慕尼黑路德维希-马克西米利安大学物理系，德国慕尼黑
44机构：德国癌症联盟（DKTK），慕尼黑分站，德国慕尼黑
55机构：慕尼黑机器学习中心（MCML），德国慕尼黑

###### 摘要

荧光显微镜中的自监督学习通常依赖于二维投影，尽管细胞本质上是三维的。我们在体积显微数据上系统比较了二维和三维掩码自编码器（MAE-2D vs. MAE-3D）。在匹配架构和训练协议下，MAE-3D 在下游单细胞任务上始终优于基于二维最大投影和切片的方法。我们进一步将视觉表征与预训练的蛋白质语言模型（ESM2）对齐，并展示跨模态监督对体积模型带来的增益更大。通道交叉注意力和频域正则化对于利用三维空间上下文至关重要。在蛋白质-蛋白质相互作用任务上，MAE-3D 的 ROC-AUC 达到 0.865，优于先前方法最多 +0.025。在蛋白质定位任务上，我们最好的三维模型达到了最先进的 AUC_micro（0.952）和 F1_micro（0.742），分别比先前方法绝对提升 +0.003 和 +0.010。总体而言，这些结果证明了原生三维建模和多模态对齐在单细胞显微表征学习中的优势。

††脚注：通讯作者：
{amirhossein.kardoost, tingying.peng, carsten.marr}@helmholtz-munich.de

## 1 引言

细胞是构成组织和器官的基本单位。其功能与亚细胞结构和空间组织密切相关。理解细胞组织仍然是生物学中的核心挑战。尽管已有大量研究[12,14,6,8,1]，但解析亚细胞结构和蛋白质定位仍然复杂，尤其是在高维成像数据中。荧光显微镜[5,2,20]通过用荧光标记标记蛋白质和细胞器，实现了细胞内结构的可视化。诸如 JUMP[2]、OpenCell[5]、WTC-11[20] 和 Human Protein Atlas (HPA)[16] 等大规模资源提供了捕获丰富亚细胞组织的多通道成像数据。值得注意的是，OpenCell 和 WTC-11 包含体积 Z 堆栈。然而，许多表征学习方法（如 Subcell[8] 和 DINO4Cell[6]）对这些体积的二维投影进行操作，丢弃了深度分辨的结构信息。我们研究了体积建模在细胞表征学习中的作用。在 OpenCell[5] 上，我们系统比较了二维和三维掩码自编码器（MAE）[22,12] 模型。我们证明，保留完整的三维结构能够产生更具信息量的表征，并且在下游性能上始终优于基于二维最大投影甚至基于切片的方法。除了纯视觉建模，我们还通过预训练的蛋白质语言模型（PLM）如 ESM2[13] 整合蛋白质序列信息，探索多模态集成。通过将图像特征与蛋白质嵌入对齐，我们将生物学基础的结构先验注入表征空间。我们证明，序列级监督增强了表征质量，尤其是在与体积建模结合时。

我们的贡献有三方面：(1) 我们证明 3D MAE 模型在两个下游任务上均优于 2D 对应模型。(2) 我们展示了通道交叉注意力和频域（FFT）正则化进一步增强了体积表征学习。(3) 我们证实将蛋白质语言模型集成到视觉框架中能够提升表征质量和下游性能，突显了多模态对齐对细胞成像的益处。代码可在 https://github.com/marrlab/mae3d-opencell 获取。

## 2 相关工作

蛋白质的荧光成像，结合 DNA 或膜参考标记，使得能够进行蛋白质定位和功能的单细胞分析[5,16,11]。OpenCell[5] 包含 1,310 种内源标记的人源蛋白质和 29,922 个实验测量的蛋白质-蛋白质相互作用，以高分辨率三维 Z 堆栈形式获取。其蛋白质多样性和体积成像使其非常适合研究三维表征学习以及整合编码结构信息的序列级嵌入。WTC-11 数据集[20] 包含 25 种内源标记蛋白质（细胞结构）的三维荧光图像，带有 DNA 和膜参考通道，支持蛋白质定位和细胞周期阶段分类等任务。与 OpenCell 相比，WTC-11 的蛋白质多样性明显较低（25 种 vs. 1,310 种蛋白质），而是专注于单细胞层面的详细结构表征。Cytoself[11] 在二维最大投影图像上训练向量量化变分自编码器[18]，并以蛋白质身份为条件，证明学习到的表征根据亚细胞定位进行聚类。DINO4Cell[6] 将自监督 DINO 训练[3] 应用于 WTC-11 的二维最大投影图像，并在下游任务（如蛋白质定位和细胞周期预测）上评估学习到的表征。Subcell[8] 使用掩码重建和基于相似性的目标从 HPA[16] 图像中学习表征，并表明图像衍生特征与蛋白质序列和结构嵌入（例如 ESM2[13]）互补，这些嵌入通过第二阶段的多模态模型进行集成。现有的基于图像的表征学习方法主要对二维图像或体积数据的二维投影进行操作。虽然二维架构可以通过逐片处理[21] 扩展到三维，但我们发现原生三维表征学习在性能上始终优于最大投影和基于切片的方法。我们的模型建立在 SelfMedMAE[22] 的基础上，这是一个最初为医学成像开发的三维掩码自编码器，我们将其适配到多通道荧光显微镜，并进一步增强了通道交叉注意力和三维频域正则化。受 Subcell[8] 启发，我们通过与源于 ESM2[13] 的蛋白质序列嵌入对齐，进一步改善了表征学习。与 Subcell 在第二阶段学习联合图像-序列表示不同，我们保留纯粹的基于图像的编码器，并通过两种机制在训练过程中整合 ESM2：(i) 用序列令牌调节解码器进行掩码重建，(ii) 使用对称的 InfoNCE[4] 目标对齐图像和序列嵌入。

详见标题图 1：MAE-3D⋆ 模型通过应用对比性 L_CLIP 损失来整合蛋白质表征，该损失在来自掩码三维 OpenCell 体积的投影图像令牌与对应的 ESM2 嵌入[13] 之间计算。ESM2 令牌还被馈送给解码器以指导重建。

## 3 方法论

所提出的模型基于掩码自编码器（MAE）框架，包括二维和三维版本[9,22]。输入是一个 Z 堆栈体积 I ∈ R^{C Z X Y}，其中 C、Z、X 和 Y 分别表示通道数、深度、宽度和高度。对于二维变体（MAE-2D_base），体积通过沿 z 轴的最大强度投影进行压缩，得到 I_2D ∈ R^{C X Y}。I_2D 图像被分割成不重叠的 p_x × p_y 大小的块。通过卷积层进行块嵌入，然后加上二维正弦位置嵌入。随机掩码比例为 m 的块，编码器处理可见块。潜在表示与掩码令牌结合，传递到解码器进行重建。对于三维变体（MAE-3D_base），保留深度维度，完整的 I_3D 体积被分割成大小为 p_z × p_x × p_y 的不重叠块。块嵌入通过三维卷积层实现，然后加上三维正弦位置嵌入。掩码和编码遵循与二维设置相同的程序。模型通过掩码块上的均方误差（MSE）损失进行训练。在两个模型中，所有通道都被联合处理。

### 3.1 通道交叉注意力（CCA）

基础模型扩展为双流编码器-解码器，以显式建模通道间交互[1]。每个通道作为独立的令牌流处理，通道之间进行交叉注意力。在通道间应用共享的随机掩码模式，以防止从互补通道中未掩码位置进行简单重建。因此，块嵌入为每个通道独立进行。在每个通道内，对所有可见令牌应用标准的多头自注意力，以捕获全局空间上下文。此外，引入了逐位置交叉注意力[19]：通道 c_i 中空间索引 i 处的令牌仅查询通道 c_j (i ≠ j) 中相同空间索引处的令牌。由于每个查询位置恰好有一个键，softmax 操作退化为 1。因此，在注意力计算中，softmax 被 sigmoid 门控机制取代[19]。解码器镜像编码器架构。带有通道交叉注意力的模型记为 CCA。

### 3.2 FFT 损失

对完全重建的输出应用基于 FFT[10] 的频率损失。该损失在每个通道 c 上针对重建的二维（MAE-2D）或三维（MAE-3D）图像计算，以保留精细的亚细胞结构。虽然这种频域损失最初是在 [12] 中为二维重建引入的，但我们将其扩展到三维设置。频域损失定义为

L_FFT = 1/2 ∑_c L_1( log(1 + |F(Î^c)|), log(1 + |F(I^c)|) ) (1)

其中 I^c 是原始图像，Î^c 是通道 c 处的重建图像，F 表示具有正交归一化的 N 维离散傅里叶变换，|·| 是幅度谱，log(1 + ·) 压缩动态范围。L_1 距离优于 L_2 距离，以增强对频域异常值的鲁棒性，并且仅使用幅度谱。对于 MAE-2D，变换应用于两个空间轴；而对于 MAE-3D，则计算所有三个体积轴上的变换。FFT 损失与 MSE 损失结合，使用权重因子 w_FFT，该因子在预热期间设置为零，并在斜坡上升阶段线性增加到目标值，从而稳定早期训练，并在重建变得结构有意义时引入频率约束。

### 3.3 与 ESM2 的多模态对齐

为了建模蛋白质序列信息与细胞形态之间的对齐，通过对比性图像-序列对齐目标[4] 将蛋白质嵌入注入 MAE 解码器。编码器架构保持不变。预训练的 ESM2[13] 蛋白质语言模型在 MAE 训练期间保持冻结，并且仅在重建损失稳定之后才引入。ESM2 嵌入通过学习到的线性层投影到解码器维度，生成一个单一的蛋白质令牌。该令牌被插入到每个通道特定的解码器序列中，与编码器的可见令牌和掩码令牌一起，使用零位置嵌入（因为它不携带空间信息）（见图1）。解码后，在图像重建之前丢弃蛋白质令牌。这种设计允许掩码块通过解码器的自注意力机制关注蛋白质上下文。对于跨模态对齐，来自编码器的图像嵌入被投影到 ESM2 嵌入维度，并使用对称的 InfoNCE[4] 目标以及余弦相似度：

L_CLIP = 1/2 ( CE(τ E_I E_P^⊤, diag) + CE(τ E_P E_I^⊤, diag) ) (2)

其中 E_I, E_P 表示归一化的图像和蛋白质投影，τ 是可学习的温度参数，目标对应于批次内匹配的对角线对。交叉熵（CE）项强制图像和蛋白质模态之间的双向对齐。最终损失变为

L = L_MSE + w_FFT L_FFT + w_CLIP L_CLIP. (3)

蛋白质-图像对齐仅对可见令牌计算，使得对比目标具有挑战性并促进鲁棒的表征。集成模型