scFusionTTT：通过测试时训练层进行单细胞转录组与蛋白质组融合.

paper：scFusionTTT: Single-cell transcriptomics and proteomics fusion with Test-Time Training layers

0. TL; DR

scFusionTTT 是一个用于单细胞多组学融合的方法，将TTT（测试时训练）层应用于单细胞多组学分析的模型。模型的核心是一个基于TTT的掩码自编码器（masked autoencoder），采用了一个精心设计的三阶段训练策略：1) 多组学预训练；2) 多组学微调；3) 将从多组学数据中学到的知识迁移到单组学（仅scRNA-seq）数据的分析中，以增强其性能。

在四个多组学（CITE-seq）数据集和四个单组学（scRNA-seq）数据集上的实验表明，scFusionTTT在大多数评估指标上都取得了最佳性能，证明了其优越性。

1. 背景介绍

单细胞多组学技术，特别是CITE-seq，让我们能够同时检测同一个细胞的基因表达（transcriptomics）和表面蛋白（proteomics）。整合这两种信息，对于精准地识别细胞身份、理解疾病机制至关重要。

近年来，基于深度学习的方法，特别是利用注意力机制（attention mechanism）的方法，在多组学整合中展现了巨大的潜力。然而，当我们将这些方法直接应用于基因组数据时，会遇到困难。

计算复杂度：一个细胞的转录组通常包含数万个基因。对于自注意力（self-attention）机制来说，其计算复杂度是输入序列长度的二次方 ($O(N^2)$)。这意味着，当基因数量从一千增加到一万时，计算量会增加一百倍。这种指数级的增长使得传统的注意力模型在处理全基因组尺度的序列数据时，变得不切实际。
序列信息：基因在染色体上的排列并非随机，其线性顺序（sequential order）蕴含着重要的调控信息（例如邻近基因的共表达）。大多数现有的整合方法，如VAE或简单的MLP，都将基因视为一个无序的特征集合，完全忽略了这种宝贵的空间序列关系。

作者将目光投向了一种全新的序列建模范式：测试时训练（Test-Time Training, TTT）层。TTT层是一种新颖的RNN变体，它通过在推理时（test-time）利用自监督学习来动态更新其隐藏状态（hidden state），从而实现了线性计算复杂度 ($O(N)$)，并且在长上下文建模任务中表现出色。

作者认为，TTT层的这些特性，使其成为处理和建模具有内在序列关系的高维基因组数据的理想工具。基于这一洞见，作者开发了scFusionTTT，旨在利用TTT层的优势，来革新单细胞多组学数据的融合方式。

2. scFusionTTT 方法

scFusionTTT的核心是一个基于TTT层的掩码自编码器（masked autoencoder），它通过一个精巧的三阶段训练策略，实现了从多组学融合到单组学增强的完整流程。

2.1 TTT层与TTT块

TTT层的核心思想是其隐藏状态 $W_t$ 不再是像传统RNN那样通过固定的门控机制更新，而是在处理每个输入token（如一个基因）$x_t$ 时，通过一个自监督学习任务，进行梯度下降式的更新。

\[W_t = W_{t-1} - \eta \nabla \ell(W_{t-1}; x_t)\]

这里的自监督损失 $\ell$ 是一个重构损失，它鼓励模型用当前的隐藏状态，去预测输入token的某个变换。

\[\ell(W; x_t) = \| f(\theta_K x_t; W) - \theta_V x_t \|^2\]

通过这种方式，TTT层能够将序列信息动态地压缩并存储到其隐藏状态（即网络权重）中。

作者将TTT层与MLP和RMSNorm层组合，构建了一个类似于Transformer Block的TTT块。这是构成scFusionTTT所有编码器和解码器的基本单元。

\[X'_l = \text{TTTLayer}(\psi(X_{l-1})) + X_{l-1}\\ X_l = \text{MLP}(\psi(X'_l)) + X'_l\]

2.2 模型架构与三阶段训练

第一阶段：多组学预训练 (Multi-omics Pre-training)

在一个无监督的自监督任务中，让模型学习RNA和ADT（蛋白质）数据的基本表达模式和序列关系。

输入包含基因和蛋白质的顺序信息。作者首先根据基因在人类基因组中的物理位置进行排序，蛋白质则根据其编码基因的顺序排序。

对输入的RNA和ADT序列进行随机掩码（random masking）。使用两个独立的、基于TTT块的编码器，分别对RNA和ADT的未掩码部分进行编码，得到潜在表示 $E_{RNA}^{out}$ 和 $E_{ADT}^{out}$。

通过两个融合TTT模块（FusionTTT modules）进行信息交换。例如，为了更新RNA的表示，作者将ADT的表示与RNA的表示拼接后，输入到一个FusionTTT模块中，并将其输出与原始RNA表示进行残差连接。

\[FT_{RNA}^{in} = E_{RNA}^{out} + \lambda \cdot \text{FusionTTT}(E_{ADT}^{out}, E_{RNA}^{out})\]

最后，使用两个独立的解码器，从融合后的表示中重构出被掩码的原始RNA和ADT数据。最小化RNA和ADT的重构均方误差。

\[L_1 = \alpha L_{MSE}^{RNA} + \beta L_{MSE}^{ADT}\]

第二阶段：多组学微调 (Multi-omics Fine-tuning)

在预训练好的模型基础上，利用少量的带标签CITE-seq数据，进行有监督的微调，以学习区分不同细胞类型。

移除解码器，将FusionTTT模块的输出作为最终的细胞嵌入，并添加一个分类头。最小化交叉熵损失（cross-entropy loss）。

\[L_2 = \text{CrossEntropyLoss}(y, \hat{y})\]

第三阶段：知识迁移与单组学预测 (Unimodal Omics Predicting)

将从多组学数据中学到的知识，迁移到只需要单组学（如scRNA-seq）输入的任务中，以增强其性能。

固定预训练好的RNA编码器和FusionTTT模块的参数。当输入一个scRNA-seq样本时，模型会通过一个投影层（projection layer）来模拟出其缺失的ADT模态的潜在表示。然后，将真实的RNA表示和模拟的ADT表示一起送入FusionTTT模块，进行与第二阶段相同的分类预测。损失函数同样是交叉熵损失。

3. 实验分析

作者在四个公开的CITE-seq数据集和四个公开的scRNA-seq数据集上，对scFusionTTT的性能进行了全面的评估。

3.1 实验一：在CITE-seq多组学数据集上的聚类性能

作者将scFusionTTT与CiteFuse, TotalVI, SCOIT等六种SOTA的多组学整合方法进行了比较，使用了多达10种不同的聚类评估指标。

在所有四个CITE-seq数据集（SPL111, PBMC5K, PBMC10K, MALT10K）上，scFusionTTT在绝大多数评估指标上都取得了最佳性能。例如，在PBMC10K数据集上，scFusionTTT的ARI达到了0.90，而次优的CiteFuse仅为0.79。在MALT10K上，DBI（越低越好）达到了0.38，远低于其他方法。

UMAP可视化结果也直观地显示，scFusionTTT学习到的细胞嵌入，其簇结构最清晰，与真实标签的吻合度最高。

这些结果强有力地证明了scFusionTTT在整合CITE-seq数据方面的卓越性能。

3.2 实验二：在scRNA-seq单组学数据集上的聚类性能

作者进一步评估了scFusionTTT通过知识迁移来增强单组学分析的能力，并与Scanpy, Seurat, scVI等四种经典的scRNA-seq分析工具进行了比较。

在所有四个scRNA-seq数据集（IFNB, PBMC3K, CBMC, BMCITE）上，经过知识迁移的scFusionTTT，其聚类性能同样在绝大多数指标上都取得了最佳性能。例如，在CBMC数据集上，其ARI达到了0.76，而表现最好的基线方法scVI仅为0.53。

这证明了作者提出的三阶段训练策略是成功的：从多组学数据中学到的知识，确实可以被有效迁移，并显著提升单组学分析的准确性。

3.3 实验三：消融研究与参数分析

作者比较了TTT融合模块、注意力融合模块和简单的逐元素相加三种融合方式。结果清晰地显示，基于TTT的融合模块性能最佳，显著优于其他两种方式。这直接证明了TTT层在信息融合中的核心贡献。

作者测试了三种不同的基因/蛋白质输入顺序：正常顺序、反向顺序和随机打乱顺序。结果显示，随机打乱顺序后，模型的性能出现了显著的下降。这无可辩驳地证明了，基因和蛋白质的线性顺序信息对于模型学习是至关重要的，而scFusionTTT的TTT层架构能够有效地利用这种序列信息。

这些实验全面地证明，scFusionTTT通过其创新的TTT层架构和精巧的三阶段训练策略，无论是在多组学整合还是单组学增强任务中，都展现了SOTA的性能，为处理高维序列化的组学数据提供了全新的范式。