KTN:ラベルが少なくたって異種混合グラフニューラルネットだって転移学習したいです(2/2)

１．KTN:ラベルが少なくたって異種混合グラフニューラルネットだって転移学習したいです(2/2)まとめ

・ラベルが豊富な転移元に対してHGNNモデルと分類器を事前に訓練し、同じHGにあるラベルが少ない転移先に対してモデルを再利用する手法を考案
・HGNNは多重性に対応するため、各ノードタイプに特化したモジュールで構成されており通常はノードタイプが異なるモデル間で転移は難しい
・どちらかのembeddingsが与えられた場合、もう一方のノードembeddingsを計算できるように連結ノードとしてembeddingsを集約して対処した

２．KTNの性能

以下、ai.googleblog.comより「Teaching old labels new tricks in heterogeneous graphs」の意訳です。元記事の投稿は2023年3月1日、Minji YoonさんとBryan Perozziさんによる投稿です

アイキャッチ画像はカスタムStable Diffusionによる生成

私たちのアプローチ：異種グラフ内のノードタイプ間の転移学習

ここでは、より実用的な学習元領域である、同じHG上にある豊富なラベルを持つ他のノードタイプに光を当てます。

HGを追加する代わりに、単一のHG内の知識(実務者が完全に所有していると仮定)を異なるタイプのノード間で転移します。

具体的には、ラベルが豊富なノードタイプ(転移元)に対して異種グラフニューラルネットワーク(HGNN:Heterogeneous graph neural networks)モデルと分類器を事前に訓練し、同じHGにあるラベルがゼロのノードタイプ(転移先)に対してモデルを再利用します。

1つの要件は、転移先と転移元のノードタイプが同じラベルセットを共有することです。(例えば、電子商取引のHGでは、商品ノードは商品カテゴリを記述するラベルセットを持ち、ユーザーノードはお気に入りの商品カテゴリを記述する同じラベルセットを共有しています)。

なぜ、これは難しいのでしょうか？

残念ながら、事前に訓練されたHGNNと分類器を転移先のノードタイプに直接再利用することはできません。HGNNアーキテクチャの重要な特徴の1つは、HGの多重性を完全に学習するために、各ノードタイプに特化したモジュールで構成されていることです。

HGNNは、ノードタイプごとに異なるモジュールのセットを用いてembeddings計算を行います。下図では、青色と赤色のモジュールが、それぞれ転移元ノードと転移先ノードのembeddingsを計算するために使用されています。

HGNNは、各ノードタイプに特化したモジュールで構成され、異なるノードタイプのembeddingsを計算するために、異なるモジュールセットを使用します。詳細は論文に記載されています。

HGNNを転移元ノードタイプで事前学習する際、HGNNの転移元に特化したモジュールはよく学習できますが、転移先に特化したモジュールは、少量の勾配しか流入しないため、学習不足となります。これは下図に示すように、転移元ノードタイプ(すなわちMtt)の勾配のL2ノルムは、転移先タイプ(すなわちMss)よりもはるかに低いことがわかります。この場合、HGNNモデルは転移先ノードタイプに対して貧弱なノードembeddings を出力し、その結果、タスクのパフォーマンスが低下しています。

HGNNでは、転移先のノードタイプに特化したモジュールは、転移元ノードタイプの事前学習時にゼロまたは少量の勾配しか受け取らず、これは、転移先ノードタイプでの性能低下につながります。

KTN:HGNNのための学習可能な交差型転移学習(cross-type transfer learning)

私たちの研究は、事前に訓練されたHGNNモデルによって計算された(貧弱な)転移先ノードのembeddingsを、転移元ノードのembeddingsの分布に従うように変換することに焦点を当てています。

そうすれば、転移元ノードのタイプについて事前に訓練された分類器を、転移先ノードのタイプに再利用することができます。

では、どのようにして転移先ノードのembeddingsを転移元領域にマッピングすればよいのでしょうか？

この問いに答えるため、HGNNがどのようにノードembeddingsを計算し、転移元と転移先の分布の関係を学習しているかを調査します。

HGNNは、各層で転移先ノードのembeddingsを補強するために、連結ノード(connected node)のembeddingsを集約します。

つまり、転移元ノードと転移先ノードのノードembeddingsは両方とも、同じ入力、つまり前の層の任意の連結ノードタイプのノードembeddingsを使って更新されます。これは、どちらかのノードタイプのノードembeddingsが与えられた場合、もう一方のノードタイプのノードembeddingsを計算できるということです。

この関係を理論的に証明し、転移先領域から転移元領域へのマッピング行列(HGNNパラメータで定義)が存在することを見出しました。(詳細は論文中の定理1に記載)

この定理に基づき、私たちは知識転移ネットワーク(KTN:Knowledge Transfer Network)と呼ぶ補助ニューラルネットワークを導入し、転移先ノードのembeddingsを受け取り、(訓練可能な)マッピング行列と掛け合わせることで変換します。そして、KTNを学習させるために、事前学習段階で性能損失とともに最小化される正則化器を定義します。テスト時には、学習済みのHGNNから計算されたターゲットembeddingsを、学習済みのKTNを使って転移元ドメインにマッピングし、分類します。

HGNNでは、転移元型と転移先型の最終的なノードembeddingsは、同じ入力である前の層のノードembeddingsを使用する異なる数学関数(f()：転移元,g()：転移先)から計算されます。

実験結果

KTNの有効性を検証するため、2つの公開異種グラフであるOpen Academic GraphとPubmedで、18種類のゼロショット転移学習タスクを実行しました。KTNを8つの最先端の転移学習法(DAN, JAN, DANN, CDAN, CDAN-E, WDGRL, LP, EP)と比較しました。以下に示すように、KTNはすべてのタスクで常にすべてのベースラインを上回り、転移学習ベースラインを最大140%(ランキング指標である正規化割引累積利得(Normalized Discounted Cumulative Gain)で測定)上回りました。

Open Academic Graph(OAG-CS)とPubmedデータセットに対するゼロショット転移学習
色は、転移学習のベースラインの異なるカテゴリを表しています。
黄色：分布の統計的性質(例：平均、分散)を利用
緑色：敵対的なモデルを使って知識を転移す
オレンジ：ラベル伝搬を利用し、グラフ構造から直接知識を転移

最も重要なことは、KTNはノードとエッジの種類に応じたパラメータを持つほぼ全てのHGNNモデルに適用でき、対象領域でのゼロショット性能を向上させることができることです。以下に示すように、KTNは6種類のHGNNモデル(R-GCN、HAN、HGT、MAGNN、MPNN、H-MPNN)において、ラベルがゼロのノードタイプに対する精度を最大190%向上させました。

KTNは6種類のHGNNモデルに適用でき、対象領域でのゼロショット性能を向上させることができます。

お持ち帰り頂きたい事

産業界における様々な産業関連のつながりや関係性は異種グラフとして表現されることがあります。HGNNは異種グラフの情報を効果的な表現に要約します。しかし、ある種のノードにおけるラベルの希少性の問題が、HGNNの幅広い応用を妨げています。

本投稿では、HGNNのために設計された最初の異種間転移学習法であるKTNを紹介しました。KTNを用いれば、ラベルの希少性に関係なく、HGNNを介して異種グラフの豊かさを十分に活用することができます。詳しくは論文をご覧ください。

謝辞

本稿は、共著者のJohn Palowitch(Google Research)、Dustin Zelle(Google Research)、Ziniu Hu(Intern, Google Research)、Russ Salakhutdinov(CMU)との共同研究です。このブログ記事のアニメーション図を作成してくれたTom Smallに感謝します。