DiffPool

Paper : Hierarchical Graph Representation Learning with Differentiable Pooling
Code :

摘要

對於Graph Classification 任務來說，GNN的最後一步需要關注圖數據的全局信息，既包含圖的結構信息，也包含各個節點的屬性信息。如果採用一次性全局池化的方法聚合點上的特徵表示，並沒有將圖的結構考慮進去，而且可能在節點重編號後聚合的結果會發生改變，也就是說聚合過程中依賴節點順序。在CNN模型中，通常採用的做法是通過層次化池化機制來逐漸提取全局信息，作者仿照CNN提出了一種層次化圖聚類的方式，支持添加到GNN模型中進行端到端的學習。

每層將該層的節點映射到若干個cluster中，作爲下一層輸入的節點。

定義

定義圖 $G$ 表示爲 $(A,F)$ ，其中 $A\in\{0,1\}^{n\times n}$ 表示鄰接矩陣，$F\in \mathbb R^{n\times d}$ 表示節點的特徵矩陣。圖分類任務定義爲學習一個映射 $f : \mathcal G \rightarrow \mathcal Y$，這需要首先學習圖的整體特徵表示，即 $f' : \mathcal G\rightarrow \mathbb R^{D}$

圖神經網絡的數學表示爲

$H^{(k)} = M(A,H^{(k-1)};\theta^{(k)})$

其中，$H^{(k)}\in \mathbb R^{n\times d}$，$H^{(0)} = F$，一種常見的GNN模型表示爲

$H^{(k)} = \text{ReLU}(\widetilde D^{-\frac{1}{2}}\widetilde A\widetilde D^{-\frac{1}{2}}H^{(k-1)}W^{(k-1)})$

其中，$\widetilde A = A+I$，$\widetilde D = \text{diag}(\sum_j \widetilde A_{ij})$，爲簡化起見，使用 $Z = \text{GNN}(A,X)$ 表示上式，其中 $Z,X \in \mathbb{R}^{n\times d}$，$Z$ 表示該層輸出的點特徵矩陣，$X$ 表示輸入的點特徵矩陣。對於池化層來說，可以表示爲
$(Z',A') = \text{Pool}(A,X)$
其中 $A'\in \mathbb R^{m\times m}$，$Z' \in \mathbb R^{m\times d}$ ，且 $m<n$

池化

假定 $n_l$ 表示經過了 $l-1$ 層池化之後的節點個數，那麼第 $l$ 層的簇分配矩陣 $S^{(l)}\in \mathbb R^{n_l\times n_{l+1}}$， $S^{(l)}$ 給出的是簇分配的概率形式，是一種軟分配矩陣。如果已知簇分配矩陣，那麼有

$\\X^{(l+1)} = (S^{(l)})^\text T Z^{(l)} \in \mathbb R^{n_{l+1}\times d} \\A^{(l+1)} = (S^{(l)})^\text T A^{(l)} S^{(l)} \in \mathbb R^{n_{l+1}\times n_{l+1}}$

使用兩個獨立的GNN網絡層，一個用來學習點的特徵矩陣，一個用來學習簇分配矩陣

$\\Z^{(l)} = \text{GNN}_\text {embed}^{(l)}(A^{(l)},X^{(l)}) \\S^{(l)} = \text{softmax}(\text{GNN}_\text {pool}^{(l)}(A^{(l)},X^{(l)}))$

對於最後一層池化層，強制 $n_{l+1} = 1$，$S^{(l)}$ 爲全1矩陣。

容易證明只要GNN具有排列不變性(Permutation invariance)，那麼DiffPool具有排列不變性。

輔助損失函數

對於第 $l$ 層，定義損失函數

$L_\text{LP}^{(l)} = ||A^{(l)},S^{(l)}(S^{(l)})^\text T||_F$

矩陣A的Frobenius範數定義爲矩陣A各項元素的絕對值平方的總和
$||A||_F = \sqrt{\sum_i\sum_j|a_{i,j}|^2}$

另一個輔助損失函數是爲了使簇分配函數儘量接近one-hot向量，即定義正則項

$L_\text E^{(l)} = \frac{1}{n_l}\sum_{i=1}^{n_l} H(S_{i}^{(l)})$

其中，$S_{i}^{(l)}$ 表示矩陣 $S^{(l)}$ 的第 $i$ 行，$H$ 表示熵函數

實驗

作者在實驗中觀察到，儘管使用DiffPool性能顯着提高，但可能無法穩定訓練，並且即使使用相同的超參數設置，不同次運行之間的精度也存在顯着差異。 可以看出，添加輔助損失函數可以使訓練更加穩定，並減少不同次運行之間準確性的標準偏差。

作者在GraphSAGE和Structure2Vec兩種算法上進行了實驗，證明了DiffPool的通用性。

作者將不同層的聚類結果可視化顯示出來，用來證明算法學到了有意義的聚類方式。

作者觀察到DIFFPOOL學會了以非均勻的方式將節點摺疊爲簇，並且傾向於將密集連接的子圖摺疊爲簇。由於GNN可以有效地對密集的，類似團的子圖執行消息傳遞(直徑較小)，因此在這樣一個密集的子圖中將節點集中在一起不太可能導致結構信息的丟失。這直觀地解釋了爲什麼摺疊密集子圖是DIFFPOOL的有用合併策略。

由於分配網絡基於輸入節點及其鄰居的特徵來計算簇分配矩陣，因此具有相似輸入特徵和鄰域結構的節點將具有相似的集羣分配。 儘管某兩個節點相距較遠，但它們對於自己和所有鄰居都具有完全相同的鄰域結構和特徵。 在這種情況下，DIFFPOOL網絡將它們分配到同一羣集中，這與其他體系結構中的池化概念不同。

DIFFPOOL對於超參數最大集羣數量是敏感的，作者發現簇分配根據網絡深度和最大集羣數而變化較大。使用較大的最大集羣數，合併的GNN可以對更復雜的層次結構進行建模，但是很大的參數會導致更大的噪聲和更低的效率。 特別的，簇分配矩陣可能不會使用某些羣集。 對應於未使用集羣的列所有節點的值都較低。

總結

作者提出的DIFFPOOL給出了一種有效的層次化聚類的方法，並給出了兩個未來的研究方向

• 學習硬簇分配矩陣(one-hot)
• 將DIFFPOOL應用到其他領域中

個人認爲，作者在實驗中觀察到的三個現象都可以繼續往下研究，尤其是第三點需要改進爲一個對超參數不敏感的聚類方法。