ICML 2024 || 圖外部注意力機制捕捉圖之間的關系

大家好，今天給大家介紹一篇最ICML 2024關于Graph Transformer研究論文，這篇論文提出了一種新的圖外部注意力機制(Graph External Attention, GEA)和基于GEA的圖外部注意力增強Transformer(Graph External Attention Enhanced Transformer, GEAET)模型。這項工作解決了現有圖表示學習方法忽視圖間相關性的問題，通過引入外部注意力單元來隱式學習圖間相關性，從而獲得更全面的圖表示。

1. 基本信息

• 論文題目：Graph External Attention Enhanced Transformer
• 作者：Jianqing Liang, Min Chen, Jiye Liang
• 作者研究單位：Key Laboratory of Computational Intelligence and Chinese Information Processing of Ministry of Education, School of Computer and Information Technology, Shanxi University, Taiyuan 030006, Shanxi, China
• 代碼鏈接：https://github.com/icm1018/GEAET

2. 研究背景

圖表示學習在近年來受到廣泛關注，在社交網絡分析、藥物發現、蛋白質設計和醫療診斷等多個領域發揮著重要作用。早期的圖表示學習研究主要集中在圖神經網絡(GNNs)上，如GCN、GraphSage、GatedGCN和GIN等。然而，基于消息傳遞的GNNs存在一些關鍵限制，包括有限的表達能力、過平滑、過度壓縮和長距離依賴性差等問題。

相對于GNN，圖Transformer（GT）通過自注意力機制捕捉節點間的全局交互信息，解決了一些GNN的典型問題。然而，現有的研究大多只關注圖的內部信息，忽略了圖與圖之間的潛在相關性。本文提出了一種新的外部圖注意力機制（GEA），并基于此設計了一種圖外部注意力增強Transformer（GEAET），以捕捉圖間的相關性并結合局部結構和全局交互信息，從而實現更全面的圖表示。

本文作者認為，在許多實際場景中（如分子圖數據），不同圖之間通常存在強相關性。直觀上，利用圖間相關性可以提高圖表示學習的效果。因此，本文提出了一個關鍵問題：如何將外部信息納入圖表示學習中？

3. 方法

本文提出了兩個主要的創新方法：圖外部注意力（Graph External Attention, GEA）和圖外部注意力增強Transformer（Graph External Attention Enhanced Transformer, GEAET）。

1. 圖外部注意力（GEA）

GEA的核心思想是引入外部的可學習單元來隱式學習圖間相關性。具體來說，GEA計算輸入圖的節點特征與外部單元之間的注意力：

其中，是輸入圖的節點特征矩陣，是一個可學習的參數，可以看作是具有個節點的外部單元，作為所有輸入圖的共享存儲。表示輸入圖節點與外部單元之間的相似度。

為了處理注意力矩陣對輸入特征尺度的敏感性，作者采用了雙重歸一化技術：

在實際應用中，為了提高網絡的能力，作者使用了兩個不同的外部單元分別作為鍵和值。此外，為了利用輸入圖中的邊信息，還引入了額外的外部單元用于邊特征，以及一個共享單元來存儲邊和節點之間的連接：

其中，是存儲邊和節點之間連接的共享單元；是節點的外部鍵值單元，是邊的外部鍵值單元。

2. 圖外部注意力增強Transformer（GEAET）

GEAET框架由兩個主要組件組成：圖嵌入和特征提取層。

圖嵌入部分首先對輸入節點特征和邊特征進行線性投影，然后使用位置編碼增強節點特征：

特征提取層在每一層中通過GEANet捕獲外部特征信息，然后與通過消息傳遞GNN和Transformer獲得的圖內信息聚合，更新節點特征：

最后，使用前饋網絡(FFN)塊聚合節點信息，得到節點表示：

這種設計使GEAET能夠同時利用圖間相關性、局部結構和全局交互信息，從而獲得更全面的圖表示。

4. 實驗發現

本文在多個基準數據集上進行了廣泛的實驗，包括CIFAR10、MNIST、PATTERN、Peptides-Struct、PascalVOC-SP和COCO-SP等。主要實驗發現如下：

1. GEAET性能優越：在多個數據集上，GEAET取得了state-of-the-art的性能。特別是在PascalVOC-SP和COCO-SP數據集上，GEAET的F1分數分別達到0.4585和0.3895，顯著超過其他模型。

2. GEANet對基礎GNN的改進：實驗表明，將GEANet與基礎消息傳遞GNN（如GCN、GINE、GatedGCN）結合，可以顯著提高性能。例如，在Peptides-Struct數據集上，GCN+GEANet的MAE降低到0.2512，而原始GCN的MAE為0.3496。

3. 注意力頭數的影響：在Peptides-Struct數據集上的實驗顯示，GEANet在8個注意力頭時達到最佳性能，而Transformer在單頭注意力時表現最好。這表明GEANet更能有效利用多頭注意力機制。

4. 位置編碼的影響：實驗結果表明，GEANet對位置編碼的依賴性較低。在沒有位置編碼的情況下，GEANet仍然能夠取得良好的性能，而Transformer的性能則顯著下降。

5. 解決過度壓縮問題：在TreeNeighbourMatch數據集上的實驗顯示，GEAET能夠有效緩解過度壓縮問題，在問題半徑r=7時仍能保持良好的泛化能力，而標準的消息傳遞GNN在r=4時就開始失效。

6. 注意力可視化：在ZINC分子圖數據集上的注意力可視化實驗表明，GEANet學習到的注意力分數更加稀疏和信息豐富，更關注重要原子（如N和O）以及連接不同結構單元的原子。這種注意力分布更接近原始分子圖的結構分布，有助于更準確地預測受限溶解度。

5. 結論

本文提出的圖外部注意力（GEA）機制和基于GEA的圖外部注意力增強Transformer（GEAET）模型，為圖表示學習領域提供了一個新的視角。通過引入外部注意力單元，GEAET能夠有效地捕捉圖間相關性，同時整合局部結構和全局交互信息，從而獲得更全面的圖表示。

實驗結果證明了GEAET在多個基準數據集上的卓越性能，特別是在處理復雜的圖結構和長程依賴性方面表現出色。此外，GEAET還展示了良好的可解釋性和對位置編碼的低依賴性，這為未來的圖表示學習研究提供了新的思路。

盡管GEAET取得了顯著的成果，但作者也指出了未來的研究方向，包括降低高內存消耗和時間復雜度，以及解決表達能力上界缺失的問題。總的來說，這項工作為圖表示學習領域提供了一個有潛力的新方法，有望在各種實際應用中發揮重要作用。

請加小助理加入AIGC技術交流群

備注公司/學校+昵稱+研究方向