【算法详解】基于图注意力神经网络的空间转录组数据整合模型STAligner

目录

0 参考文献

1 背景介绍

2 STAligner算法解析

2.1 数据预处理

2.2 空间邻接图

2.3 空间感知点三元组学习（Spatially aware spot triplet learning）

2.3.1 图注意力自编码器

2.3.2 点三元组学习

2.4 多切片对齐与3D重构

---

0 参考文献

STAligner原论文：Integrating spatial transcriptomics data across different conditions, technologies, and developmental stages

1 背景介绍

随着空间转录组学（Spatial Transcriptomics, ST）技术的进步，对于组织切片的检测逐渐具有更高的空间分辨率，更大的视野域以及更多的点数（根据技术而定，每个用于防止切片的捕获区都由实际起到检测作用的点阵构成，例如10x Visium）。

目前大多数对于转录组学数据的分析局限于对单细胞RNA序列分析或者在二维ST切片内的分析，这忽略的组织切片的空间结构信息对于探索实际细胞分布和生理活动机理的作用。通过同时考虑细胞基因表达水平和空间位置信息，构建多个ST切片之间的关联，有益于增强下游分析任务的效果。因此对多个ST切片数据进行整合和比较分析至关重要。

然而来自不同个体、不同技术、不同生理状态的切片往往区别较大，并且受到批次效应（batch effect）的影响。即使是来自同一个个体同一个器官的切片，由于受到实验时的操作影响（如切片放置位置、放置方向等），也可能存在两切片在垂直切片方向上不能直接完全对齐的情况，这使得联合分析多个ST切片面临困难。因此需要使用合适的方法实现ST切片之间的对齐（alignment）和ST数据的整合（integration）。

STAligner是一个基于图注意力神经网络（graph attention neural network），用于整合和对齐ST数据集的方法。

之前已经介绍过几个具有类似目的的算法如PASTE、PASTE2、GPSA等。

2 STAligner算法解析

2.1 数据预处理

STAligner使用多ST切片的基因表达矩阵和空间坐标矩阵作为输入。其中基因表达矩阵是每一个检测点中不同基因的表达数量，每一行代表一个点，每一列代表一个基因。空间坐标矩阵通常会使用相对距离矩阵，若有n的点，则为n*n矩阵，第i行第j列元素为第i个点到第j个点的距离。

预处理时，先将基因表达矩阵进行行归一化，即每个点对应基因表达向量的元素和为1。随后进行对数变换。

对于每一个ST切片，选取最能够区分不同点之间基因表达情况区别的前5000个基因，再选择所有切片的5000个基因的交集作为所选基因，组成基因表达矩阵。通常以方差大小作为衡量标准。

2.2 空间邻接图

空间邻接图是一个无向邻接图，以一个对称矩阵表示，矩阵元素为0和1。构建邻接图时，基于空间坐标矩阵，以欧氏距离为标准，设置阈值r，若两点之间的欧式距离小于r，则认为两点之间存在一条无向边。将矩阵对应位置设置为1。

通过调整阈值r，使得在不同ST场景和平台的数据中，每个点平均具有4~15个邻居。

2.3 空间感知点三元组学习（Spatially aware spot triplet learning）

这里我是直接将Spatially aware spot triplet learning翻译为了空间感知点三元组学习，担心翻译得不太恰当，因此同时提供英文。空间感知意味着模型学习了空间位置信息，点三元组用于指导目标函数的设置。

STAligner由两部分组成：用于空间感知嵌入学习的图注意力自编码器（graph attention autoencoder）和用于批次校正（batch correction）的点三元组学习。

2.3.1 图注意力自编码器

图注意力自编码器是一个使用图结构信息和节点属性的无监督学习框架。它采用自注意力机制自适应地学习节点表示。与NLP领域的带掩码的自注意力机制类似，只有该节点的邻居节点才会对所学的节点表示产生影响。

STAligner的图注意力自编码器使用预处理后的基因表达矩阵X和对应的空间邻接图作为输入，进而获得节点表示z。以L表示编码器层数，N表示节点数量，则第k层对第i个结点的编码器输出如下

其中代表第i个节点的邻居集合，是第k层编码器的可训练参数矩阵，,是为了使得节点与邻居的相关程度具有可比性而设置的参数，通过下面公式计算

其中代表第i个点和其邻居j的相关程度，其中是第k层的可训练参数

第L层并不使用注意力机制，输出直接如下计算得到

解码器数量与编码器相同，用于重建基因表达特征。因此解码器与编码器采用相反的过程，使用编码器的输出作为解码器的输入，对于第k层，，使用下面方法重建点i在第k-1层的表示

与编码器类似，解码器最后一层计算方式如下，其输出作为重建的基因表达特征。

为了防止过拟合，令。

优化目标是最大化输入特征和输出的重建特征的相似程度，可以用过最小化下面的均方误差来实现

2.3.2 点三元组学习

STAligner使用三元组学习（Triplet learning）来进一步增强编码器对于多ST切片数据的具有批次校正效果的嵌入表示。每个三元组包括一个锚点（anchor）、一个正例点（positive spot）和一个负例点（negative spot）。其中锚点-正例对用于克服批次效应，而锚点-负例对用于指导数据对齐。

为了避免批次效应和原始数据的部分点缺失，三元组在嵌入空间中定义，而非在原始基因表达空间定义。

使用互近邻（mutual nearest neighbor, MNN）来从一对ST切片中确定锚点-正例对，其中锚点和正例分别来自两个切片，锚点和正例具有相似的基因表达情况。负例从锚点所在切片随机选取，由于切片中的点数往往远大于k近邻数量（一个正常的k），因此所选负例在锚点k近邻中的概率很低。

Triplet损失用于最小化锚点与正例的距离并最大化锚点与负例的距离，该距离在嵌入的z空间中计算，如下所示

其中a、p、n分别代表锚点、正例、负例的索引，是如上定义的三元组的集合，一共包含个三元组。

最终整体损失函数被写为如下形式

2.4 多切片对齐与3D重构

与PASTE等方法一样，STAligner也可以进行3D组织结构重建。多切片的对齐是实现3D重构的重要工作。

ICP算法是一个可以将多个切片数据通过刚性变换（平移和旋转）进行对齐的算法，但是鉴于ST切片存在“部分重叠性”和原始数据的缺失，因此难以直接使用。

STAligner从一组需要进行对齐的切片中选取landmark区域，landmark区域在每个切片相互之间应具有足够的重叠，并且应该是不对称的（这与ICP算法的性质有关）。使用landmark区域的MNN点对作为输入，从而得到对于每个切片的刚性变换。再将这个变换应用到对应的整个切片，从而实现3D结构的重建。

简单来说，STAligner的多切片对齐，是基于所有切片共有部分的对齐实现的。

可视化的展示如下图。