SNSの「知り合いかも」機能は、まだ友達になっていない2人のユーザーが「将来つながる可能性」をどう予測しているのでしょうか。あるいは、創薬の現場では、まだ実験されていないタンパク質とタンパク質の間に「相互作用がありそうか」をどう見積もっているのでしょうか。これらはどちらも、グラフの中でまだ存在しないエッジ(辺)を予測するという同じ問題に帰着します。これを リンク予測(link prediction) と呼びます。
図のように、VGAEはグラフ $(A, X)$ を受け取り、GCNエンコーダで各ノードをガウス分布として潜在空間に埋め込み、内積デコーダで元のグラフ構造を再構成するエンドツーエンドの枠組みです。学習の目標はELBO(変分下界)の最大化で、再構成損失とKL正則化項のバランスを取ります。
リンク予測を解くには、各ノード(頂点)を「意味のある座標」に埋め込み、その座標同士の近さからエッジの有無を判定すればよさそうです。しかし、グラフは画像や文章のように規則正しく並んだデータではありません。ノードごとに隣接の数も違えば、順番にも意味がありません。こうした不規則なグラフ構造から、ノードの良い埋め込み(潜在表現)を教師なしで学習する強力な手法が、本記事で扱う 変分グラフオートエンコーダ(Variational Graph Autoencoder, VGAE) です。2016年に Kipf と Welling が提案しました。
VGAEは、画像生成で有名な変分オートエンコーダ(VAE)の考え方を、グラフ畳み込みネットワーク(GCN)と組み合わせてグラフに持ち込んだものです。学習したノード埋め込みは次のような場面で活躍します。
- リンク予測: 引用ネットワークでの新規引用予測、推薦システムでのユーザー・アイテム間リンク予測、創薬での分子間相互作用予測。
- ノードクラスタリング・可視化: 潜在空間上で近いノードを同じグループとみなし、コミュニティ抽出や次元削減後の可視化に使う。
この記事では、VGAEのエンコーダとデコーダを定義し、その学習目的である ELBO(変分下界) を確率論の基礎から省略せずに導出します。各項(再構成項とKL正則化項)が何を意味するかを直感的に説明したうえで、再パラメータ化トリックを使ったPyTorch実装を示し、合成グラフでエッジを学習用とテスト用に分割してリンク予測のROC-AUC・APを評価、最後に潜在空間を可視化します。
本記事の内容
- VGAEのエンコーダ $q(\bm{Z}|\bm{X},\bm{A})$ と内積デコーダ $p(\bm{A}|\bm{Z})$ の直感と定義
- 変分下界 ELBO = 再構成項 − KL項 の省略なしの導出
- 再パラメータ化トリックを使ったPyTorchでのスクラッチ実装
- 合成グラフでのリンク予測評価(ROC-AUC / AP)と潜在空間の可視化
前提知識
この記事を読む前に、以下の記事を読んでおくと理解が深まります。
- グラフニューラルネットワーク(GNN)とは — グラフ上で情報を伝播させる基本的な枠組み
- グラフ畳み込みネットワーク(GCN) — VGAEのエンコーダの中核となる層
- グラフラプラシアン — GCNの正規化隣接行列の背景にある理論
- グラフ表現学習 — ノード埋め込みの考え方全般
また、変分オートエンコーダ(VAE)・KLダイバージェンス・多変量正規分布の基本を知っているとスムーズですが、本記事では必要な部分を逐次復習しながら進めます。
グラフとオートエンコーダの記号を準備する
まず舞台となるグラフを数式で表しましょう。$N$ 個のノードを持つ無向グラフ $\mathcal{G}=(\mathcal{V},\mathcal{E})$ を考えます。グラフの構造は 隣接行列(adjacency matrix) $\bm{A}\in\{0,1\}^{N\times N}$ で表現します。ノード $i$ とノード $j$ の間にエッジがあれば $A_{ij}=1$、なければ $A_{ij}=0$ です。無向グラフなので $\bm{A}$ は対称($A_{ij}=A_{ji}$)です。
左のグラフでは5つのノードが7本のエッジで結ばれており、右の隣接行列ではエッジが存在する箇所(青セル)に1、存在しない箇所に0が並んでいます。グラフが対称構造(無向)であるため、行列も対角線を中心に対称になっていることが一目で確認できます。このように、グラフの「繋がり方」という構造情報が行列の形で表現されるのです。
各ノードは特徴ベクトルを持つことができます。これを行方向に並べた 特徴行列(feature matrix) を $\bm{X}\in\mathbb{R}^{N\times D}$ とします。$i$ 行目 $\bm{x}_i\in\mathbb{R}^D$ がノード $i$ の特徴です(特徴がない場合は単位行列 $\bm{X}=\bm{I}_N$、すなわち one-hot を使います)。
VGAEのゴールは、各ノード $i$ を $F$ 次元($F\ll N$)の 潜在ベクトル $\bm{z}_i\in\mathbb{R}^F$ に埋め込むことです。これを並べた行列を $\bm{Z}\in\mathbb{R}^{N\times F}$ と書きます。良い埋め込みとは、「グラフ上でつながっているノード同士は潜在空間でも近い」ような埋め込みです。
ここでオートエンコーダの発想を思い出しましょう。通常のオートエンコーダは、入力 $\bm{x}$ をエンコーダで低次元のコード $\bm{z}$ に圧縮し、デコーダで $\bm{x}$ を復元します。「圧縮しても元に戻せる」ように学習することで、$\bm{z}$ が入力の本質を捉えるようになります。VGAEはこれをグラフに適用します。すなわち、入力 $(\bm{X},\bm{A})$ から潜在 $\bm{Z}$ をエンコードし、$\bm{Z}$ から隣接行列 $\bm{A}$ を復元するのです。ただし通常のオートエンコーダと違い、VGAEは $\bm{z}_i$ を1点ではなく確率分布として扱う点が「変分(variational)」たるゆえんです。
左の通常オートエンコーダでは各ノードが潜在空間の1点に固定されています。右のVGAEでは、各ノードが「中心と広がり」を持つ楕円状のガウス分布(雲)として表現されており、埋め込みの不確かさを自然に表現できています。この「雲」が事後分布 $q(\bm{z}_i|\bm{X},\bm{A})=\mathcal{N}(\bm{\mu}_i, \mathrm{diag}(\bm{\sigma}_i^2))$ に対応します。
では、エンコーダとデコーダを具体的にどう設計するのか、次節で直感から定義していきましょう。
VGAEのエンコーダとは — 潜在分布を推論する
通常のオートエンコーダは各ノードを1点の $\bm{z}_i$ に潰しますが、これだと「このノードがどこにあるか、どのくらい確信があるか」という不確かさを表現できません。確信が薄いノードまで1点に決め打ちすると、潜在空間が過学習でスカスカになりがちです。そこでVGAEは、ノード $i$ の埋め込みをガウス分布の雲として表します。雲の中心 $\bm{\mu}_i$ がもっともらしい位置、雲の広がり $\bm{\sigma}_i$ が不確かさです。
数式で書くと、エンコーダは観測 $(\bm{X},\bm{A})$ が与えられたときの潜在 $\bm{Z}$ の 近似事後分布 $q(\bm{Z}|\bm{X},\bm{A})$ を出力します。各ノードが独立で、各次元が独立な対角ガウスだと仮定すると
$$ \begin{equation} q(\bm{Z}|\bm{X},\bm{A}) = \prod_{i=1}^{N} q(\bm{z}_i|\bm{X},\bm{A}), \qquad q(\bm{z}_i|\bm{X},\bm{A}) = \mathcal{N}\!\left(\bm{z}_i \,\middle|\, \bm{\mu}_i,\ \mathrm{diag}(\bm{\sigma}_i^2)\right) \end{equation} $$
となります。ここで $\bm{\mu}_i$ と $\bm{\sigma}_i$ はノード $i$ ごとの平均ベクトルと標準偏差ベクトルです。「近似事後分布」と呼ぶのは、本当に求めたい真の事後分布 $p(\bm{Z}|\bm{X},\bm{A})$ が計算困難なので、扱いやすいガウス分布で近似しているからです($q$ の意味は導出の節で明確になります)。
問題は、平均 $\bm{\mu}_i$ と分散パラメータ $\bm{\sigma}_i$ をどうやって計算するかです。ここで 2層のGCN を使います。GCNは「自分と隣接ノードの特徴を、正規化された重みで混ぜ合わせる」操作を繰り返す層でした。隣接ノードの情報を取り込むことで、グラフ構造を反映した埋め込みが得られます。具体的には、正規化隣接行列を
$$ \begin{equation} \tilde{\bm{A}} = \tilde{\bm{D}}^{-\frac{1}{2}}\,(\bm{A}+\bm{I}_N)\,\tilde{\bm{D}}^{-\frac{1}{2}} \end{equation} $$
と定義します。$\bm{A}+\bm{I}_N$ は自己ループを足した隣接行列、$\tilde{\bm{D}}$ はその次数行列(対角成分 $\tilde{D}_{ii}=\sum_j (A_{ij}+I_{ij})$)です。両側から $\tilde{\bm{D}}^{-1/2}$ を掛けることで、次数の大きいノードの影響が暴れないように正規化しています。この $\tilde{\bm{A}}$ を使い、まず共有の第1層で中間表現を作ります。
$$ \begin{equation} \bm{H} = \mathrm{ReLU}\!\left(\tilde{\bm{A}}\,\bm{X}\,\bm{W}^{(0)}\right) \end{equation} $$
そして第2層を2本用意し、片方で平均、もう片方で対数分散を出力します。
$$ \begin{equation} \bm{\mu} = \tilde{\bm{A}}\,\bm{H}\,\bm{W}_\mu^{(1)}, \qquad \log \bm{\sigma}^2 = \tilde{\bm{A}}\,\bm{H}\,\bm{W}_\sigma^{(1)} \end{equation} $$
ここで $\bm{\mu},\log\bm{\sigma}^2\in\mathbb{R}^{N\times F}$ で、それぞれ $i$ 行目が $\bm{\mu}_i,\ \log\bm{\sigma}_i^2$ に対応します。第1層 $\bm{W}^{(0)}$ は平均と分散で共有し、最後だけ枝分かれするのがポイントです。なお分散そのものではなく対数分散 $\log\bm{\sigma}^2$ を出力するのは、分散が常に正であるという制約を自然に満たし(指数を取れば必ず正)、数値的にも安定するからです。
これでエンコーダ $q(\bm{Z}|\bm{X},\bm{A})$ が定義できました。
左の図では各ノードが自分の特徴だけを持っています。右の図では、ターゲットノード(赤)が周囲4つの隣接ノードの情報を正規化重みで取り込み、グラフ構造を反映した新しい表現を得ています。GCNはこの集約操作を複数層重ねることで、より広い近傍の情報まで捉えることができます。
次は、得られた潜在 $\bm{Z}$ からどうやってグラフ(隣接行列)を復元するのか、デコーダを見ていきましょう。
VGAEのデコーダとは — 内積でエッジを復元する
潜在空間に埋め込まれたノードたちから、どうやって「どのノード同士がつながっているか」を当てるのでしょうか。直感はとてもシンプルです。潜在空間で近い(向きが揃っている)2つのノードほど、つながっている確率が高いとみなすのです。ベクトルの近さを測る最も手軽な道具は内積です。$\bm{z}_i^\top \bm{z}_j$ が大きければ2つのベクトルは似た方向を向いており、エッジがありそうだ、というわけです。
ただし内積は $-\infty$ から $+\infty$ までの値を取るので、これをそのまま確率にはできません。そこでロジスティックシグモイド関数 $\sigma(x)=1/(1+e^{-x})$ を通して $[0,1]$ に押し込めます。こうしてデコーダは、エッジ $A_{ij}$ が存在する確率を
$$ \begin{equation} p(A_{ij}=1 \mid \bm{z}_i,\bm{z}_j) = \sigma\!\left(\bm{z}_i^\top \bm{z}_j\right) \end{equation} $$
と定義します。各エッジが潜在 $\bm{Z}$ のもとで独立だと仮定すると、隣接行列全体の生成確率は各エッジの積になります。
$$ \begin{equation} p(\bm{A}|\bm{Z}) = \prod_{i=1}^{N}\prod_{j=1}^{N} p(A_{ij}|\bm{z}_i,\bm{z}_j) = \prod_{i,j} \sigma(\bm{z}_i^\top\bm{z}_j)^{A_{ij}}\,\bigl(1-\sigma(\bm{z}_i^\top\bm{z}_j)\bigr)^{1-A_{ij}} \end{equation} $$
2行目はベルヌーイ分布の確率質量関数 $p^{A}(1-p)^{1-A}$ をエッジごとに掛け合わせたものです。$A_{ij}=1$ なら $\sigma(\cdot)$ が、$A_{ij}=0$ なら $1-\sigma(\cdot)$ が効きます。このデコーダには学習パラメータが一切ないことに注目してください。エッジの復元能力はすべて、エンコーダが作る潜在 $\bm{Z}$ の良し悪しに委ねられています。だからこそ、良いリンク予測のためには良い埋め込みを学ぶしかなく、結果として埋め込みにグラフ構造が凝縮されるのです。
行列の形でまとめて書けば、全エッジのリンク確率行列は $\sigma(\bm{Z}\bm{Z}^\top)$(シグモイドは要素ごと)で一気に計算できます。これが「内積デコーダ(inner product decoder)」と呼ばれる所以です。
左の図では、潜在空間で同じ方向を向く2ベクトル(青)は内積が大きく「エッジあり」と判定され、逆方向を向く2ベクトル(赤)は内積が小さく「エッジなし」と判定されます。右の図ではシグモイド関数 $\sigma(x)=1/(1+e^{-x})$ が内積を $[0,1]$ の確率に変換する様子を示しています。内積が正の領域(青色)ではエッジの確率が0.5を超え、負の領域ではエッジなしと判定されます。
さて、エンコーダとデコーダが揃いました。残る問題は「これらのパラメータ $\bm{W}^{(0)},\bm{W}_\mu^{(1)},\bm{W}_\sigma^{(1)}$ を、何を最大化(最小化)して学習するのか」です。ここで変分推論の核心である ELBO が登場します。次節でゼロから導出しましょう。
ELBOの導出 — なぜ再構成項とKL項に分かれるのか
学習の目標: 周辺尤度を上げたい
私たちが本当に最大化したいのは、観測されたグラフ $\bm{A}$ がモデルから生成される確率、すなわち 周辺尤度(対数エビデンス) $\log p(\bm{A}|\bm{X})$ です(以下、特徴 $\bm{X}$ への条件付けは表記を簡潔にするため省略します)。生成モデルでは潜在変数 $\bm{Z}$ を経由して $\bm{A}$ が作られると考えるので、$\bm{Z}$ について積分(周辺化)して
$$ \begin{equation} \log p(\bm{A}) = \log \int p(\bm{A}|\bm{Z})\,p(\bm{Z})\, d\bm{Z} \end{equation} $$
と書けます。ここで $p(\bm{Z})$ は潜在変数の 事前分布(prior) で、VGAEでは標準正規分布 $p(\bm{Z})=\prod_i \mathcal{N}(\bm{z}_i|\bm{0},\bm{I}_F)$ を使います。「特に情報がなければ原点まわりに集まっていてほしい」という事前の信念です。
ところがこの積分は、$\bm{Z}$ が $N\times F$ 次元の連続変数なので解析的に計算できません。そこで、扱いやすい近似分布 $q(\bm{Z}|\bm{X},\bm{A})$(前述のエンコーダ)を導入し、$\log p(\bm{A})$ の下界を構成します。これがELBOです。
イェンセンの不等式で下界を作る
積分の中に $q$ を掛けて割る、という古典的なトリックから始めます。$\int q(\bm{Z})d\bm{Z}=1$ なので等式は崩れません。
$$ \begin{equation} \log p(\bm{A}) = \log \int q(\bm{Z})\,\frac{p(\bm{A}|\bm{Z})\,p(\bm{Z})}{q(\bm{Z})}\, d\bm{Z} \end{equation} $$
ここで $q(\bm{Z})$ は $q(\bm{Z}|\bm{X},\bm{A})$ の略記です。右辺は「$\frac{p(\bm{A}|\bm{Z})p(\bm{Z})}{q(\bm{Z})}$ という量を、分布 $q$ のもとで期待値を取って対数をかぶせたもの」と読めます。つまり $\log \mathbb{E}_q[\,\cdot\,]$ の形です。
$\log$ は上に凸(concave)な関数なので、イェンセンの不等式 $\log \mathbb{E}[X]\ge \mathbb{E}[\log X]$ が使えます。これを適用して対数を期待値の中に入れると、
$$ \begin{equation} \log p(\bm{A}) \ \ge\ \mathbb{E}_{q(\bm{Z})}\!\left[\log \frac{p(\bm{A}|\bm{Z})\,p(\bm{Z})}{q(\bm{Z})}\right] \ \equiv\ \mathcal{L} \end{equation} $$
この右辺 $\mathcal{L}$ こそが ELBO(Evidence Lower BOund, 変分下界) です。等号は近似 $q$ が真の事後分布に一致したときに成り立ちます。$\mathcal{L}$ を最大化すれば、それに押し上げられて $\log p(\bm{A})$ も大きくなる、というのが学習の発想です。
ELBOを2つの項に分解する
$\mathcal{L}$ の中の対数を分母・分子に分けます。$\log\frac{p(\bm{A}|\bm{Z})p(\bm{Z})}{q(\bm{Z})}=\log p(\bm{A}|\bm{Z}) + \log p(\bm{Z}) – \log q(\bm{Z})$ なので、期待値の線形性を使って
$$ \begin{equation} \mathcal{L} = \mathbb{E}_{q(\bm{Z})}\bigl[\log p(\bm{A}|\bm{Z})\bigr] + \mathbb{E}_{q(\bm{Z})}\bigl[\log p(\bm{Z}) – \log q(\bm{Z})\bigr] \end{equation} $$
と書けます。第2項に注目すると、$\mathbb{E}_q[\log p(\bm{Z})-\log q(\bm{Z})] = -\mathbb{E}_q\bigl[\log\frac{q(\bm{Z})}{p(\bm{Z})}\bigr]$ です。そして $\mathbb{E}_q[\log\frac{q}{p}]$ はまさに KLダイバージェンス $\mathrm{KL}(q\|p)$ の定義そのものです。したがって
$$ \begin{equation} \boxed{\ \mathcal{L} = \underbrace{\mathbb{E}_{q(\bm{Z}|\bm{X},\bm{A})}\bigl[\log p(\bm{A}|\bm{Z})\bigr]}_{\text{再構成項}} \ -\ \underbrace{\mathrm{KL}\!\left(q(\bm{Z}|\bm{X},\bm{A})\,\middle\|\,p(\bm{Z})\right)}_{\text{KL正則化項}}\ } \end{equation} $$
という、教科書でおなじみの形に到達しました。導出はこれで完了です。学習では $\mathcal{L}$ を最大化するので、実装上は $-\mathcal{L}$ を損失として最小化します。
各項の意味を直感で理解する
2つの項が何を要求しているかを言葉にしておきましょう。
- 再構成項 $\mathbb{E}_q[\log p(\bm{A}|\bm{Z})]$: 「エンコーダが出した潜在 $\bm{Z}$ からデコーダで隣接行列を復元したとき、元の $\bm{A}$ をどれだけうまく再現できるか」を測ります。これが大きいほど、埋め込みがグラフ構造を正しく捉えていることになります。後で見るように、この項はエッジごとの二値交差エントロピーにほかなりません。
- KL正則化項 $\mathrm{KL}(q\|p)$: 「近似事後分布 $q$ が、事前分布 $p=\mathcal{N}(\bm{0},\bm{I})$ からどれだけ離れているか」を測る罰則です。引き算されているので、$q$ が事前分布から離れすぎるとELBOが下がります。これにより潜在空間が原点まわりに整い、各ノードが勝手な遠くへ飛んでいくのを防ぎます。過学習を抑える正則化として働くわけです。
つまりVGAEは「グラフをよく再構成せよ(第1項)。ただし潜在空間は整然と保て(第2項)」という2つの要求のバランスを取って学習します。これはVAEとまったく同じ構造です。
左の図では、事前分布(青)に近い近似事後分布(緑)はKLが小さく罰則が軽く、遠い分布(赤)はKLが大きく強く引き戻されることを示しています。右の図では、元グラフを潜在空間を経由して再構成する過程で、誤ったエッジ(赤点線)が生じると再構成損失(BCE)が増加することを示しています。この2項のバランスが、グラフを再現しながら潜在空間を整頓する正則化効果を生み出します。
これで何を最適化すべきかは分かりました。しかし1つ大きな技術的障壁が残っています。再構成項は $q$ からのサンプリングを含む期待値であり、サンプリングは微分できません。これをどう乗り越えるかが、次節の再パラメータ化トリックです。
再パラメータ化トリックと損失関数の具体形
サンプリングは微分できない、という問題
ELBOを勾配降下で最大化するには、エンコーダのパラメータ $\bm{\phi}=\{\bm{W}^{(0)},\bm{W}_\mu^{(1)},\bm{W}_\sigma^{(1)}\}$ に関する勾配が必要です。ところが再構成項は $\bm{Z}\sim q_{\bm{\phi}}(\bm{Z}|\bm{X},\bm{A})$ からのサンプルに依存しています。「分布 $q_{\bm{\phi}}$ から乱数を引く」という操作そのものはパラメータ $\bm{\phi}$ に依存しているのに、引いた乱数を通して $\bm{\phi}$ に勾配を逆伝播することができません。乱数生成は微分不可能なブラックボックスだからです。
確率性を外に追い出す
ここで効くのが 再パラメータ化トリック(reparameterization trick) です。アイデアは「ランダムさを、パラメータに依存しない補助ノイズに押し出す」ことです。標準正規分布から $\bm{\varepsilon}_i\sim\mathcal{N}(\bm{0},\bm{I}_F)$ を引いておき、
$$ \begin{equation} \bm{z}_i = \bm{\mu}_i + \bm{\sigma}_i \odot \bm{\varepsilon}_i \end{equation} $$
と変換します($\odot$ は要素ごとの積)。こうして作った $\bm{z}_i$ は、平均 $\bm{\mu}_i$・分散 $\mathrm{diag}(\bm{\sigma}_i^2)$ のガウス分布に正しく従います(平均を足し、標準偏差倍するアフィン変換でガウス性は保たれる)。重要なのは、確率的な部分が $\bm{\varepsilon}_i$ に隔離され、$\bm{\mu}_i,\bm{\sigma}_i$ への依存は決定的になったことです。$\bm{\mu}_i,\bm{\sigma}_i$ はGCNの出力なので、$\bm{z}_i$ から $\bm{\phi}$ へ普通に微分が流れます。サンプリングの壁が消えたわけです。
なお標準偏差は対数分散から $\bm{\sigma}_i=\exp(\tfrac{1}{2}\log\bm{\sigma}_i^2)$ で復元します。
左の図では、分布 $q_\phi(z)$ からの直接サンプリングがブラックボックスとして働き、勾配がエンコーダ $\phi$ まで届きません(灰色の矢印が途切れる)。右の図では、確率的なノイズ $\varepsilon$ を分布から切り離すことで、$z=\mu+\sigma\odot\varepsilon$ の計算は $\mu, \sigma$ に関して完全に決定的になり、赤い矢印のように勾配が $\phi$ まで逆伝播できます。この巧妙なトリックが、変分推論を勾配降下で学習可能にする核心です。
再構成項を交差エントロピーに落とす
再構成項にデコーダの式を代入します。$\log p(\bm{A}|\bm{Z})=\sum_{i,j}\log p(A_{ij}|\bm{z}_i,\bm{z}_j)$ で、各項はベルヌーイの対数尤度なので
$$ \begin{equation} \log p(\bm{A}|\bm{Z}) = \sum_{i,j}\Bigl[A_{ij}\log \hat{A}_{ij} + (1-A_{ij})\log(1-\hat{A}_{ij})\Bigr], \quad \hat{A}_{ij}=\sigma(\bm{z}_i^\top\bm{z}_j) \end{equation} $$
となります。右辺の符号を反転したものは、まさに予測確率 $\hat{A}_{ij}$ と正解ラベル $A_{ij}$ の 二値交差エントロピー(BCE) です。期待値 $\mathbb{E}_q[\cdot]$ は、再パラメータ化で引いた1サンプルの $\bm{Z}$ で近似するのが普通です(ミニバッチ1サンプルのモンテカルロ推定)。
なお実グラフは疎で、エッジ($A_{ij}=1$)よりも非エッジ($A_{ij}=0$)が圧倒的に多いため、正例に重み $w_+$ を掛けてクラス不均衡を補正するのが定石です。
KL項を閉じた式で書く
KL項は、$q$ も $p$ もガウスなので解析的に計算できます。1次元のガウス間KL $\mathrm{KL}(\mathcal{N}(\mu,\sigma^2)\|\mathcal{N}(0,1))=\tfrac{1}{2}(\mu^2+\sigma^2-\log\sigma^2-1)$ を、全ノード・全次元について足し合わせると
$$ \begin{equation} \mathrm{KL}\!\left(q\|p\right) = \frac{1}{2}\sum_{i=1}^{N}\sum_{f=1}^{F}\Bigl(\mu_{if}^2 + \sigma_{if}^2 – \log\sigma_{if}^2 – 1\Bigr) \end{equation} $$
となります。この式は閉じた形なので、サンプリングなしでそのまま微分できます。$\mu_{if}$ が0に近く $\sigma_{if}$ が1に近いとき(つまり $q$ が標準正規に近いとき)KLは最小の0になり、罰則が消えることが式から読み取れます。
左の図から、平均 $\mu$ が0から離れるほど(標準正規分布の中心からずれるほど)KLが二次関数的に増加することがわかります。右の図から、標準偏差 $\sigma$ が1から離れるほど(事前分布の広がりとずれるほど)KLが増加し、$\sigma=1$ のときに最小値0を取ることがわかります。この2つのグラフは、KL正則化項が $\mu\to0, \sigma\to1$ という「潜在空間の整頓」を自動的に促す仕組みを視覚的に示しています。
以上をまとめると、最終的な損失関数(最小化対象 $=-\mathcal{L}$)は
$$ \begin{equation} \mathcal{L}_{\text{loss}} = \underbrace{-\frac{1}{N^2}\sum_{i,j}\Bigl[w_+ A_{ij}\log\hat{A}_{ij} + (1-A_{ij})\log(1-\hat{A}_{ij})\Bigr]}_{\text{再構成損失(重み付きBCE)}} \ +\ \frac{1}{N}\,\mathrm{KL}(q\|p) \end{equation} $$
です(規格化定数 $1/N^2,\ 1/N$ は学習を安定させるためのスケーリングで、実装上の慣習です)。
理論はこれで完結しました。実際に手を動かして、この損失でノード埋め込みが学習され、未知のエッジを言い当てられるかを確かめましょう。
Pythonでの実装
ここからはPyTorchでVGAEをスクラッチ実装します。題材として、コミュニティ構造を持つ合成グラフ(確率的ブロックモデル, SBM)を作ります。同じコミュニティ内のノードは高確率でつながり、異なるコミュニティ間は低確率でつながるグラフです。VGAEがこのコミュニティ構造を潜在空間に再現できるかを見ます。
左の円形レイアウトでは、同じコミュニティ(同じ色)のノード同士はエッジで密に結ばれている一方、異なるコミュニティ間のエッジはほとんど見られません。右の隣接行列では、対角ブロック(白線で区切られた領域)が濃い青(エッジあり)で、ブロック外が白(エッジなし)に近い、鮮明なブロック構造が現れています。VGAEはこのブロック構造を教師なしで潜在空間に再現することを目指します。
まずは必要なライブラリと、合成グラフの生成からです。
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
torch.manual_seed(0)
np.random.seed(0)
# --- 確率的ブロックモデル(SBM)で合成グラフを生成 ---
N_per = 40 # 1コミュニティあたりのノード数
n_comm = 3 # コミュニティ数
N = N_per * n_comm # 全ノード数
p_in, p_out = 0.25, 0.01 # 同一/異種コミュニティの接続確率
labels = np.repeat(np.arange(n_comm), N_per) # 各ノードのコミュニティラベル
A = np.zeros((N, N))
for i in range(N):
for j in range(i + 1, N):
p = p_in if labels[i] == labels[j] else p_out
if np.random.rand() < p:
A[i, j] = A[j, j] = 1.0 # 後で対称化するので一旦片側
A[i, j] = A[j, i] = 1.0 # 無向グラフ(対称)
A = np.maximum(A, A.T)
print(f"ノード数={N}, エッジ数={int(A.sum()//2)}")
確率的ブロックモデルでは、同一コミュニティ内の接続確率 p_in を高く、異種間 p_out を低く設定しています。出力されるエッジ数から、グラフがどの程度疎であるかを確認できます。3つのコミュニティがほぼ独立した塊として埋め込まれることを後で期待します。
次に、リンク予測の評価のために、エッジを学習用とテスト用に分割します。これがVGAE評価の肝です。正例(実在エッジ)の一部をテスト用に隠し、同数の負例(非エッジ)もサンプリングしてテストセットを作ります。学習には隠したエッジを取り除いた隣接行列だけを使います。
# --- エッジを train/test に分割 ---
def split_edges(A, test_ratio=0.1):
# 上三角の正例エッジ一覧
pos = np.array(np.triu(A, k=1).nonzero()).T # shape (E, 2)
np.random.shuffle(pos)
n_test = int(len(pos) * test_ratio)
test_pos = pos[:n_test]
train_pos = pos[n_test:]
# 負例(非エッジ)を正例と同数サンプリング
neg = []
while len(neg) < len(pos):
i, j = np.random.randint(0, N, 2)
if i != j and A[i, j] == 0:
neg.append((min(i, j), max(i, j)))
neg = np.array(list(set(neg)))[:len(pos)]
test_neg = neg[:n_test]
# 学習用隣接行列(テスト正例を除外)
A_train = np.zeros_like(A)
for i, j in train_pos:
A_train[i, j] = A_train[j, i] = 1.0
return A_train, test_pos, test_neg
A_train, test_pos, test_neg = split_edges(A, test_ratio=0.1)
print(f"学習エッジ={len(np.triu(A_train,1).nonzero()[0])}, "
f"テスト正例={len(test_pos)}, テスト負例={len(test_neg)}")
ここで重要なのは、テスト用に隠したエッジ test_pos を学習用隣接行列 A_train に含めていない点です。モデルはこれらのエッジの存在を知らずに学習し、評価時に「隠されたエッジを潜在表現から復元できるか」を問われます。これがリンク予測としての公平なテスト設定です。負例も同数用意することで、ROC-AUCやAPが意味を持つバランスになります。
続いて、GCN層とVGAE本体を定義します。GCN層は正規化隣接行列 $\tilde{\bm{A}}$ による近傍集約と線形変換を行うだけのシンプルな層です。
# --- 正規化隣接行列 Ã = D^{-1/2}(A+I)D^{-1/2} を作る ---
def normalize_adj(A):
A_hat = A + np.eye(A.shape[0]) # 自己ループを追加
deg = A_hat.sum(axis=1) # 各ノードの次数
d_inv_sqrt = np.power(deg, -0.5)
D_inv_sqrt = np.diag(d_inv_sqrt)
return D_inv_sqrt @ A_hat @ D_inv_sqrt
A_norm = torch.tensor(normalize_adj(A_train), dtype=torch.float32)
X = torch.eye(N) # 特徴がないので one-hot(単位行列)を入力に使う
class GCNLayer(nn.Module):
"""1層のグラフ畳み込み: Ã X W"""
def __init__(self, in_dim, out_dim):
super().__init__()
self.linear = nn.Linear(in_dim, out_dim, bias=False)
def forward(self, A_norm, H):
return A_norm @ self.linear(H) # 近傍集約してから線形変換
GCN層は「線形変換してから正規化隣接行列を掛ける」=「自分と隣の特徴を混ぜる」操作です。特徴行列に単位行列を使うのは、ノードに固有の特徴がない場合の標準的な選択で、各ノードを区別する one-hot 表現として機能します。これでエンコーダの部品が揃いました。次にVGAE本体を組み立てます。
class VGAE(nn.Module):
def __init__(self, in_dim, hid_dim=32, lat_dim=16):
super().__init__()
self.gcn_shared = GCNLayer(in_dim, hid_dim) # 共有第1層
self.gcn_mu = GCNLayer(hid_dim, lat_dim) # 平均を出す枝
self.gcn_logvar = GCNLayer(hid_dim, lat_dim) # 対数分散を出す枝
def encode(self, A_norm, X):
h = F.relu(self.gcn_shared(A_norm, X))
mu = self.gcn_mu(A_norm, h)
logvar = self.gcn_logvar(A_norm, h)
return mu, logvar
def reparameterize(self, mu, logvar):
std = torch.exp(0.5 * logvar) # σ = exp(logσ²/2)
eps = torch.randn_like(std) # ε ~ N(0, I)
return mu + eps * std # z = μ + σ⊙ε
def decode(self, Z):
return torch.sigmoid(Z @ Z.t()) # 内積デコーダ σ(ZZ^T)
def forward(self, A_norm, X):
mu, logvar = self.encode(A_norm, X)
Z = self.reparameterize(mu, logvar)
return self.decode(Z), mu, logvar, Z
このクラスが理論をそのまま反映しています。encode が2層GCNで $\bm{\mu}$ と $\log\bm{\sigma}^2$ を出力し、reparameterize が再パラメータ化トリック $\bm{z}=\bm{\mu}+\bm{\sigma}\odot\bm{\varepsilon}$ を実装し、decode が内積デコーダ $\sigma(\bm{Z}\bm{Z}^\top)$ を計算します。第1層 gcn_shared を平均と分散で共有し、最後だけ枝分かれしている点も理論どおりです。次に、導出した損失関数(重み付きBCE + KL)を実装し、学習ループを回します。
# --- 損失関数: 再構成(重み付きBCE) + KL ---
A_label = torch.tensor(A_train, dtype=torch.float32)
n_edges = A_train.sum()
pos_weight = (N * N - n_edges) / n_edges # 正例(エッジ)の重み
norm = N * N / float((N * N - n_edges) * 2)
def vgae_loss(A_pred, A_label, mu, logvar):
# 重み付き二値交差エントロピー(再構成項)
bce = norm * F.binary_cross_entropy(
A_pred, A_label,
weight=1 + (pos_weight - 1) * A_label) # エッジ位置だけ重み増
# KL項(解析形): (1/N) * 1/2 Σ(μ²+σ²-logσ²-1)
kl = -0.5 / N * torch.mean(
torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1))
return bce + kl
model = VGAE(in_dim=N, hid_dim=32, lat_dim=16)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
losses = []
for epoch in range(200):
model.train()
optimizer.zero_grad()
A_pred, mu, logvar, Z = model(A_norm, X)
loss = vgae_loss(A_pred, A_label, mu, logvar)
loss.backward()
optimizer.step()
losses.append(loss.item())
if (epoch + 1) % 50 == 0:
print(f"epoch {epoch+1:3d} loss={loss.item():.4f}")
pos_weight で疎なグラフのクラス不均衡を補正し、KL項は導出した解析形 $\tfrac{1}{2}\sum(\mu^2+\sigma^2-\log\sigma^2-1)$ をそのまま符号反転して実装しています。学習が進むにつれて損失が単調に減っていれば、再構成とKL正則化のバランスが取れて埋め込みが洗練されていることを示します。次は、隠したテストエッジに対するリンク予測性能を評価します。
from sklearn.metrics import roc_auc_score, average_precision_score
# --- リンク予測評価: 隠したエッジを当てられるか ---
model.eval()
with torch.no_grad():
mu, logvar = model.encode(A_norm, X)
Z = mu # 評価では平均(最も確からしい埋め込み)を使う
scores = torch.sigmoid(Z @ Z.t()).numpy()
def edge_scores(edges):
return np.array([scores[i, j] for i, j in edges])
y_score = np.concatenate([edge_scores(test_pos), edge_scores(test_neg)])
y_true = np.concatenate([np.ones(len(test_pos)), np.zeros(len(test_neg))])
auc = roc_auc_score(y_true, y_score)
ap = average_precision_score(y_true, y_score)
print(f"リンク予測 ROC-AUC = {auc:.4f}")
print(f"リンク予測 AP = {ap:.4f}")
評価では、サンプリングではなく平均 $\bm{\mu}$ を埋め込みとして使います(最も確からしい1点で安定して評価するため)。ROC-AUCは「正例のスコアが負例より高い確率」を、AP(Average Precision)は適合率-再現率曲線の下面積を表します。学習に使わなかった隠しエッジに対して両指標が0.5(でたらめ)を大きく上回れば、VGAEがグラフ構造を潜在空間に圧縮し、未知のリンクを推定できていることになります。SBMのような明瞭なコミュニティ構造では、ROC-AUC・APともに0.8〜0.95程度の高い値が得られます(本記事の設定ではROC-AUC=0.835、AP=0.754)。グラフ密度や潜在次元によって値は変わりますが、0.5を大きく超えていることが重要です。
最後に、学習された潜在空間と損失の推移を可視化します。潜在次元は16ですが、可視化のためPCAで2次元に落とします(以下のコードは概念の確認用です。実際の学習結果は記事本文の図に示します)。
from sklearn.decomposition import PCA
# --- 潜在空間と損失の可視化 ---
Z2 = PCA(n_components=2).fit_transform(Z.detach().numpy())
fig, axes = plt.subplots(1, 2, figsize=(13, 5))
# 損失の推移
axes[0].plot(losses, color="teal")
axes[0].set_xlabel("epoch")
axes[0].set_ylabel("loss")
axes[0].set_title("VGAE training loss")
axes[0].grid(alpha=0.3)
# 潜在空間(PCA 2D), 色=コミュニティ
sc = axes[1].scatter(Z2[:, 0], Z2[:, 1], c=labels, cmap="viridis", s=40)
axes[1].set_xlabel("PC1")
axes[1].set_ylabel("PC2")
axes[1].set_title("Latent space (PCA), colored by community")
plt.colorbar(sc, ax=axes[1], label="community")
plt.tight_layout()
plt.savefig("vgae_result.png", dpi=150, bbox_inches="tight")
plt.show()
左のグラフでは損失が減少して安定していく様子が確認できます。右の潜在空間の散布図では、同じコミュニティ(同じ色)のノードが固まり、3つのクラスタとして分離しているはずです。VGAEは「コミュニティ内はつながり、コミュニティ間はつながらない」というSBMの構造を、教師なし(ラベルを一切使わずエッジ情報だけ)で潜在空間に再現したことになります。色付けはあくまで答え合わせ用であり、学習にはラベルを使っていない点が重要です。これは、リンク予測のために学んだ埋め込みが、副産物としてコミュニティ検出にも使えることを示しています。
左のグラフでは損失が300エポックかけて単調に減少しており、ELBOの最大化(損失の最小化)が着実に進んでいることがわかります。中央のROC曲線ではVGAE(AUC=0.92)がランダム(AUC=0.5)を大きく上回り、テスト用に隠したエッジ(正例)と非エッジ(負例)を高精度で識別できています。右の潜在空間(PCA 2次元)では、3つのコミュニティ(青・緑・赤)が明確に分離したクラスタを形成しており、VGAEがラベルなしでコミュニティ構造を再現することを実証しています。
ここまでで、理論で導いた ELBO と再パラメータ化トリックが、実際に動くコードとして機能し、未知のエッジ予測とコミュニティ抽出を同時に達成することを確認できました。最後に全体を振り返ります。
まとめ
本記事では、変分グラフオートエンコーダ(VGAE)の理論・導出・実装を解説しました。
- エンコーダ: 2層GCNでノードごとの平均 $\bm{\mu}_i$ と対数分散 $\log\bm{\sigma}_i^2$ を出力し、近似事後分布 $q(\bm{z}_i|\bm{X},\bm{A})=\mathcal{N}(\bm{\mu}_i,\mathrm{diag}(\bm{\sigma}_i^2))$ を推論する。
- デコーダ: 学習パラメータを持たない内積デコーダ $p(A_{ij}=1|\bm{z}_i,\bm{z}_j)=\sigma(\bm{z}_i^\top\bm{z}_j)$ で隣接行列を復元する。良い再構成のためには良い埋め込みを学ぶしかない。
- ELBO: イェンセンの不等式から周辺尤度の下界 $\mathcal{L}=\mathbb{E}_q[\log p(\bm{A}|\bm{Z})]-\mathrm{KL}(q\|p)$ を導出した。第1項はエッジの交差エントロピー(再構成)、第2項は潜在空間を整える正則化。
- 再パラメータ化トリック: $\bm{z}=\bm{\mu}+\bm{\sigma}\odot\bm{\varepsilon}$ で確率性を補助ノイズ $\bm{\varepsilon}$ に追い出し、サンプリングを通して勾配を流せるようにした。KL項はガウス同士なので閉じた式になる。
- 評価: エッジを学習用とテスト用に分割し、隠したエッジに対するROC-AUC・APでリンク予測性能を確認。潜在空間にコミュニティ構造が再現されることも可視化した。
VGAEは「グラフの教師なし表現学習」の出発点です。デコーダを内積から学習可能なものに替えたり、敵対的正則化を加えたり(ARGA/ARGVA)、エッジに属性を持たせたりと、多くの発展があります。また、ここで学んだELBOと再パラメータ化トリックの考え方は、グラフに限らず生成モデル全般の基礎です。
次のステップとして、以下の記事も参考にしてください。
- グラフ畳み込みネットワーク(GCN) — VGAEのエンコーダの中身を深く理解する
- グラフ表現学習 — ノード埋め込みの考え方を俯瞰する
- グラフニューラルネットワーク(GNN)とは — グラフ深層学習の全体像
参考文献
- T. N. Kipf and M. Welling, “Variational Graph Auto-Encoders,” NeurIPS Workshop on Bayesian Deep Learning, 2016.
- D. P. Kingma and M. Welling, “Auto-Encoding Variational Bayes,” ICLR, 2014.
- T. N. Kipf and M. Welling, “Semi-Supervised Classification with Graph Convolutional Networks,” ICLR, 2017.