はじめに

グラフ埋め込み（Graph Embedding）は、グラフの要素を低次元のベクトル空間にマッピングする手法です。グラフ構造をコンパクトで連続的なベクトル表現に変換することにより、機械学習などの下流タスクで扱いやすくなります。

埋め込みの手法は、主に「ノードの埋め込み」と「グラフ全体の埋め込み」の二つに分類されます。本記事では、後者のグラフ全体の埋め込みの代表的な手法である、Graph2vec [1] について、簡単な解説とPythonでの実行例を紹介します。

目次

前半では理論をザックリと解説します。後半では、タンパク質のグラフ構造を例に、実際にPythonでGraph2vecを実行した際の結果を評価します。

解説： Graph2vecとは
 実行： Graph2vecをPythonで実行する

Graph2vecとは

Graph2vecは、グラフ集合$\{G_1, G_2, ...\}$が得られた際に、各グラフの特徴ベクトルを学習します。ここで、各グラフは$G_i = (N, E, \lambda)$のように表現され、$V$と$E$はノードとエッジの集合を、$\lambda: V \rightarrow L $は，アルファベット$L$からすべてのノード$v\in V$ にラベルを割り当てる関数を表します。この関数λについては後述します。

■ Doc2vecに着想

Graph2vecの中核を担うアルゴリズムはDoc2vec [2] という手法に着想を得ています。

Doc2vecの原著 [2] より抜粋

Doc2vecでは単語の集合である文書を入力とし、skip-gramを拡張したモデルを用いて文書の分散表現を学習します。この際、文書中に出現する単語を予測するよう学習しつつ、入力層と中間層を更新します。

Graph2vecでは同様に、グラフ全体を「文書」と見なし、グラフの各ノード周辺の根付き部分グラフを「単語」と見なすことで特徴ベクトルを学習します。類似文書で共通の単語が共起するように、類似グラフでは共通の根付き部分グラフを持つことを学習に組み込みます。

根付き部分グラフとは

さて、「根付き部分グラフ」とは何でしょうか？

木 (数学) のwikiより抜粋

特定の一つのノードを特別視したグラフ構造を根付きグラフ（rooted graph）と定義します。Graph2vecでは、事前に最大深さ$H$を設定し、グラフ$G_i$内の全てのノード$n$に対して根付き部分グラフ$sg_n^d$を抽出します。ここで添え字の$d$は$0 \leq d \leq H$の深さを表しており、各ノードに対して$H+1$個の部分グラフを集合として獲得します。

根付き部分グラフのリラベルと識別IDマッピング

上述した根付き部分グラフを集合として扱う際には、部分グラフに一意なラベルを割り当て、数値に変換する方が扱いやすいです。
そこで本モデルでは、グラフの同型性判定でも頻用される、Weisfeiler-Lehman Kernel (WL) [3] を用いたノードのリラベル手法を採用しています。

WL Kernel [3] の原著より抜粋; 1回のリラベル操作の図

詳細は割愛しますが、WL Kernelを用いてリラベルと識別IDへのマッピングを行い、多重集合を作成します。これにより、大元のグラフ（文書）を構成する根付き部分グラフ（単語）集合と見なすことができ、Doc2vecの応用が可能です。

WL Kernelの論文解説は以下のブログなどが参考になりました。Pythonでの実装確認も含めていつかまとめようと思います。

学習プロセス

各グラフのIDを one-hot ベクトルの形でモデルに入力します。出力層では、サブグラフ$sg_j$がグラフ$G_i$から抽出されたという、条件付き確率分布を出力します。この際に、以下の尤度を最大化するように学習を更新します。

$\sum_{j=1}^{n_i(H+1)} log \frac{exp(\vec{G_{i}} \cdot \vec{sg_i})}{\sum_{sg\in Voc} exp(\vec{G_{i}} \cdot \vec{sg_i})}$

かなり端折ったので、導出は原著を参照してください。

学習完了後には、類似の特徴ベクトルに写像されるグラフ同士は多くの共通の根付き部分グラフを持ち、潜在空間内の近くに位置します。

Graph2vecをPythonで実行する

karateclubというライブラリーを使用します。

データの準備

今回は、タンパクの立体構造から酵素活性の有無を予測するベンチマークデータを用いて検証してみます。

import networkx as nx
from torch_geometric.datasets import TUDataset

# Load protein dataset
dataset = TUDataset('.',name='PROTEINS', use_node_attr=True)

graphs = []
labels = []
for data in dataset:
    # Create Networkx.classes
    e_list = []
    tensor_edgelist = data['edge_index']
    for i in range(len(tensor_edgelist[0])):
        e_list.append((int(tensor_edgelist[0][i]), int(tensor_edgelist[1][i])))
    g = nx.from_edgelist(e_list)

    # Load features (amino acids)
    x = data['x']
    #nx.set_node_attributes(g, {j: x[j] for j in range(g.number_of_nodes())}, "feature")
    nx.set_node_attributes(g, {j: str(j) for j in range(g.number_of_nodes())}, "feature")

    # Checking the consecutive numeric indexing.
    node_indices = sorted([node for node in g.nodes()])
    numeric_indices = [index for index in range(g.number_of_nodes())]

    if numeric_indices == node_indices:
        graphs.append(g)
        labels.append(int(data["y"]))
    else:
        pass

最終的に1108のタンパク質を解析対象としました。アミノ酸配列由来のノード特徴量を使用することもできますが(上記のコメントアウトの箇所)、今回は簡単のためノードのIDをそのままノード特徴量として使用しています。

Graph2vecの実行

早速、Graph2vecを実行します。各グラフについて256次元の特徴ベクトルを取得し、ヒートマップで可視化しました。

import seaborn as sns
import matplotlib.pyplot as plt
from karateclub.graph_embedding import Graph2Vec

model = Graph2Vec(wl_iterations=2, use_node_attribute="feature", dimensions=256, 
                  down_sampling=0.0001, epochs=100, learning_rate=0.025, min_count=10)
model.fit(graphs)
emb = model.get_embedding() # (1108, 128)

sns.clustermap(emb)
plt.show()

はっきりとしたクラスターが得られると期待していたのですが、得られた特徴量は比較的均質でした。

各グラフの相関

特徴ベクトルに基づき、グラフ間のピアソンの相関をヒートマップで可視化しました。

import pandas as pd

emb_corr = pd.DataFrame(emb.T).corr()
sns.heatmap(emb_corr)
plt.show()

やはり、比較的に均質な特徴量であったため、多くのグラフ間で高い相関を示しています。特異な部分構造を反映したベクトル表現を獲得できているのか少し怪しいです。

似ている・似ていないグラフを可視化

上のセクションで示したように、グラフ間の類似度を特徴量の相関という観点で評価できます。そこで、ID=0のグラフに対して、類似度の高いグラフと低いグラフをそれぞれ描画してみます。

● グラフID=0と類似度の高いグラフ: ID=49 (Corr=0.993)

● グラフID=0と類似度の低いグラフ: ID=620 (Corr=0.151)

ノード数や細長い外観が類似しているID=49と高い相関を示し、グラフのスケールが異なるID=620とは低い相関を示した結果は直観とも合致します。

以上の結果から、Graph2vec自体は機能しており、意味のある特徴量を獲得出来ていることが分かります。

二値分類タスクに応用可能か？

得られた特徴量をUMAPで可視化して、タンパク質の酵素活性の有無を分類できるか試してみます。

import numpy as np
import umap
import umap.plot

# Fit UMAP
mapper = umap.UMAP(random_state=0,n_neighbors=10, min_dist=0.5, n_components=2)
mapper.fit(emb)

# Visualize with umap.plot
umap.plot.points(mapper,labels=np.array(labels))
plt.show()

ラベルの0と1がそれぞれ酵素活性の有無を表しています。左下に0ラベルを濃縮するなど、部分的には層別化できていそうですが、二値分類タスクとしては、イマイチな印象です。

今回の検証では、Graph2vecを用いて取得した各タンパク質の特徴ベクトルは、酵素活性の有無を予測するには不十分な情報量であると示唆されます。

おわりに

Graph2vecという単語を見た時、GNNのreadoutのように全頂点の特徴量の総和を返すような手法かと思ったのですが、文章（グラフ）を構成する単語（サブグラフ）と見なし、Doc2vecの問題設定へと落とし込む点に感心しました。

グラフの同型性の判定もそうですが、教師なしでグラフの特徴ベクトルを要求する場面は多いと思うので、今後も積極的に活用していこうと思います。一方で、今回試したようなラベル情報が与えられた分類タスクでは、GNNの学習を通じてノード特徴量を更新し、readoutで集約すれば良い気もします。

本記事の作成にあたり、Gl2vec [4] の著者でもあるChen氏の修士論文の和文要旨は非常に参考になりました。Gl2vecも本領域で重要な立ち位置を築いているため、Graph2vecと併せて読むことをお勧めします。

コードはGitHubで公開しています。パラメータを少し弄ったくらいなので、チューニングすると分類タスクの性能は多少上がると思います。

参考文献

[1] Narayanan, Annamalai, et al. "graph2vec: Learning distributed representations of graphs." arXiv preprint arXiv:1707.05005 (2017).
[2] Le, Quoc, and Tomas Mikolov. "Distributed representations of sentences and documents." International conference on machine learning. PMLR, 2014.
[3] Shervashidze, Nino, et al. "Weisfeiler-lehman graph kernels." Journal of Machine Learning Research 12.9 (2011).
[4] Chen, Hong, and Hisashi Koga. "Gl2vec: Graph embedding enriched by line graphs with edge features." Neural Information Processing: 26th International Conference, ICONIP 2019, Sydney, NSW, Australia, December 12–15, 2019, Proceedings, Part III 26. Springer International Publishing, 2019.