论文114：Dynamic graph representation with knowledge-aware attention for WSI (CVPR‘24)

1 要点

题目：用于WSI分类的知识感知动态图表征 (WiKG)

代码：https://github.com/WonderLandxD/WiKG

背景：

1. 已有方法重点关注重要的实例，但难以捕捉实例之间的交互性；
2. 传统的图表征方法利用显式的空间位置构建拓扑结构，但限制了任意位置的实例之间的灵活交互能力，尤其是空间距离远时；

方法：利用动态图构建来量化WSI中区块 (实例) 之间的位置关系：

1. 动态构建邻居和有向边缘嵌入：
   a) 使用Otsu阈值函数区分前景组织区域，并将其划分为非重叠的实例；
   b) 使用特征编码器获取实例嵌入，并使用两个单独的线性层对其嵌入，分别获得头部和尾部；
   c) 使用点积和softmax函数计算头部和尾部的相似度；
   d) 选择具有最高相似度得分的前$k$个实例作为每个实例的邻居；
   e) 使用头部和尾部嵌入分配有向边嵌入；
2. 设计了一种知识感知注意力机制：
   a) 计算邻居的尾部嵌入的线性组合，以描述一阶连续结构；
   b) 使用softmax函数对这些组合归一化；
3. 生成图级嵌入及分类：
   a) 将聚合的邻居信息与原始头部融合，以获得新的头部表征；
   b) softmax函数获得WSI的概率分数；
4. 交叉熵损失引导训练；

数据集：

1. TCGA：包括食管癌、肾癌，以及肺癌；

2 方法

本节主要介绍动态图构建和用于节点更新的知识感知注意力机制，如图2。

2.1 动态图构建

与传统方法利用空间关系构建图表征，WiKG基于可学习的潜在特征来评估实例之间的位置关系：

1. 对于给定的一个WSI，利用Otsu阈值方法来区分前景组织区域，并将其划分为不交叠的实例 X = { x 1 , x 2 , … , x n } X=\{x_1,x_2,\dots,x_n\} X={x1​,x2​,…,xn​}，这也同时表示图节点；
2. 利用特征编码器，例如在ImageNet上训练的Vision Transformer来获取实例嵌入；
3. 使用两个独立的线性层来来分别将实例嵌入为头部$h_i$和尾部$t_i$，其中头部用于探索其它实例与其本身的相关性，尾部用于探索当前实例对其它实例的贡献：
   $$h_i=W_{h}f(X),t_i=W_{t}f(X),\text{\hspace{1em}(1)}$$
4. 计算头部和尾部的点乘，并通过softmax来获取最终的相似性：
   $${\omega }_{i,j}=\frac{h_i^T{t}_j}{{\sum }_{j=1}^N(h_i^T{t}_j)},\text{\hspace{1em}(2)}$$其中${\omega }_{i,j}$表示第$i$个头部和第$j$个尾部之间的相似性。对于每个实例，其相似性最高的$k$个实例被选择作为其邻居实例：
   N ( i ) = { j ∈ V : ω i , j ∈ Topk { w i , j } j = 1 N } , (3) \tag{3} \mathcal{N}(i)=\{ j\in V:\omega_{i,j}\in \text{Topk}\{w_i,j\}_{j=1}^N \}, N(i)={j∈V:ωi,j​∈Topk{wi​,j}j=1N​},(3)其中$V$表示实例集、$|V|=N$，以及$|\mathcal{N}(i)|=k$；
5. 依据拓扑结构将嵌入分配给有向边：
   $$r_{i,j}={\omega }_{i,j}{t}_j+(1-{\omega }_{i,j})h_i,\text{for every }j\in \mathcal{N}(i),\text{\hspace{1em}(4)}$$其中$r_{i,j}$表示实例$j$到实例$i$的边嵌入；
6. 通过以上步骤，每个WSI可以被描述为一个动态图表征$G=(V,\mathcal{E},\mathcal{F},\mathcal{R})$，其中$V$表示节点的集合、$\mathcal{E}$表示边的集合、$\mathcal{F}$表示头部和尾部嵌入的集合，以及$\mathcal{R}$表示有向边嵌入的集合。具体地， E = { ( h , r , t ) : ( h , t ) ∈ F , r ∈ R } \mathcal{E}=\{(h,r,t):(h,t)\in\mathcal{F},r\in\mathcal{R}\} E={(h,r,t):(h,t)∈F,r∈R}是头部、尾部，以及在每个有向边上的高位相似性的集合。
7. 算法1提供了动态图构建的Pytorch风格的伪代码。
   

2.2 知识感知注意力机制

为了充分利用图结构中图节点的关联性，设计用于节点之间信息传递和汇聚的知识感知注意力机制：

1. 对于每个实例$i$，计算头部嵌入和其邻居${\mathcal{N}}_i$之间的线性组合，以刻画其一阶连接结构：
   $$h_{\mathcal{N}(i)}=\sum _{j\in \mathcal{N}(i)}\pi (h_i,r_{i,j},t_j)t_j,\text{\hspace{1em}(5)}$$其中$\pi (h,r,t)$是引导尾部信息汇聚至头部的权重因子，这里使用三元组的非线性组合来计算：
   $$u(h_i,r_{i,j},t_j)=t_j^T\text{tanh}(h_i+r_{i,j}).\text{\hspace{1em}(6)}$$这种组合方式
2. 利用softmax标准化：
   π ( h i , r i , j , t j ) = exp ⁡ { u ( h i , r i , j , t j ) } ∑ j ∈ N ( i ) exp ⁡ { u ( h i , r i , j , t j ) } . (7) \tag{7} \pi(h_i,r_{i,j},t_j)=\frac{\exp\{ u(h_i,r_{i,j},t_j) \}}{\sum_{j\in\mathcal{N}(i)}\exp\{ u(h_i,r_{i,j},t_j) \}}. π(hi​,ri,j​,tj​)=∑j∈N(i)​exp{u(hi​,ri,j​,tj​)}exp{u(hi​,ri,j​,tj​)}​.(7)
   通过建模三元组的关系，并将它们描述为关于边的知识信息，头节点可以高效地度量来自尾部节点的信号。图3展示了所提出的知识感知注意力机制的实现。
   
3. 使用一个双重互动机制来促进节点之间的信息交换：
   $$h_i={\sigma }_1(W_1(h_i+h_{{\mathcal{N}}_{(i)}}))+{\sigma }_2(W_2(h_2\odot h_{N_{(i)}})),\text{\hspace{1em}(8)}$$其中$\sigma$是激活函数，类似于LeakyReLU；
4. 使用Readout函数生成图级别嵌入并用softmax获得最终的WSI预测概率：
   $$\hat{Y}=\text{Softmax}(\text{Readout}(G)),\text{\hspace{1em}(9)}$$其中Readout是一个全局池化函数，类似于均值和最大池化；
5. 在训练期间，选用交叉熵作为损失函数：
   $${\mathcal{L}}_{ce}=1\frac{1}{M}\sum _{m=1}^M\sum _{c=1}^C{Y}_{m,c}\ln ({\hat{Y}}_{m,c}),\text{\hspace{1em}(10)}$$其中$C$是类别数、$M$是训练样本数，以及$Y$是one-hot标签。