GNN Part 1 - 教師なし頂点表現学習問題

Date: 2024-04-30 | Category: gnn

本記事からしばらくの間, 「グラフニューラルネットワーク」佐藤竜馬著を教科書としてGNNの基礎について学んでいきます.

本記事では, 「グラフニューラルネットワーク」Chapter2.3 古典的なグラフ機械学習手法において紹介されている教師なし頂点表現学習問題を, JAXを用いて実装してみます. なお, 今回実装したコードはGitHubにて公開しています.

教師なし頂点表現学習問題
最適化問題の求め方
- 勾配降下法による求解
- 行列分解による求解
結果
- 埋め込み次元d=5の場合
- 埋め込み次元dを変えてみる
まとめ
参考文献

教師なし頂点表現学習問題

教師なし頂点表現学習問題は以下のような問題です.

入力: 重み付きグラフ $G = (V, E, w)$ ($V$は頂点集合, $E$は辺集合, $w$は辺の重み)
出力: グラフ構造を考慮した頂点の埋め込み $Z \in \mathbb{R}^{n \times d}$ ($n$は頂点数, $d$は埋め込み次元)

この問題は, グラフ構造を考慮した頂点の埋め込みを学習する問題です.

埋め込みとはデータの情報を低次元のベクトルに変換したもので, 機械学習においてデータの特徴を扱いやすいベクトルとして抽出するために用いられます. 自然言語処理においては単語の埋め込み(Word Embedding)が有名です. 今回はそのグラフ版です.

この問題の出力として得られる$Z$は, 各行$Z_i$が頂点$i$の埋め込みベクトルとなります. よって, $Z$を求めることで, グラフデータを, 扱いやすいベクトルデータに変換し, 機械学習の問題に持ってくることができます.

ただ, “グラフ構造を考慮した”という部分が曖昧です. ここでは, 教師なし頂点表現学習問題を以下の最適化問題として定式化することで, グラフ構造を考慮した頂点の埋め込みを学習します.

\[\min_{Z\in \mathbb{R}^{n \times d}} \frac{1}{2} \sum_{u\in V} \sum_{v\in V}((D_{uv} + W_{uv})-Z_u^T Z_v)^2\]

ここで, $D \in \mathbb{R}^{n \times n}$は次数行列, $W \in \mathbb{R}^{n \times n}$は隣接行列です. 次数行列は各頂点の次数を対角成分に持つ行列で, 隣接行列は隣接する頂点間の重みを持つ行列です.

この最適化問題は, $Z_u^T Z_v$が$D_{uv} + W_{uv}$に近づくように$Z$を学習する問題です. 各頂点$u$の埋め込みベクトルは$Z_u$であることから, この最適化問題は, 隣接している頂点同士の埋め込みベクトルの内積が大きくなるように設計されています.

最適化問題の求め方

教師なし頂点表現学習問題の最適化問題は, 「グラフニューラルネットワーク」では行列分解を用いて解く方法が紹介されています.

今回は, JAXを用いて, 勾配降下法及び行列分解で最適化問題を解くコードを装してみます.

具体例として, 以下のグラフを用います.

graph

このグラフの次数行列$D$と隣接行列$W$は以下の通りです.

D = jnp.diag(jnp.array([2, 4, 2, 1, 1]))
print("D", D)
# [[2 0 0 0 0]
# [0 4 0 0 0]
# [0 0 2 0 0]
# [0 0 0 1 0]
# [0 0 0 0 1]]

# adjacency matrix
W = jnp.array([
    [0, 1, 1, 0, 0],
    [1, 0, 1, 1, 1],
    [1, 1, 0, 0, 0],
    [0, 1, 0, 0, 0],
    [0, 1, 0, 0, 0]
])

この時の正解データである, $D + W$は以下のような行列になります. D+W

$Z$は, $Z@Z.T$が$D+W$に近づくように学習されます.

勾配降下法による求解

勾配降下法では, パラメータ$Z$をランダムに初期化し, $D+W$と$Z@Z.T$の差を最小化するようにパラメータ$Z$を更新していきます.

まず, 最適化問題を解くための損失関数を定義します.

def loss_fn(Z):
    loss = 0
    #愚直な方法
    #for u in range(5):
    #    for v in range(5):
    #        loss += ((D[u,v] + W[u,v]) - Z[u].T @ Z[v])**2
    # 行列演算を用いた方法
    loss = jnp.sum((D + W - Z @ Z.T)**2)

    return loss

続いて, 勾配降下法のupdate関数を定義します.

def update(theta, lr=0.01):
    loss, grad = jax.value_and_grad(loss_fn)(theta)
    return theta - lr * grad, loss

最後に, 勾配降下法で最適化問題を解きます.

Z = jax.random.uniform(key = jax.random.PRNGKey(0), shape=(D.shape[0], d))
for epoch in range(1000):
    Z, loss = update(Z)
    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss {loss}")

行列分解による求解

「グラフニューラルネットワーク」でも紹介されているように, 今回の最適化問題は行列分解によって最適解が求められます. 実装は以下のようにして行えます.

e, V = jnp.linalg.eigh(D + W)
Z = jnp.zeros((D.shape[0], d))
for i in range(min(d, D.shape[0])):
    Z = Z.at[:,i].set(V[:, -1 - i] * jnp.sqrt(e[-1 - i]))