隨著深度學(xué)習(xí)模型的應(yīng)用和推廣，人們逐漸發(fā)現(xiàn)模型常常會利用數(shù)據(jù)中存在的虛假關(guān)聯(lián)（Spurious Correlation）來獲得較高的訓(xùn)練表現(xiàn)。但由于這類關(guān)聯(lián)在測試數(shù)據(jù)上往往并不成立，因此這類模型的測試表現(xiàn)往往不盡如人意 [1]。其本質(zhì)是由于傳統(tǒng)的機(jī)器學(xué)習(xí)目標(biāo)（Empirical Risk Minimization，ERM）假設(shè)了訓(xùn)練測試集的獨(dú)立同分布特性，而在現(xiàn)實(shí)中該獨(dú)立同分布假設(shè)成立的場景往往有限。在很多現(xiàn)實(shí)場景中，訓(xùn)練數(shù)據(jù)的分布與測試數(shù)據(jù)分布通常表現(xiàn)出不一致性，即分布偏移（Distribution Shifts），旨在提升模型在該類場景下性能的問題通常被稱為分布外泛化（Out-of-Distribution）問題。關(guān)注學(xué)習(xí)數(shù)據(jù)中的相關(guān)性而非因果性的 ERM 等一類方法往往難以應(yīng)對分布偏移。盡管近年涌現(xiàn)了諸多方法借助因果推斷（Causal Inference）中的不變性原理（Invariance Principle）在分布外泛化（Out-of-Distribution）問題上取得了一定的進(jìn)展，但在圖數(shù)據(jù)上的研究依然有限。這是因?yàn)閳D數(shù)據(jù)的分布外泛化比傳統(tǒng)的歐式數(shù)據(jù)更加困難，給圖機(jī)器學(xué)習(xí)帶來了更多的挑戰(zhàn)。本文以圖分類任務(wù)為例，對借助因果不變性原理的圖分布外泛化進(jìn)行了探究。

近年來，借助因果不變性原理，人們在歐式數(shù)據(jù)的分布外泛化問題中取得了一定的成功，但對圖數(shù)據(jù)的研究仍然有限。與歐式數(shù)據(jù)不同，圖的復(fù)雜性對因果不變性原理的使用以及克服分布外泛化難題提出了獨(dú)特的挑戰(zhàn)。

為了應(yīng)對該挑戰(zhàn)，我們在本工作中將因果不變性融入到圖機(jī)器學(xué)習(xí)中，并提出了因果啟發(fā)的不變圖學(xué)習(xí)框架，為解決圖數(shù)據(jù)的分布外泛化問題提供了新的理論和方法。

論文已在 NeurIPS 2022 發(fā)表，本工作由香港中文大學(xué)、香港浸會大學(xué)， 騰訊 AI Lab 以及悉尼大學(xué)合作完成。 

• 論文標(biāo)題：Learning Causally Invariant Representations for Out-of-Distribution Generalization on Graphs
• 論文鏈接：https://openreview.net/forum?id=A6AFK_JwrIW
• 項(xiàng)目代碼：https://github.com/LFhase/CIGA

圖數(shù)據(jù)的分布外泛化

圖數(shù)據(jù)的分布外泛化難在哪？

圖神經(jīng)網(wǎng)絡(luò)近年來在涉及圖結(jié)構(gòu)的機(jī)器學(xué)習(xí)應(yīng)用，如推薦系統(tǒng)、AI 輔助制藥等領(lǐng)域，取得了很大的成功。然而，因現(xiàn)有的大部分的圖機(jī)器學(xué)習(xí)算法都依賴于數(shù)據(jù)的獨(dú)立同分布假設(shè)，使得當(dāng)測試數(shù)據(jù)和訓(xùn)練數(shù)據(jù)出現(xiàn)偏移（Distribution Shifts）時，算法的性能會極大下降。同時，因?yàn)閳D數(shù)據(jù)結(jié)構(gòu)的復(fù)雜性，導(dǎo)致圖數(shù)據(jù)的分布外泛化相比于歐式數(shù)據(jù)更普遍且更具挑戰(zhàn)性。

圖 1. 圖上的分布偏移示例。

首先，圖數(shù)據(jù)的分布偏移可以出現(xiàn)在圖的節(jié)點(diǎn)特征分布中（Attribute-level Shifts）。例如，在推薦系統(tǒng)中，訓(xùn)練數(shù)據(jù)涉及的商品可能采自一些比較流行的類別，涉及到的用戶也可能來自于某些特定地區(qū)，而在測試階段，系統(tǒng)則需要妥善處理所有類別以及地區(qū)的用戶和商品 [2,3,4]。此外，圖數(shù)據(jù)的分布偏移還可以出現(xiàn)在圖的結(jié)構(gòu)分布中（Structure-level Shifts）。早在 2019 年，人們就注意到，在較小的圖上進(jìn)行訓(xùn)練得到的圖神經(jīng)網(wǎng)絡(luò)難以學(xué)到有效的注意力（Attention）權(quán)重以泛化到更大的圖上 [5]，這也推動了一系列相關(guān)工作的提出 [6,7]。在現(xiàn)實(shí)場景中，這兩類分布偏移往往可能同時出現(xiàn)，并且這些不同層級的分布偏移還可以和所要預(yù)測的標(biāo)簽具有不同的虛假關(guān)聯(lián)模式。如在推薦系統(tǒng)中，來自特定類別的商品與特定地區(qū)的用戶往往會在商品用戶交互圖上展現(xiàn)獨(dú)特的拓?fù)浣Y(jié)構(gòu) [4]。在藥物分子屬性預(yù)測中，訓(xùn)練時涉及的藥物分子可能偏小，同時預(yù)測的結(jié)果也會受到實(shí)驗(yàn)測定環(huán)境的影響 [8]。

此外，歐式空間的分布外泛化往往會假設(shè)數(shù)據(jù)來自于多個環(huán)境（Environment）或者域（Domain），并進(jìn)一步假設(shè)訓(xùn)練期間模型能夠獲取訓(xùn)練數(shù)據(jù)中每個樣本所屬的環(huán)境，以此來發(fā)掘跨越環(huán)境的不變性。然而，要獲得數(shù)據(jù)的環(huán)境標(biāo)簽往往需要和數(shù)據(jù)相關(guān)的一些專家知識，而由于圖數(shù)據(jù)的抽象性，使得圖數(shù)據(jù)的環(huán)境標(biāo)簽獲得更加昂貴。因此，大部分現(xiàn)有的圖數(shù)據(jù)集如 OGB 都不含此類環(huán)境標(biāo)簽信息，即便少部分如 DrugOOD 數(shù)據(jù)集存在環(huán)境標(biāo)簽，但也存在不同程度的噪聲。

現(xiàn)有方法能否解決圖上的分布外泛化問題？

為了對圖數(shù)據(jù)分布外泛化的挑戰(zhàn)有一個直觀的理解，我們基于 Spurious-Motif [9] 數(shù)據(jù)集構(gòu)建新的數(shù)據(jù)以進(jìn)一步實(shí)例化上述幾大挑戰(zhàn)，并嘗試使用現(xiàn)有的方法如歐式數(shù)據(jù)上分布外泛化的訓(xùn)練目標(biāo) IRM [10]，或者具有更強(qiáng)表達(dá)能力的 GNN [11]，分析能否通過已有的方法解決圖數(shù)據(jù)的分布外泛化問題。

圖 2. Spurious Motif 數(shù)據(jù)集示例。

Spurious Motif 任務(wù)如圖 2 所示，主要根據(jù)輸入的圖中是否含有特定結(jié)構(gòu)的子圖（如 House，或者 Cycle）對圖標(biāo)簽進(jìn)行判斷，其中節(jié)點(diǎn)顏色代表節(jié)點(diǎn)的屬性。使用該數(shù)據(jù)集可以比較清晰地測試不同層級的分布偏移對圖神經(jīng)網(wǎng)絡(luò)性能的影響。對于一個使用 ERM 進(jìn)行訓(xùn)練的普通 GNN 模型：

• 如果訓(xùn)練階段大部分有 House 子圖的樣本都節(jié)點(diǎn)大部分綠色，而 Cycle 則是藍(lán)色，那么在測試階段，模型則傾向于預(yù)測任何含大量綠色節(jié)點(diǎn)的圖為 “House”，而藍(lán)色節(jié)點(diǎn)的圖為 “Cycle”。
• 如果訓(xùn)練階段大部分有 House 子圖的樣本都與一個六邊形子圖共同出現(xiàn)，那么在測試階段，模型則傾向于判定任何含有六邊形結(jié)構(gòu)的圖為 “House”。

此外，模型在訓(xùn)練時無法獲得任何和環(huán)境標(biāo)簽相關(guān)的信息，得到實(shí)驗(yàn)結(jié)果如圖 3 所示（更多結(jié)果可以查閱論文附錄 D）。

圖 3. 現(xiàn)有方法在不同圖分布偏移下的表現(xiàn)。

如圖 3 所示，普通的 GCN 不論是在使用 ERM 或者 IRM 訓(xùn)練，都無法應(yīng)對圖的結(jié)構(gòu)偏移（Struc）；而在增加了圖節(jié)點(diǎn)屬性偏移（Mixed）以及圖大小分布偏移后（圖 3 中），模型性能將進(jìn)一步降低；此外即便使用具有更強(qiáng)表達(dá)能力的 kGNN 也難以避免嚴(yán)重的性能損失（平均性能的降低，或更大的方差）。

由此，我們自然地引出所要研究的問題：如何才能獲得一個具有應(yīng)對多種圖分布偏移的 GNN 模型？

面向圖數(shù)據(jù)分布外泛化的因果模型

為了解決上述問題，我們需要對學(xué)習(xí)目標(biāo)，即不變圖神經(jīng)網(wǎng)絡(luò)（Invariant GNN），進(jìn)行定義，即在最糟糕的環(huán)境下仍舊表現(xiàn)良好的模型（嚴(yán)謹(jǐn)?shù)亩x參見論文）：

定義 1（不變圖神經(jīng)網(wǎng)絡(luò)）給定一系列收集自不同的具有因果關(guān)聯(lián)的環(huán)境的圖分類數(shù)據(jù)集，其中包含被認(rèn)為是來自環(huán)境 e 的獨(dú)立同分布樣本，考慮一個圖神經(jīng)網(wǎng)絡(luò)，其中和分別是作為輸入的圖空間和樣本空間，f 是不變圖神經(jīng)網(wǎng)絡(luò)，當(dāng)且僅當(dāng)，即最小化所有環(huán)境的最壞經(jīng)驗(yàn)損失 (worst empirical risk)，其中為模型在環(huán)境中的經(jīng)驗(yàn)損失。

模型在訓(xùn)練時只能獲得部分的訓(xùn)練環(huán)境中的數(shù)據(jù)，如果不對數(shù)據(jù)的過程進(jìn)行任何假設(shè)，不變圖神經(jīng)網(wǎng)絡(luò)定義所要求的 minmax 最優(yōu)性是很難做到的。因此，我們從因果推斷（Causal Inference）的角度使用因果模型（Structural Causal Model）對圖的生成過程進(jìn)行建模，并對環(huán)境之間的關(guān)聯(lián)進(jìn)行刻畫，以嘗試定義圖數(shù)據(jù)上的因果不變性。

圖 4. 圖數(shù)據(jù)生成過程的因果模型。

不失一般性，我們將所有影響圖生成的隱變量納入隱空間，并將圖的生成過程建模為。此外，對于隱變量，根據(jù)其是否受環(huán)境 E 影響，我們將其劃分成不變隱變量（invariant latent variable）以及虛假隱變量（spurious latent variable）。對應(yīng)地，隱變量 C 與 S 分別會影響 G 的某個子圖的生成，分別記作不變子圖以及虛假子圖，如圖 4 (a) 所示，而 C 主要控制了圖的標(biāo)簽 Y。這也可以進(jìn)一步推出，即 C 與 Y 相比于 S 有更高的互信息。這樣的生成過程與許多實(shí)際例子相對應(yīng)，如一個分子的藥化屬性通常由某個關(guān)鍵的基團(tuán)（分子子圖）決定（如羥基 - HO 之于分子的水溶性）。

此外，C 與Y，S以及 E 在隱空間有多種類型的交互，主要跟進(jìn)虛假隱變量 S 與標(biāo)簽 Y 是否在有不變隱變量 C 之外額外的關(guān)聯(lián)，即，可以概括為兩種：如圖 4 (b) 的 FIIF（Fully Informative Invariant Feature）以及圖 4 (c) 的 PIIF（Partially Informative Invariant Feature）。其中 FIIF 表示給定不變信息后標(biāo)簽與虛假相關(guān)量獨(dú)立。PIIF 則相反。需要說明的是，為了盡可能地覆蓋更多的圖分布偏移，我們的因果模型致力于對各種圖生成模型的廣泛的建模。如有更多關(guān)于圖生成過程的知識，圖 4 所示的因果模型則可以進(jìn)一步泛化到更具體的例子。如在附錄 C.1 中，我們展示了如何通過增加額外圖極限（graphon）的假設(shè)，將因果圖泛化至先前 Bevilacqua 等人用于分析圖大小分布偏移的工作 [7]。

基于上述的因果分析，我們可以知道，當(dāng)模型只使用不變子圖 進(jìn)行預(yù)測的時，即只使用之間的關(guān)聯(lián)，模型的預(yù)測才不會受到環(huán)境 E 的改變而影響；反之，如果模型的預(yù)測依賴于任何與 S 或有關(guān)的信息，其預(yù)測結(jié)果將會因?yàn)?E 的變化發(fā)生極大的改變，從而出現(xiàn)性能損失。因此，我們的目標(biāo)可以從學(xué)習(xí)一個不變圖神經(jīng)網(wǎng)絡(luò)，進(jìn)一步細(xì)化至：a) 識別潛在的不變子圖；b) 用識別的子圖預(yù)測 Y。為了進(jìn)一步與數(shù)據(jù)生成的算法過程相對應(yīng)，我們進(jìn)一步把圖神經(jīng)網(wǎng)絡(luò)拆分為子圖識別網(wǎng)絡(luò)（Featurizer GNN）和分類網(wǎng)絡(luò)（Classifier GNN），且，其中為的子圖空間。那么模型的學(xué)習(xí)目標(biāo)則可表示為如公式 (1) 所示：

其中，，為子圖識別網(wǎng)絡(luò)對不變子圖的預(yù)測；為與 Y 的互信息，通常，最大化可以通過最小化使用預(yù)測 Y 的經(jīng)驗(yàn)損失實(shí)現(xiàn)。然而，由于 E 的缺失，我們難以直接使用 E 對進(jìn)行獨(dú)立性的驗(yàn)證，為此，我們必須尋求其他等價條件以識別需要的不變子圖。

因果啟發(fā)的不變圖學(xué)習(xí)

 為了解決在缺失時的不變子圖識別問題，基于公式 (1) 的框架，我們希望尋求一個公式 (1) 的易于實(shí)現(xiàn)的等價條件。特別地，我們首先考慮一種比較簡單的情況，即潛在的不變子圖大小固定且已知，。在這樣的條件下，考慮最大化，盡管與有同樣的大小，但因?yàn)?/span>與 Y 也存在關(guān)聯(lián)，所以在沒有任何其他約束的情況下，最大化可能會使得估計得到的不變子圖中包含部分與 Y 有互信息的虛假子圖。

為了將中可能的虛假子圖部分 “擠” 出去，我們將進(jìn)一步從因果模型中尋求更多關(guān)于特有的屬性。注意到，不論是 PIIF 還是 FIIF 的虛假關(guān)聯(lián)類型，對于最大化與標(biāo)簽 Y 互信息的子圖，我們有：

• 不同環(huán)境，中與相同不變隱變量 C 的不變子圖是這兩個環(huán)境中互信息最大的兩個子圖，即；
• 同一個環(huán)境中對應(yīng)不同不變隱變量 C 的不變子圖兩個不變子圖是這個環(huán)境中互信息最小的兩個子圖，即；

結(jié)合上述兩個性質(zhì)，我們可以推出

由于在實(shí)踐中我們難以直接觀察得到，我們則可以通過作為在公式 (2) 中的代理使用。

同時，當(dāng)和同時達(dá)到最大化時，將自動最小化，否則模型的預(yù)測將坍縮至平凡解。由此，我們得到了在一種簡單情況下的不變子圖等價條件，結(jié)合公式 (1)，我們得到了第一版因果啟發(fā)的不變圖學(xué)習(xí)（Causality-inspired Invariant Graph leArning）框架，即 CIGAv1:

其中，且，即與 G 來自同個類別 Y。我們在論文中進(jìn)一步證明了 CIGAv1 在已知圖大小情況下能成功識別圖 4 對應(yīng)的因果模型中潛在的不變子圖。然而，由于先前的假設(shè)過于理想化，在實(shí)踐中，不變子圖的大小可能會發(fā)生改變同時對應(yīng)的大小我們也往往無法得知。在沒有子圖大小的假設(shè)下，只需要將全圖識別為不變子圖即能滿足 CIGAv1 的要求。因此，我們考慮進(jìn)一步尋求關(guān)于不變子圖別的性質(zhì)用于去除這一假設(shè)。

注意到，在最大化時，可能出現(xiàn)? ?中的虛假子圖部分與被去除的不變子圖部分享有同樣的和相關(guān)的互信息。那么，我們能否反其道而行之，同時最大化以去除中可能的虛假子圖部分呢？答案是肯定的，我們可以利用與 Y 的關(guān)聯(lián)令其與的估計互相競爭。需要注意的是，在最大化時需要保證不會超過，否則將預(yù)測的又將陷入平凡解。結(jié)合這一額外的條件，我們則可以將關(guān)于不變子圖大小的假設(shè)從公式 (3) 去除，得到如下 CIGAv2：

圖 5. 因果啟發(fā)的不變圖學(xué)習(xí)框架示意。

CIGA 的實(shí)現(xiàn)：在實(shí)踐中，估計兩個子圖的互信息通常比較困難，而監(jiān)督式的對比學(xué)習(xí) [11] 則提供了一種可行的解法：

其中對應(yīng)著公式 (4) 中的正樣本，而則是對應(yīng)于的圖表示。當(dāng)時，公式 (5) 提供了對于的一種基于 von Mises-Fisher kernel density 的非參數(shù)再代入熵估計（Nonparameteric Resubstitution Entropy Estimator ）[13,14]。最終 CIGA 核心部分的實(shí)現(xiàn)如圖 5 所示，即通過在隱表示空間拉近同個類別不變子圖的圖表示，同時最大化不同類別不變子圖的圖表示，以最大化。此外，對于公式 (4) 中的另一個約束，我們則可以通過鉸鏈損失（hinge loss）的思路進(jìn)行實(shí)現(xiàn)，即，只優(yōu)化預(yù)測時經(jīng)驗(yàn)損失大于對應(yīng)的不變子圖的虛假子圖。?

實(shí)驗(yàn)與討論

在實(shí)驗(yàn)中，我們使用 16 個合成或來自真實(shí)世界的數(shù)據(jù)集，對 CIGA 在不同圖分布偏移下進(jìn)行了充分的驗(yàn)證。在實(shí)驗(yàn)中，我們使用可解釋 GNN 框架 [9] 實(shí)現(xiàn)了 CIGA 的原型，而實(shí)際上 CIGA 有更多實(shí)現(xiàn)的方式。具體的數(shù)據(jù)集以及實(shí)驗(yàn)細(xì)節(jié)詳見文中實(shí)驗(yàn)部分。

合成數(shù)據(jù)集上圖結(jié)構(gòu)分布偏移以及混合分布偏移的表現(xiàn)

我們首先基于 SPMotif 數(shù)據(jù)集 [9] 構(gòu)造了 SPMotif-Struc 以及 SPMotif-Mixed 數(shù)據(jù)集，其中 SPMotif-Struc 包含了特定子圖與圖中其他子圖結(jié)構(gòu)的虛假關(guān)聯(lián)，以及圖大小的分布偏移；而 SPMotif-Mixed 則在 SPMotif-Struc 的基礎(chǔ)上新增了圖節(jié)點(diǎn)屬性層級的分布偏移。表中第一欄為 ERM 以及可解釋 GNN 的基線，第二欄則為歐式空間最先進(jìn)的分布外泛化算法。從結(jié)果中可以發(fā)現(xiàn)，不論是更好的 GNN 框架還是歐式空間的分布外泛化算法，都受制于圖上的分布偏移，且當(dāng)更多的分布偏移出現(xiàn)時，性能損失（更小的平均分類性能或更大的方差）將進(jìn)一步增強(qiáng)。相對的，CIGA 則能在不同強(qiáng)度的分布偏移下保持良好的性能，并極大超越最好的基線表現(xiàn)。

真實(shí)數(shù)據(jù)集上各類圖分布偏移的表現(xiàn)

我們接著在真實(shí)數(shù)據(jù)集和各種真實(shí)數(shù)據(jù)中存在的圖分布偏移進(jìn)一步測試了 CIGA 的表現(xiàn)，包括來自 AI 輔助制藥中藥物分子屬性預(yù)測的 DrugOOD 中三種不同環(huán)境劃分（實(shí)驗(yàn)環(huán)境 Assay，分子骨架 Scaffold，分子大小 Size）的三個數(shù)據(jù)集，包含了各種真實(shí)應(yīng)用場景的圖分布偏移；基于歐式空間中經(jīng)典的圖像數(shù)據(jù)集 ColoredMNIST [10] 轉(zhuǎn)換得到的 CMNIST-SP，主要包含圖節(jié)點(diǎn)屬性的 PIIF 類型分布偏移；基于自然語言情感分類數(shù)據(jù)集 SST5 以及 Twitter 轉(zhuǎn)化得到的 Graph-SST5 以及 Twitter [15]，并且額外添加了圖度數(shù)的分布偏移。此外，我們還使用了先前研究較多的 4 個分子圖大小分布偏移數(shù)據(jù)集 [7]，

測試結(jié)果如上表所示，可以發(fā)現(xiàn)，在真實(shí)數(shù)據(jù)中，由于任務(wù)難度增加，使用更好架構(gòu)的 GNN 或者歐式空間的分布外泛化優(yōu)化目標(biāo)訓(xùn)練得到的模型性能甚至弱于使用 ERM 訓(xùn)練得到的普通 GNN 模型。這一現(xiàn)象也與歐式空間中更難任務(wù)下的分布外泛化實(shí)驗(yàn)觀察得到的現(xiàn)象類似 [16]，反應(yīng)了真實(shí)數(shù)據(jù)上的分布外泛化難度以及現(xiàn)有方法的不足。與之相對地，CIGA 則能在所有的真實(shí)數(shù)據(jù)和圖分布偏移上獲得提升，甚至在某些數(shù)據(jù)集如 Twitter、PROTEINS 中達(dá)到經(jīng)驗(yàn)最優(yōu)的 Oracle 水準(zhǔn)。在最新的圖分布外泛化測試基準(zhǔn) GOOD 上圖分類數(shù)據(jù)集的初步測試也顯示了 CIGA 是目前最好且能應(yīng)對各種個樣的圖分布偏移的圖分布外泛化算法。

由于使用了可解釋 GNN 作為 CIGA 的原型實(shí)現(xiàn)架構(gòu)，我們也對模型識別得到的 DrugOOD 中的進(jìn)行了可視化，發(fā)現(xiàn) CIGA 確實(shí)發(fā)現(xiàn)了一些比較一致的分子基團(tuán)用于分子屬性預(yù)測。這可以為后續(xù) AI 輔助制藥提供更好的依據(jù)。

圖 6. DrugOOD 中 CIGA 識別得到的部分不變子圖。

總結(jié)及展望

 本文通過因果推斷的角度，首次將因果不變性引入至多種圖分布偏移下的圖分布外泛化問題中，并提出了一個全新的具有理論保證的解決框架 CIGA。大量實(shí)驗(yàn)也充分驗(yàn)證了 CIGA 優(yōu)秀的分布外泛化性能。放眼未來，基于 CIGA，我們可以進(jìn)一步探索更好的實(shí)現(xiàn)框架 [17]，或?yàn)?CIGA 引入更好的具有理論保障的數(shù)據(jù)增強(qiáng)方法 [3,18]，并在理論上建模納入圖上的協(xié)變量偏移（Covariate Shift）[19]，以進(jìn)一步提升 CIGA 識別不變子圖的能力，促進(jìn)圖神經(jīng)網(wǎng)絡(luò)在 AI 輔助制藥等真實(shí)應(yīng)用場景的真實(shí)落地使用。