VAEに基づくグラフの表現学習

こんにちは。年末年始はファイアーエムブレムで可処分時間が消えてしまった DSOC R&D Group の橋本です。
さて、今回の記事では、変分オートエンコーダ (Variational Auto-Encoder, VAE) [1]に基づくグラフの表現学習について紹介したいと思います。
近年グラフに対する深層学習手法の発展が目覚ましく、応用先としては材料科学（分子や結晶をグラフと見る）*1や、ソーシャルネットワークなどが挙げられます。ここではソーシャルネットワークのようなグラフに対するものに限定して、ノードの潜在表現を得る手法を紹介します。*2

Variational Graph Auto-Encoder

変分グラフオートエンコーダ (Variational Graph Auto-Encoder, VGAE) とは、VAEにおけるencoderの部分にグラフ畳み込みネットワーク (Graph Convolutional Networks, GCN) [2]を使うことにより、グラフデータに対しても潜在表現の学習が行える手法です[3]。*3 これまでVAEをベースとして、encoderやdecoderに種々のニューラルネットワークのアーキテクチャが組み込まれたり、他の確率モデルと組み合わせるといった形で活用されてきましたが、本手法ではencoderにGCNを導入するだけでノードの低次元潜在変数を得ることができます。

f:id:w-hash52:20191224001049p:plain:w400 — Cora citation networkの潜在表現。ノードが論文、色が論文のラベル、エッジが引用を表す。[3] Figure 1より引用。

今回対象とする非有向グラフを $\mathcal {G} = ( \mathcal {V} , \mathcal {E} )$ とします。ここで、 $\mathcal {V}$ はノード集合、 $\mathcal {E}$ はエッジの集合を示します。ノード数は $N = | \mathcal {V} |$ とし、隣接行列を $\boldsymbol { \rm A}$ とします。ノードの潜在変数を $\boldsymbol { \rm z }_{i}$ とし、潜在変数の次元を $F$ とし、 $\boldsymbol { \rm Z}$ ( $N \times F$ )を潜在変数全体とします。また、入力にはノードの特徴も用います。ノードの特徴量の次元を $D$ とすると、特徴全体は $N \times D$ である $\boldsymbol { \rm X}$ で表現されます。

まずは、 $\boldsymbol { \rm X}$ と $\boldsymbol { \rm A}$ を用いたノードの潜在変数の推論を考えてみます。

\[
q( \boldsymbol { \rm Z} | \boldsymbol { \rm X}, \boldsymbol { \rm A}) = \prod_{i=1}^N \mathcal {N} (\boldsymbol { \rm z }_{i} | \boldsymbol { \mu }_{i}, \rm {diag} ( \boldsymbol { \rm \sigma }_{i}))\tag{1}
\]

ここで潜在変数のパラメータである $\boldsymbol { \mu }_{i}$ と $\boldsymbol { \rm \sigma }_{i}$ はGCNから学習します。

潜在変数から再構築した隣接行列は、潜在変数の内積をsigmoid関数にかけることで生成できます。

\[
p( \boldsymbol { \rm A} | \boldsymbol { \rm Z}) = \prod_{i=1}^N \prod_{j=1}^N p (A_{ij} | \boldsymbol { \rm z }_{i}, \boldsymbol { \rm z }_{j})\tag{2}
\]

ここで、 $p(A_{ij} = 1 | \boldsymbol { \rm z }_{i}, \boldsymbol { \rm z }_{j}) = \sigma ( \boldsymbol { \rm z }_{i}^{\rm T} \boldsymbol { \rm z }_{j})$ であり、 $\sigma (\cdot)$ はsigmoid関数です。

学習は、VAEと同様に変分下限 $\mathcal {L}$ を、encoderにおける変分パラメータ (ここではencoderのGCNのパラメータを指す)に対して最大化させることで行います。

\[
\mathcal {L} = \mathbb {E}_{q( \boldsymbol { \rm Z} | \boldsymbol { \rm X}, \boldsymbol { \rm A}) } [{\rm log} p(\boldsymbol { \rm A} | \boldsymbol { \rm Z})] - {\rm KL}[q( \boldsymbol { \rm Z} | \boldsymbol { \rm X}, \boldsymbol { \rm A}) || p( \boldsymbol { \rm Z}))]\tag{3}
\]

ここで ${\rm KL} [q || p$ ] は分布 $q$ と $p$ の近さを表す尺度であるKLダイバージェンス (Kullback-Leibler divergence) であり、 $p( \boldsymbol { \rm Z } ) = \prod_{i=1}^N \mathcal {N} ( \boldsymbol { \rm z }_{i} | 0, \boldsymbol { \rm I })$ は潜在変数の事前分布です。通常のVAEではdecoderにも適当なニューラルネットワークを置き、そのパラメータに対しても同時に最適化させていますが、VGAEについてはdecoderにパラメータがないことに注意してください。VAEと同様に、再構築したデータ（ここでは隣接行列）の尤度が高くなるように学習しつつ、一方で第2項のKLダイバージェンスにおいて潜在変数の事後分布が事前分布から離れすぎないように正則化をかけていることがわかります。

Graphite

VGAEでは潜在表現からグラフを生成する場合に、潜在表現の内積をsigmoid関数にかけることで再構築された隣接行列を生成していましたが、decoderにもニューラルネットを使うことで表現力が向上し、潜在変数から隣接行列を再構築した際の尤度が高くなると考えられます。Graphiteでは、潜在表現からdecodeする際にもニューラルネットを使うことで、VGAEのdecoderにおける制約を取り除いています[4]。
ニューラルネットを用いて潜在変数からグラフを生成するには、まず隣接行列の"種"を生成する必要があります。そのために、VGAEと同様にして潜在変数の内積を使って定義します。

\[
\hat { \boldsymbol { \rm A}} = \frac{ \boldsymbol { \rm Z}^{\rm T} \boldsymbol { \rm Z} } {|| \boldsymbol { \rm Z} ||} + \boldsymbol { \rm I} \tag{4}
\]

内積を使っている点ではVGAEのdecoderと同様ですが、種となるグラフに明示的に非負性を保証するため、単位行列を足しています。
種である隣接行列 $\hat { \boldsymbol { \rm A}}$ を生成した後は、以下の式によって特徴行列を更新します。

\[
\boldsymbol { \rm Z}^{*} = {\rm GNN} (\hat { \boldsymbol { \rm A}}, [ \boldsymbol { \rm Z} | \boldsymbol { \rm X} ]) \tag{5}
\]

ここで $[ \boldsymbol { \rm Z} | \boldsymbol { \rm X}$ ] は $\boldsymbol { \rm Z}$ と $\boldsymbol { \rm X}$ の結合を意味します。decoderの層の数だけ更新した後、Eq.(2) とほぼ同様に、特徴 $\boldsymbol { \rm Z}^{*}$ の内積にsigmoid関数をかけることで、潜在変数 $\boldsymbol { \rm Z}$ と入力した特徴量 $\boldsymbol { \rm X}$ がある場合の条件付き確率 $p( \boldsymbol { \rm A} | \boldsymbol { \rm Z}, \boldsymbol { \rm X} )$ が得られます（式はEq.(2) において $\boldsymbol { \rm z }_{i}$ 、 $\boldsymbol { \rm z }_{j}$ を $\boldsymbol { \rm Z}^{*}$ と入れ替えれば同様なので略）。

Deep Generative Latent Feature Rational Model

VGAEやGraphiteではGCNを用いてグラフデータにおける潜在表現の学習は行えますが、潜在変数自体の解釈性は高くありません。そこでDeep Generative Latent Feature Rational Model (DGLFRM) では、混合メンバーシップ確率的ブロックモデル (Mixed-Membership Stochastic Blockmodels, MMSB)[5]と呼ばれる潜在変数モデルをGraphiteに取り入れることにより、潜在表現の解釈性を高めています[6]。具体的には、各ノードにおけるコミュニティの混合比率がわかります。さらに、そのノードにおけるコミュニティ数自体を、棒折り過程 (stick-breaking process, SBP) を用いることで推定できます。説明の都合上、再構築した隣接行列 $\boldsymbol { \rm {\hat A}}$ を生成するdecoderについて、先に説明したいと思います。この論文では潜在変数全体の次元を $N \times K$ としていますので、コミュニティ数を $K$ として説明していきます。

f:id:w-hash52:20191224174224p:plain:w400 — decoderのグラフィカルモデル。 [6] Figure 1の左図から引用。

ノード $n$ における潜在変数（ノード埋め込み）が $\boldsymbol z_n = \boldsymbol b_n \odot \boldsymbol r_n$ のように $\boldsymbol b_n$ と $\boldsymbol r_n$ の要素積で書ける、とします。ここで $\boldsymbol b_n$ はK次元のバイナリベクトルであり、 $\boldsymbol r_n$ は $K$ 次元の実数ベクトルであり、事前分布として $p( \boldsymbol r_n ) = \mathcal{N} ( \boldsymbol r_n | 0, \sigma^2 \boldsymbol {\rm I})$ のようにガウス分布が仮定されているとします。また、SBPによる $b_n$ のモデリングは次のようにして書けます。

\begin{align}
v_k & \sim \rm {Beta} ( \alpha, 1) \tag{5} \\\
\pi_k & = \prod_{i=j}^k v_k \tag{6} \\\
b_{nk} & \sim \rm {Bernoulli} (\pi_k ) \tag{7} \\\
\end{align}

DGLFRMにおいてはSBPは有限次元の $K$ を用いているため、実際は打ち切り棒折り過程となります。ここで $\alpha$ は集中度パラメータとよばれるハイパーパラメータです。 $v_k$ は $\rm {Beta} (\alpha, 1)$ からサンプルされる確率変数ですが、Bernoulli分布のパラメータとなっている $\pi_k$ は $K$ について総和をとったとき必ず1となります。この保証のもとでは確率が割り当てられたコミュニティ以外は $\pi_k$ がほぼゼロに縮退するため、実質的にノードのサイズが決まることになります。
$\boldsymbol b_n$ および $\boldsymbol r_n$ が決まれば $\boldsymbol z_n = \boldsymbol b_n \odot \boldsymbol r_n$ に従って潜在変数を生成します。その後、多層ニューラルネットワークのような非線形関数を用いて $f_n$ に変換し、 $f_n$ 同士の内積をsigmoid関数に通すことで $p (A_{ij} | \boldsymbol { f }_{n}, \boldsymbol { f }_{m})$ が求まります。

f:id:w-hash52:20191224174037p:plain:w400 — DGLFRMにおけるencoderのグラフィカルモデル。[6] Figure 1の右図から引用。

次にencoderについて説明します。decoderで先に説明した ${ v_n, b_n, r_n}$ を推定するため、以下のように変分パラメータ ${c_k, d_k, \pi_k, \mu_k, \sigma_k}$ を用いて ${ v_n, b_n, r_n}$ を近似する変分分布を定義します。

\begin{align}
q(v_{nk} ) &= {\rm Beta} ( c_{nk}, d_{nk} ) \tag{8} \\\
q(v_{nk} ) &= {\rm Bernoulli} ( \pi_k) \tag{9} \\\
q(\boldsymbol {r}_n) &= \mathcal {N} ( \boldsymbol {r}_n | \boldsymbol {\mu}_n, {\rm diag} ( \boldsymbol {\sigma}_n)) \tag{10} \\\
\end{align}

以上の変分分布に出現している変分パラメータ ${c_k, d_k, \pi_k, \mu_k, \sigma_k}$ はVGAEと同様の枠組みでGCN( $\boldsymbol { \rm A}, \boldsymbol { \rm X}$ ) から得られ、VAEと同様に変分下限を最大化することで学習されます。*4

\begin{align}
\mathcal{L} &= \sum_{n=1}^N \Bigl( \mathbb {E}_q [ {\rm log} p(X_n | \boldsymbol z_n)] \Bigr) \tag{11}\\\
& + \sum_{n=1}^N \sum_{m=1}^N \Bigl( \mathbb {E}_q [ {\rm log} p(A_{nm} | \boldsymbol z_n, \boldsymbol z_m)] \Bigr) \tag {12}\\\
& - \sum_{n=1}^N \Bigl({\rm KL} [q( \boldsymbol b_n | \boldsymbol v_n) || p(\boldsymbol b_n | \boldsymbol v_n)] + {\rm KL} [q( \boldsymbol r_n) || p( \boldsymbol r_n)] \\\
&+ {\rm KL} [q( \boldsymbol v_n) || p( \boldsymbol v_n )] \Bigr) \tag {13}
\end{align}

最後に、具体的な潜在表現を見てみましょう。

f:id:w-hash52:20191226172342p:plain — （左）DGLFRM （右）VGAE の潜在表現。 $K=10$ で固定。[6] Figure 2から引用。

（左）DGLFRM （右）VGAE の潜在表現。 $K=10$ で固定。[6] Figure 2から引用。

上の図は100個のノードと10個のコミュニティを持つグラフの合成データに対して、DGLFRMとVGAEを適用したものです。右のVGAEではコミュニティ数 $K=10$ の各次元に値が割り振られていますが、左のDGLFRMでは該当するコミュニティ以外はほぼゼロになっており、そのノードがどのコミュニティにどれだけの比率で属しているか、という潜在表現の解釈性が向上しています。

最後に

本記事では、ソーシャルネットワーク向けにVAE（の派生形）を応用した表現学習についてまとめ、以下の3つを紹介しました。

VAEにおけるencoderにGCNを使ったVGAE
VGAEのdecoderにもニューラルネットワークを使えるようにしたGraphite
潜在変数の解釈性を高める DGLFRM

今回は時間の都合上実データから具体的に潜在表現を得るような実験はできなかったので、今後実験していければと思います。

参考文献

[1]: Diederik P Kingma, Max Welling, "Auto-Encoding Variational Bayes", in ICLR, 2014.
[1312.6114] Auto-Encoding Variational Bayes

[2]: Thomas N. Kipf, Max Welling, "Semi-Supervised Classification with Graph Convolutional Networks", in ICLR, 2017. Semi-Supervised Classification with Graph Convolutional Networks

[3]: Thomas N. Kipf, Max Welling, "Variational Graph Auto-Encoders", in NIPS Workshop "Bayesian Deep Learning", 2016. Graphite: Iterative Generative Modeling of Graphs

[4]: Aditya Grover, Aaron Zweig, Stefano Ermon, "Graphite: Iterative Generative Modeling of Graphs", in ICML, 2019. [1611.07308] Variational Graph Auto-Encoders

[5]: Edoardo M Airoldi, David M Blei, Stephen E Fienberg, Eric P Xing, "Mixed Membership Stochastic Blockmodels", in NIPS, 2007.
Mixed Membership Stochastic Blockmodels

[6]: Nikhil Mehta, Lawrence Carin, Piyush Rai, "Stochastic Blockmodels meet Graph Neural Networks", in ICML 2019.
Stochastic Blockmodels meet Graph Neural Networks

[7]: Jiaxuan You, Bowen Liu, Rex Ying, Vijay Pande, Jure Leskovec, "Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation", in NeurIPS, 2018.
[1806.02473] Graph Convolutional Policy Network for Goal-Directed Molecular Graph Generation

[8]: Hiroshi Kajino, "Molecular Hypergraph Grammar with Its Application to Molecular Optimization", in ICML, 2019.
Molecular Hypergraph Grammar with Its Application to Molecular Optimization

[9]: Eric Nalisnick, Padhraic Smyth, "Stick-Breaking Variational Autoencoders", in ICLR, 2017.
[1605.06197] Stick-Breaking Variational Autoencoders

[10]:Chris J. Maddison, Andriy Mnih, Yee Whye Teh, "The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables", in ICLR, 2017.
[1611.00712] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables

buildersbox.corp-sansan.com

*1: 私はもともと化学を専攻していたというのもあり、材料科学に対するグラフ機械学習の手法にも興味があったりします

*2: VAEを用いたグラフの表現学習において、応用ドメイン先によって発展のアプローチが異なるようです。大規模なソーシャルネットワークについては本記事で紹介するようにノード埋め込みの解釈性向上についての研究が多いです。一方で例えばVAEを材料科学に応用する場合、新規材料の候補を生成することが一つの目的となります。それを達成するためにはできるだけ化合物が化学的に正しいように生成されることが重要となるため、いかに化合物の原子価の規則を守ったまま分子をdecodeするかについての研究などが主に進んでいます[7, 8]。

*3: GraphVAE: Towards Generation of Small Graphs Using Variational Autoencoders という非常に似た名前の論文もありますが、こちらは論文名のとおり小分子を対象にしたVAEであるため、ご注意ください。

*4:VAEやVGAEでは潜在変数がガウス分布だったため、[1]のようなreparameterization trickにより学習できました。しかしDGLFRMでは潜在変数は単純なガウス分布ではなく、Beta分布やBernoulli分布を用いているため、学習時におけるreparameterization trickにはもう少し工夫が要ります。詳しくは[9, 10] を御覧ください。