GATのマルチヘッドアテンション機構を理解して実装する

GATアテンションの概念: 隣人ごとに重みを変える

ソーシャルネットワークであなたの興味を予測したいとき、フォローしている全員を等しく扱うのは賢明でしょうか。おそらく、あなたと関心の近い数人の影響が大きく、惰性でフォローしている多数のアカウントの影響は小さいはずです。つまり「同じ隣人でも、重みが違う」のです。ところが従来のグラフ畳み込み(GCN)は、隣人の重みをグラフの構造(次数)だけで機械的に決めてしまい、隣人の特徴量がどれだけ自分に関係するかを見ていません。

この弱点を正面から解決したのが GAT(Graph Attention Network) です。GATは「どの隣人をどれだけ重視するか」を、固定式ではなく 学習可能なアテンション機構 で決めます。さらに、複数の視点から同時に重み付けする マルチヘッドアテンション を導入することで、一つの集約だけでは見落とす多面的な関係を捉えます。

GATを理解すると、次のような応用に直結します。

  • 論文引用ネットワークの分類: 引用先・引用元のうち、どの論文が分類に効いているかを注意係数として解釈できる
  • 分子グラフの物性予測: 分子内のどの原子間結合が物性に寄与しているかを、アテンションの重みから読み取れる

この記事では、注意スコア $e_{ij}$ の定義から、LeakyReLU と共有重みを使った具体的な計算、ソフトマックスによる正規化 $\alpha_{ij}$、そしてマルチヘッドの連結・平均という集約までを 一行ずつ省略せず 導出します。その上で、GCN の固定された次数正規化重みとの違いをはっきりさせ、PyTorch Geometric で Cora データセットを学習し、学習された注意係数をエッジの太さとして可視化します。最後にヘッド数を変えて精度と注意分布の違いを観察します。

本記事の内容

  • GAT のアテンション機構(注意スコアと正規化)の直感と数学的定義
  • 注意係数 $\alpha_{ij}$ とマルチヘッド集約の一行ずつの導出
  • GCN の固定重みと GAT の動的重みの違いの明確化
  • PyTorch Geometric による実装・学習と注意係数の可視化

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が一段と深まります。GAT は「メッセージパッシングの枠組みの中で、集約の重みをアテンションで決める手法」なので、メッセージパッシングと GCN を先に押さえておくと、GAT が何を新しくしたのかがクリアになります。

本記事は、既存の GAT 入門記事の「理論深掘り版」です。入門記事で概要をつかんだ方が、注意係数の導出やマルチヘッドの数式、GCN との数式レベルの違いまで踏み込みたいときに読む位置づけです。

GAT のアテンションとは

まず、私たちが日常的に行っている「重み付き平均」をイメージしてみましょう。あなたが新しいレストランを選ぶとき、3人の友人がそれぞれ別の店を勧めてきたとします。グルメな友人 A の意見は強く重視し、食にこだわりのない友人 C の意見は軽く流す——このとき、あなたは無意識に「友人の意見の信頼度」に応じた重みをつけて意思決定しています。重みの合計は 1(100%)になるように、頭の中で正規化しているはずです。

GAT のアテンションも、まさにこれと同じことをします。ノード $i$ が自分の新しい特徴を計算するとき、隣人 $j$ たちの特徴を集めて重み付き和を取りますが、その重み $\alpha_{ij}$ を「ノード $i$ にとってノード $j$ がどれだけ重要か」に応じて決めるのです。しかも、その重みは固定ではなく、ノードの特徴を見て 学習 されます。

ここで重要なのは、重み $\alpha_{ij}$ が次の2つの性質を持つことです。

  1. 非負で、隣人全体にわたる合計が 1 になる(確率分布のように振る舞う)
  2. 特徴依存である。隣人 $j$ の特徴がノード $i$ に関係が深いほど大きくなる

従来の GCN では、この重みがグラフの構造(ノードの次数)だけで決まっていました。GAT はここに「特徴を見て動的に重みを決める」という発想を持ち込みます。次の節では、この「重要度」をどうやって数値(スコア)に落とすのかを、定義から丁寧に見ていきましょう。

注意スコア $e_{ij}$ の定義

アテンションの第一歩は、「ノード $i$ にとってノード $j$ がどれだけ重要か」を表す 生のスコア $e_{ij}$ を計算することです。重み $\alpha_{ij}$ そのものではなく、まず正規化前のスコアを作り、あとでソフトマックスで正規化する、という二段構えで進めます。

スコアを作るには、ノード $i$ とノード $j$ の特徴を見比べる必要があります。GAT のオリジナル論文(Veličković ら, 2018)が採用したのは、次の形です。各ノード $i$ は特徴ベクトル $\bm{h}_i \in \mathbb{R}^{F}$ を持っているとします。$F$ は入力特徴の次元です。

まず、すべてのノードに 共有された 線形変換 $\bm{W} \in \mathbb{R}^{F’ \times F}$ を適用して、特徴を $F’$ 次元に射影します。

$$ \begin{equation} \bm{z}_i = \bm{W} \bm{h}_i \end{equation} $$

ここで「共有された」という点が大切です。すべてのノードに同じ $\bm{W}$ を使うことで、パラメータ数がノード数に依存せず、学習結果が新しいノードにも一般化できます(帰納的な学習が可能になります)。

次に、変換後のノード $i$ とノード $j$ の特徴を 連結 し、学習可能なベクトル $\bm{a} \in \mathbb{R}^{2F’}$ との内積を取って、活性化関数 LeakyReLU を通します。これが注意スコアです。

$$ \begin{equation} e_{ij} = \mathrm{LeakyReLU}\!\left( \bm{a}^\top [\, \bm{W}\bm{h}_i \,\Vert\, \bm{W}\bm{h}_j \,] \right) \end{equation} $$

ここで $[\cdot \Vert \cdot]$ はベクトルの連結(concatenation)を表します。$\bm{W}\bm{h}_i$ と $\bm{W}\bm{h}_j$ はそれぞれ $F’$ 次元なので、連結すると $2F’$ 次元になり、$\bm{a}$ の次元と一致します。

LeakyReLU を使う理由

なぜ普通の ReLU ではなく LeakyReLU なのでしょうか。ReLU は負の入力をすべて 0 にしてしまいます。もし $\bm{a}^\top[\cdots]$ が負になったノードのスコアが一律に 0 になると、ソフトマックスを取る前の段階で「どのくらい重要でないか」の情報(負の度合い)が失われてしまいます。LeakyReLU は

$$ \begin{equation} \mathrm{LeakyReLU}(x) = \begin{cases} x & (x \geq 0) \\ \alpha_{\text{slope}}\, x & (x < 0) \end{cases} \end{equation} $$

と定義され、負の側にも小さな傾き $\alpha_{\text{slope}}$(GAT 論文では 0.2)を残します。これにより、負のスコアも「弱い負」「強い負」と区別され、勾配も流れ続けるため学習が安定します。

注意スコアeijの計算フロー: 線形変換から連結内積LeakyReLUまで

この計算フロー図では、入力特徴 $\bm{h}_i, \bm{h}_j$ から始まり、共有重み $\bm{W}$ による線形変換、連結、アテンションベクトル $\bm{a}$ との内積、そして LeakyReLU による活性化という5段階が一目でわかります。各段階でテンソルの次元が変化し($F \to F’ \to 2F’ \to 1 \to 1$)、最終的に1つのスカラースコアが得られます。

ReLUとLeakyReLUの比較: 負の領域でも勾配が残る理由

このグラフから、ReLU と LeakyReLU の本質的な違いが読み取れます。ReLU では $x < 0$ の領域が一律に 0 になり、勾配消失を引き起こすのに対し、LeakyReLU は傾き 0.2 の小さな傾きが残ることで「この隣人は少し重要でない」「かなり重要でない」という負方向のグラデーションを保存します。この情報が次のソフトマックス正規化に活きてくるのです。

スコア $e_{ij}$ は、ノード $i$ から見たノード $j$ の重要度を表す実数ですが、まだ正規化されていません。隣人ごとにスケールがバラバラで、足しても 1 になりません。そこで次の節では、この生スコアを「合計 1 の重み」に変換するソフトマックス正規化を導入します。

注意係数 $\alpha_{ij}$ への正規化

生のスコア $e_{ij}$ をそのまま重みとして使うわけにはいきません。隣人が10人いるノードと2人しかいないノードでは、スコアの合計のスケールが全く違ってしまい、層を重ねるうちに値が発散したり消失したりするからです。レストラン選びの例で「友人の意見の重みの合計を 1 にする」とした直感を、ここで数式にします。

ノード $i$ の隣人集合を $\mathcal{N}_i$ とします(通常、自分自身も含めるために自己ループを加えます)。注意係数 $\alpha_{ij}$ は、$\mathcal{N}_i$ の上でソフトマックスを取って定義します。

$$ \begin{equation} \alpha_{ij} = \mathrm{softmax}_j(e_{ij}) = \frac{\exp(e_{ij})}{\displaystyle\sum_{k \in \mathcal{N}_i} \exp(e_{ik})} \end{equation} $$

この定義により、次の2つが保証されます。指数関数 $\exp$ は常に正なので $\alpha_{ij} > 0$ であり、分母が「隣人全体の和」なので

$$ \begin{equation} \sum_{j \in \mathcal{N}_i} \alpha_{ij} = \frac{\sum_{j \in \mathcal{N}_i}\exp(e_{ij})}{\sum_{k \in \mathcal{N}_i}\exp(e_{ik})} = 1 \end{equation} $$

が成り立ちます。つまり $\alpha_{ij}$ は、ノード $i$ の隣人にわたる確率分布になっています。これがアテンションを「ソフトな選択」と呼ぶ理由です。重要な隣人に大きな確率を、そうでない隣人に小さな確率を割り当てます。

スコアの全体像(展開形)

ここまでの定義を一つの式にまとめると、注意係数の完全な形が見えます。$e_{ij}$ の定義を代入すると、

$$ \begin{equation} \alpha_{ij} = \frac{\exp\!\left( \mathrm{LeakyReLU}\!\left( \bm{a}^\top [\bm{W}\bm{h}_i \Vert \bm{W}\bm{h}_j] \right) \right)}{\displaystyle\sum_{k \in \mathcal{N}_i} \exp\!\left( \mathrm{LeakyReLU}\!\left( \bm{a}^\top [\bm{W}\bm{h}_i \Vert \bm{W}\bm{h}_k] \right) \right)} \end{equation} $$

となります。少し複雑に見えますが、構造は単純です。分子が「ノード $i$ と $j$ の相性スコアの指数」、分母が「ノード $i$ と全隣人の相性スコアの指数の和」です。相性が良い隣人ほど、分子が相対的に大きくなり、$\alpha_{ij}$ が 1 に近づきます。

ソフトマックス正規化: 生スコアから確率分布への変換

この棒グラフから、ソフトマックスの本質が読み取れます。左の生スコア $e_{ij}$ は正負まちまちでスケールが揃っていませんが、右のソフトマックス後の注意係数 $\alpha_{ij}$ はすべて $[0, 1]$ に収まり合計がぴったり 1.000 になっています。生スコアが最大の隣人2(0.562)に最大の注意係数が割り当てられており、スコアの大小関係が正しく反映されていることも確認できます。これが「ソフトな選択」と呼ばれる理由です。

$\bm{a}$ を分解して内積の意味を読む

連結と内積の部分は、$\bm{a}$ を前半 $\bm{a}_1 \in \mathbb{R}^{F’}$ と後半 $\bm{a}_2 \in \mathbb{R}^{F’}$ に分けると、もっと見通しが良くなります。連結ベクトルと $\bm{a}$ の内積は、

$$ \begin{equation} \bm{a}^\top [\bm{W}\bm{h}_i \Vert \bm{W}\bm{h}_j] = \bm{a}_1^\top (\bm{W}\bm{h}_i) + \bm{a}_2^\top (\bm{W}\bm{h}_j) \end{equation} $$

と分解できます。連結したベクトルとの内積は、前半同士の内積と後半同士の内積の和に等しいからです。この形を見ると、スコアが「ノード $i$ 由来の項」と「ノード $j$ 由来の項」の和でできていることがわかります。第1項はノード $i$ が自分の状態に応じてベースラインを決め、第2項は隣人 $j$ の特徴がそれをどれだけ押し上げる(または下げる)かを表します。これは実装上も効率的で、$\bm{a}_1^\top \bm{W}\bm{h}_i$ と $\bm{a}_2^\top \bm{W}\bm{h}_j$ をノードごとに一度だけ計算し、エッジごとに足し合わせるだけで済みます。

アテンションベクトルaの分解: ソースノード項と宛先ノード項

このフロー図では、$\bm{a}$ を前半 $\bm{a}_1$ と後半 $\bm{a}_2$ に分解することで、内積が「ノード $i$ 由来の項 $\bm{a}_1^\top \bm{z}_i$」と「ノード $j$ 由来の項 $\bm{a}_2^\top \bm{z}_j$」の和になることが視覚的にわかります。この分解は実装上も非常に重要で、各ノードの項を事前計算しておき、エッジごとに足し合わせるだけで済むため、全ノード数 + エッジ数に比例した計算量(大規模グラフでも効率的)を実現できます。

正規化された注意係数 $\alpha_{ij}$ が手に入ったので、いよいよこれを使って隣人の特徴を集約します。次の節では、$\alpha_{ij}$ を重みとした集約式と、複数のアテンションを並列に走らせるマルチヘッドの仕組みを導出します。

特徴の集約とマルチヘッドアテンション

注意係数 $\alpha_{ij}$ が決まれば、あとはこれを重みにして隣人の特徴を集めるだけです。ノード $i$ の新しい特徴 $\bm{h}_i’$ は、隣人の変換済み特徴 $\bm{W}\bm{h}_j$ を $\alpha_{ij}$ で重み付き和して、非線形活性化 $\sigma$(ELU など)を通します。

$$ \begin{equation} \bm{h}_i’ = \sigma\!\left( \sum_{j \in \mathcal{N}_i} \alpha_{ij}\, \bm{W}\bm{h}_j \right) \end{equation} $$

これがアテンション一つ分、すなわち シングルヘッド の更新式です。$\alpha_{ij}$ が大きい隣人の特徴が、新しい特徴に強く反映されることがわかります。

なぜマルチヘッドが必要か

マルチヘッドアテンション: K個の独立した観点からの重み付け

このアーキテクチャ図から、マルチヘッドの全体像が読み取れます。同じ入力特徴 $\bm{h}_i, \bm{h}_j$ が $K$ 個のヘッドに並列に入り、それぞれ独立したパラメータ $\bm{W}^{(k)}, \bm{a}^{(k)}$ で処理されます。$K$ 個の出力は、中間層では連結(次元が $KF’$ に増える)、出力層では平均($F’$ 次元を保つ)で集約されます。この分岐と集約が「複数の観点で隣人を評価し、合議で決める」という直感を実現しています。

一つのアテンションだけでは、隣人の見方が一通りに固定されてしまいます。レストラン選びでも「味の好み」という観点だけでなく、「価格」「立地」「雰囲気」など複数の観点で友人の意見を重み付けしたいはずです。観点ごとに重視する友人は違うでしょう。

マルチヘッドアテンションは、これを実現します。$K$ 個の独立したアテンション機構(ヘッド)を並列に走らせ、それぞれが独自の重み行列 $\bm{W}^{(k)}$ とアテンションベクトル $\bm{a}^{(k)}$ を持ちます。ヘッド $k$ の注意係数を $\alpha_{ij}^{(k)}$ と書くと、ヘッド $k$ の出力は

$$ \begin{equation} \bm{h}_i^{(k)} = \sum_{j \in \mathcal{N}_i} \alpha_{ij}^{(k)}\, \bm{W}^{(k)}\bm{h}_j \end{equation} $$

です。各ヘッドが「異なる観点」で隣人を重み付けすることで、関係の多面性を捉えます。問題は、$K$ 個の出力 $\bm{h}_i^{(1)}, \dots, \bm{h}_i^{(K)}$ をどうまとめるか、です。GAT は層の位置に応じて2つの集約方法を使い分けます。

中間層:連結(concatenation)

中間層では、$K$ 個のヘッドの出力を 連結 します。

$$ \begin{equation} \bm{h}_i’ = \Big\Vert_{k=1}^{K} \sigma\!\left( \sum_{j \in \mathcal{N}_i} \alpha_{ij}^{(k)}\, \bm{W}^{(k)}\bm{h}_j \right) \end{equation} $$

ここで $\Vert_{k=1}^{K}$ は $K$ 個のベクトルを順に連結する操作です。各ヘッドの出力が $F’$ 次元なら、連結後は $K F’$ 次元になります。連結を使うと各ヘッドの情報がそのまま保たれ、次の層がヘッドごとの特徴を区別して使えます。中間層では情報を失いたくないので、連結が好まれます。出力次元が $K$ 倍になる点に注意してください。

出力層:平均(averaging)

最終層では、出力の次元をクラス数などに合わせたいので、連結ではなく 平均 を取ります。

$$ \begin{equation} \bm{h}_i’ = \sigma\!\left( \frac{1}{K} \sum_{k=1}^{K} \sum_{j \in \mathcal{N}_i} \alpha_{ij}^{(k)}\, \bm{W}^{(k)}\bm{h}_j \right) \end{equation} $$

平均では、$K$ 個のヘッドの出力を要素ごとに足して $K$ で割るため、出力次元は $F’$ のまま保たれます。複数ヘッドの「合議」を取ることで、最終的な予測が安定します。連結と違って次元が増えないので、分類タスクの最終層ではこちらが使われます。

連結と平均の違いを数式で押さえる

両者の本質的な違いを整理しておきましょう。連結は $K$ 個の $F’$ 次元ベクトルを並べて $KF’$ 次元にする操作です。情報は完全に保存されますが、次元とそれに続く層のパラメータ数が $K$ 倍になります。一方、平均は $K$ 個のベクトルを重ね合わせて $F’$ 次元のままにする操作で、次元は増えませんが、各ヘッド固有の細かい情報は平均化で薄まります。「中間層は情報量重視で連結、出力層は次元固定とアンサンブル効果重視で平均」という設計は、この性質から自然に導かれます。

連結と平均の次元変化: 情報保存と次元固定のトレードオフ

左の棒グラフから、連結では $K$ が増えるにつれて出力次元が $K \times F’$ に線形に増加する一方、平均では常に $F’$ 次元に固定されることが一目でわかります。右の散布図は、連結が表現の豊かさを最大化する一方で、平均がより安定した汎化性能をもたらすという設計上のトレードオフを示しています。「中間層は連結で情報を保存し、出力層は平均でクラス数に合わせた次元に収める」という GAT の設計がこのトレードオフの賢い折衷案です。

ここまでで GAT の更新式が完成しました。では、この「動的に決まる重み」は、GCN の重みと具体的に何が違うのでしょうか。次の節で両者を並べて比較します。

GCN の固定重みとの違い

GCNの固定重みとGATの学習重みの決定メカニズム比較

上の図が示すように、GCN は特徴を全く見ずにグラフ構造(次数)だけで重みを固定するのに対し、GAT は学習可能なパラメータ $\bm{W}$ と $\bm{a}$ を通じて特徴から動的に重みを決めます。この根本的な違いが、以降の数式比較でより明確になります。

GAT のありがたみは、GCN と並べて見ると一目瞭然です。GCN(Kipf & Welling, 2017)の1層の更新式は、対称正規化を使った次の形でした。

$$ \begin{equation} \bm{h}_i’ = \sigma\!\left( \sum_{j \in \mathcal{N}_i} \frac{1}{\sqrt{d_i}\sqrt{d_j}}\, \bm{W}\bm{h}_j \right) \end{equation} $$

ここで $d_i$ はノード $i$ の次数(自己ループを含む隣人の数)です。この式の重み

$$ \begin{equation} c_{ij} = \frac{1}{\sqrt{d_i}\sqrt{d_j}} \end{equation} $$

に注目してください。これは グラフの構造だけ で決まります。ノード $i$ と $j$ の次数がわかれば、特徴量 $\bm{h}_i, \bm{h}_j$ を一切見なくても重みが確定します。しかも学習中にこの重みは変わりません。つまり GCN は「次数の小さい(=希少な)隣人を相対的に重視する」という、構造ベースの固定ルールに従っているだけなのです。

これを GAT の重み $\alpha_{ij}$ と比べてみましょう。下の表に違いをまとめます。

観点 GCN の重み $c_{ij}$ GAT の重み $\alpha_{ij}$
決まり方 次数のみ(構造) 特徴 $\bm{h}_i, \bm{h}_j$ から学習
学習中の変化 変化しない(固定) パラメータ更新で変化する
対称性 $c_{ij} = c_{ji}$(対称) $\alpha_{ij} \neq \alpha_{ji}$(非対称)
同じ次数の隣人の扱い 全員同じ重み 特徴に応じて差をつける
隣人ごとの解釈 解釈困難(構造のみ) 重要な隣人を可視化できる

特に重要なのが 非対称性 です。GCN では $c_{ij} = c_{ji}$ ですが、GAT のスコア $e_{ij} = \mathrm{LeakyReLU}(\bm{a}_1^\top\bm{W}\bm{h}_i + \bm{a}_2^\top\bm{W}\bm{h}_j)$ は $i$ と $j$ を入れ替えると $\bm{a}_1$ と $\bm{a}_2$ の役割が入れ替わり、しかも分母の正規化はノードごとに異なる隣人集合で行われるため、一般に $\alpha_{ij} \neq \alpha_{ji}$ となります。これは「A にとっての B の重要度」と「B にとっての A の重要度」が違ってよい、という直感に合っています。あなたにとって有名人は重要でも、有名人にとってあなたは多数のフォロワーの一人にすぎない、という関係です。

注意係数の非対称性: GCN対称重みとGAT非対称重みの比較

この比較図から、対称性の違いが視覚的に明確になります。左の GCN では、全ノードの次数が同じ場合は全エッジに同じ重み(ここでは 0.33)が割り当てられ、方向による差がありません。右の GAT では、「ノード i から見た有名人 J の重み(0.55)」と「有名人 J から見たノード i の重み(0.10)」がまったく異なります。フォロワー10万人のアカウントにとって自分は多数の一人にすぎないが、自分にとってそのアカウントは非常に重要という非対称な関係を、GAT は自然に表現できるのです。

統一的な見方:メッセージパッシング

GCN も GAT も、実は同じ メッセージパッシング の枠組みに収まります。一般のメッセージパッシングは

$$ \begin{equation} \bm{h}_i’ = \sigma\!\left( \sum_{j \in \mathcal{N}_i} w_{ij}\, \bm{W}\bm{h}_j \right) \end{equation} $$

と書け、重み $w_{ij}$ の決め方だけが両者で異なります。GCN は $w_{ij} = 1/\sqrt{d_i d_j}$ という固定値、GAT は $w_{ij} = \alpha_{ij}$ という学習される値です。つまり GAT は「GCN の固定重みを、特徴依存の学習可能な重みに置き換えたもの」と理解できます。この一文が、GAT が GCN の自然な拡張であることの核心です。

メッセージパッシング統一フレームワーク: GCNとGATの重みの決め方の違い

この統一フレームワーク図は、GCN と GAT が実は同じ枠組みの特殊ケースであることを示しています。どちらも $\bm{h}_i’ = \sigma(\sum_{j \in \mathcal{N}_i} w_{ij} \bm{W} \bm{h}_j)$ という形をしており、違いは $w_{ij}$ の決め方だけです。GCN は $1/\sqrt{d_i d_j}$(固定・対称)、GAT は $\alpha_{ij}$(動的・非対称)です。この見方をすると、GAT は「GCN のメッセージパッシングの重みを学習可能にした自然な拡張」という位置づけが明確になります。

理論はここまでです。次は実際に PyTorch Geometric で GAT を組み、Cora データセットで学習し、本当に隣人ごとに違う重みが学習されるのかを目で確かめましょう。

Python での実装

ここからは、注意係数の計算を自分の手で確かめるところから始め、最後に PyTorch Geometric で実用的な GAT を学習します。まずは小さな例で、$e_{ij}$ から $\alpha_{ij}$ までの計算が「合計 1 の確率分布」になることを numpy で確認します。

import numpy as np

np.random.seed(0)

# LeakyReLU の定義(負の傾き 0.2)
def leaky_relu(x, slope=0.2):
    return np.where(x >= 0, x, slope * x)

# 入力特徴: 4ノード, 各3次元
F_in, F_out = 3, 2
H = np.random.randn(4, F_in)

# 共有重み W (F_out x F_in) とアテンションベクトル a (2*F_out)
W = np.random.randn(F_out, F_in)
a = np.random.randn(2 * F_out)

# 変換後特徴 z_i = W h_i
Z = H @ W.T            # 形状 (4, F_out)
a1, a2 = a[:F_out], a[F_out:]

# ノード0が隣人 {0,1,2,3}(自己ループ含む)を持つとして注意スコアを計算
i = 0
neighbors = [0, 1, 2, 3]
scores = []
for j in neighbors:
    e_ij = leaky_relu(a1 @ Z[i] + a2 @ Z[j])  # 分解形で計算
    scores.append(e_ij)
scores = np.array(scores)

# ソフトマックスで正規化
alpha = np.exp(scores) / np.sum(np.exp(scores))

print("生スコア e_ij:", np.round(scores, 3))
print("注意係数 alpha_ij:", np.round(alpha, 3))
print("合計:", round(alpha.sum(), 6))

このコードでは、本文で導出した分解形 $e_{ij} = \mathrm{LeakyReLU}(\bm{a}_1^\top \bm{z}_i + \bm{a}_2^\top \bm{z}_j)$ をそのまま計算しています。出力を見ると、生スコア $e_{ij}$ は正負まちまちのバラバラな値ですが、ソフトマックスを通した注意係数 $\alpha_{ij}$ はすべて 0 から 1 の値になり、合計がぴったり 1 になっていることが確認できます。これは数式 (5) で示した「$\alpha_{ij}$ は隣人にわたる確率分布である」という性質が、実装でも正しく成り立っていることを意味します。スコアが最も大きかった隣人に、最も大きな注意係数が割り当てられている点にも注目してください。

GCN 重みと GAT 重みの違いを数値で見る

次に、同じグラフに対して GCN の固定重み $c_{ij}$ と GAT の学習可能な重み $\alpha_{ij}$ がどう違うかを並べて計算してみます。GCN の重みは次数だけで決まり、特徴を一切見ないことを実感するのが目的です。

import numpy as np

# 5ノードの単純な無向グラフ(隣接リスト、自己ループ込み)
adj = {
    0: [0, 1, 2],
    1: [1, 0, 2, 3],
    2: [2, 0, 1],
    3: [3, 1, 4],
    4: [4, 3],
}
deg = {i: len(adj[i]) for i in adj}  # 次数

# GCN の対称正規化重み c_ij = 1/sqrt(d_i d_j)
print("=== GCN の固定重み c_0j(ノード0の隣人)===")
for j in adj[0]:
    c = 1.0 / np.sqrt(deg[0] * deg[j])
    print(f"  c_0{j} = {c:.3f}  (特徴に依存しない)")

# GAT の重み(前のコードの alpha を流用するイメージ)
np.random.seed(1)
alpha_gat = np.random.dirichlet(np.ones(len(adj[0])))  # 合計1の重みの例
print("=== GAT の学習される重み alpha_0j ===")
for j, a_val in zip(adj[0], alpha_gat):
    print(f"  alpha_0{j} = {a_val:.3f}  (特徴から学習)")

出力を見ると、GCN の重み $c_{0j}$ は隣人 $j$ の次数が大きいほど小さくなる、という構造だけのルールで決まっています。同じ次数の隣人なら必ず同じ重みになります。一方 GAT の $\alpha_{0j}$ は、ここではダミーですが、合計 1 を保ちながら隣人ごとに自由な値を取れます。実際の GAT ではこの値が特徴から学習され、「次数が同じでも特徴が違えば違う重み」が付くわけです。GCN が構造に縛られているのに対し、GAT は特徴を見て柔軟に重みを変えられる、という本質的な違いがここに現れています。

PyTorch Geometric で Cora を学習する

いよいよ実データで GAT を学習します。Cora は2708本の論文を引用関係でつないだグラフで、各論文を7つの研究分野のいずれかに分類するタスクです。PyTorch Geometric の GATConv を使い、8ヘッドの中間層(連結)と1ヘッドの出力層(平均)を持つ標準的な構成にします。

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GATConv

# Cora データセットの読み込み
dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]

class GAT(torch.nn.Module):
    def __init__(self, in_dim, hid_dim, out_dim, heads=8):
        super().__init__()
        # 中間層: heads個のヘッドを連結 -> 出力次元は hid_dim*heads
        self.conv1 = GATConv(in_dim, hid_dim, heads=heads, dropout=0.6)
        # 出力層: 連結された特徴を入力, 1ヘッドで平均(concat=False)
        self.conv2 = GATConv(hid_dim * heads, out_dim, heads=1,
                             concat=False, dropout=0.6)

    def forward(self, x, edge_index, return_attention=False):
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        if return_attention:
            x, attn = self.conv2(x, edge_index,
                                 return_attention_weights=True)
            return x, attn
        x = self.conv2(x, edge_index)
        return x

model = GAT(dataset.num_features, 8, dataset.num_classes, heads=8)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005, weight_decay=5e-4)

このモデル定義では、本文の数式 (11) と (12) がそのまま GATConv の引数に対応しています。中間層 conv1heads=8 で連結(デフォルト concat=True)なので出力が $8 \times 8 = 64$ 次元になり、出力層 conv2concat=False で平均集約を行い7クラスの次元に落とします。dropout=0.6 はアテンション係数とノード特徴にドロップアウトをかけ、Cora のような小規模グラフでの過学習を防ぐための設定です。return_attention_weights=True を指定すると、学習後に各エッジの注意係数を取り出せます。

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    # 訓練ノードのみで交差エントロピー損失
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

@torch.no_grad()
def test():
    model.eval()
    out = model(data.x, data.edge_index)
    pred = out.argmax(dim=1)
    accs = []
    for mask in [data.train_mask, data.val_mask, data.test_mask]:
        acc = (pred[mask] == data.y[mask]).float().mean().item()
        accs.append(acc)
    return accs

for epoch in range(1, 201):
    loss = train()
    if epoch % 40 == 0:
        train_acc, val_acc, test_acc = test()
        print(f"Epoch {epoch:3d} | loss {loss:.3f} | "
              f"val {val_acc:.3f} | test {test_acc:.3f}")

この学習ループを回すと、損失が単調に下がり、検証・テスト精度が上がっていきます。標準的な設定では、200エポックでテスト精度がおよそ 0.82 前後(82%程度)に達します。これは特徴を一切重み付けしない単純な平均集約より明確に高く、隣人ごとに重みを変えるアテンションが分類に効いていることを示しています。エポックが進んでも検証精度が大きく崩れないのは、ドロップアウトが過学習を抑えているおかげです。

学習された注意係数をエッジの太さで可視化する

GAT の最大の魅力は、学習された重みを「解釈」できる点です。あるノードに着目し、その隣人へのエッジを注意係数 $\alpha_{ij}$ に比例した太さで描いてみましょう。重要な隣人ほど太い線で結ばれるはずです。

import numpy as np
import matplotlib.pyplot as plt

# 学習後、注意係数付きで forward
model.eval()
with torch.no_grad():
    _, (edge_index_att, alpha) = model(
        data.x, data.edge_index, return_attention=True)

# alpha は (エッジ数, ヘッド数) -> ここは出力層なので1ヘッド分
alpha = alpha.mean(dim=1).cpu().numpy()        # ヘッド平均
edge_index_att = edge_index_att.cpu().numpy()

# 着目ノード(例: 次数が中程度のノード)を選ぶ
target = 100
mask = edge_index_att[1] == target             # target を受け取る側のエッジ
src_nodes = edge_index_att[0][mask]
weights = alpha[mask]

# 円周上に隣人を配置して放射状に描画
n = len(src_nodes)
angles = np.linspace(0, 2 * np.pi, n, endpoint=False)
plt.figure(figsize=(7, 7))
for ang, w, s in zip(angles, weights, src_nodes):
    x, y = np.cos(ang), np.sin(ang)
    # 注意係数に比例した太さでエッジを描く
    plt.plot([0, x], [0, y], color='steelblue',
             linewidth=1 + 20 * w, alpha=0.7)
    plt.scatter([x], [y], s=60, color='orange', zorder=5)
    plt.text(1.15 * x, 1.15 * y, f"{w:.2f}", ha='center', va='center')

plt.scatter([0], [0], s=200, color='crimson', zorder=6)
plt.text(0, -0.12, f"node {target}", ha='center')
plt.title(f"Learned attention weights into node {target}")
plt.axis('equal'); plt.axis('off')
plt.tight_layout()
plt.savefig('gat_attention.png', dpi=150, bbox_inches='tight')
plt.show()

この図では、中心の赤いノードが着目ノード、周囲のオレンジのノードがその隣人で、エッジの太さが注意係数 $\alpha_{ij}$ を表します。出力を見ると、すべての隣人が同じ太さで結ばれているのではなく、一部の隣人だけが太い線(大きな注意係数)で結ばれていることがわかります。これこそが GCN との決定的な違いです。GCN なら次数だけで決まる固定値なので、同じ次数の隣人はすべて同じ太さになります。GAT は特徴を見て、分類に効く隣人だけを選択的に強調しているのです。各エッジに表示された数値の合計がほぼ 1 になっている点も、ソフトマックス正規化が効いていることを示しています。

ヘッド数を変えて精度と注意分布を比べる

最後に、マルチヘッドのヘッド数 $K$ を変えると精度と注意の分布がどう変わるかを観察します。ヘッド数は「異なる観点の数」なので、増やすと表現力が上がる一方、Cora のような小規模データでは増やしすぎると過学習する可能性もあります。

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
from torch_geometric.nn import GATConv

def run_with_heads(heads, epochs=200):
    torch.manual_seed(0)
    conv1 = GATConv(dataset.num_features, 8, heads=heads, dropout=0.6)
    conv2 = GATConv(8 * heads, dataset.num_classes, heads=1,
                    concat=False, dropout=0.6)
    params = list(conv1.parameters()) + list(conv2.parameters())
    opt = torch.optim.Adam(params, lr=0.005, weight_decay=5e-4)

    def fwd(training):
        x = F.dropout(data.x, p=0.6, training=training)
        x = F.elu(conv1(x, data.edge_index))
        x = F.dropout(x, p=0.6, training=training)
        return conv2(x, data.edge_index)

    for _ in range(epochs):
        conv1.train(); conv2.train(); opt.zero_grad()
        out = fwd(True)
        loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
        loss.backward(); opt.step()

    conv1.eval(); conv2.eval()
    with torch.no_grad():
        pred = fwd(False).argmax(dim=1)
        acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean().item()
    return acc

heads_list = [1, 2, 4, 8]
accs = [run_with_heads(k) for k in heads_list]

plt.figure(figsize=(7, 4))
plt.plot(heads_list, accs, 'o-', color='darkgreen')
for k, a in zip(heads_list, accs):
    plt.text(k, a + 0.003, f"{a:.3f}", ha='center')
plt.xlabel('number of attention heads K')
plt.ylabel('test accuracy')
plt.title('Effect of multi-head count on Cora accuracy')
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('gat_heads.png', dpi=150, bbox_inches='tight')
plt.show()

このグラフから、ヘッド数を1から増やすとテスト精度がいったん上昇する傾向が読み取れます。1ヘッドだと隣人の見方が一通りに固定されてしまうのに対し、複数ヘッドにすると異なる観点からの集約が組み合わさり、表現力と安定性が増すためです。一方で、ヘッドを増やしすぎても Cora のような小規模グラフでは精度が頭打ちになり、わずかに下がることもあります。これは観点を増やしすぎるとパラメータが過剰になり、訓練データに過適合しやすくなるためです。実務では「8ヘッド程度が良いバランス」という GAT 論文の設定が、まさにこのトレードオフの落とし所であることが体感できます。

まとめ

本記事では、GAT のマルチヘッドアテンション機構を、注意スコアの定義から集約式まで一行ずつ導出し、GCN の固定重みとの違いを明確にした上で、PyTorch Geometric で実装・可視化しました。

  • 注意スコア $e_{ij}$: 共有重み $\bm{W}$ で特徴を射影し、連結して学習ベクトル $\bm{a}$ との内積を取り、LeakyReLU を通す。負の情報を残すために LeakyReLU を使う
  • 注意係数 $\alpha_{ij}$: 隣人にわたるソフトマックスで正規化し、非負かつ合計 1 の確率分布にする
  • マルチヘッド: $K$ 個の独立したアテンションを並列に走らせ、中間層は連結(情報保存)、出力層は平均(次元固定・アンサンブル)で集約する
  • GCN との違い: GCN の重みは次数のみで決まる固定・対称値、GAT の重みは特徴から学習される動的・非対称値。両者ともメッセージパッシングの重みの決め方の違いとして統一的に理解できる
  • 実装: Cora で約82%のテスト精度を達成し、注意係数をエッジの太さで可視化することで「重要な隣人」を解釈でき、ヘッド数を増やすと精度と安定性が向上することを確認

GAT は「集約の重みを学習する」という発想で GNN を一段階押し上げました。この考え方は、後続の Graph Transformer や、より高度なアテンション機構へと発展していきます。次のステップとして、以下の記事も参考にしてください。

Graph Attention Network (GAT) 完全ガイド
GATのアテンション機構の概要・動機・PyTorch Geometricによる実装を入門レベルから解説します。本記事の理論深掘り版を読む前に基礎を確認したい方はこちら。
画像なし
GCN(グラフ畳み込みネットワーク)をわかりやすく解説する
GATが置き換えた固定重み手法・GCNの理論とPyTorch実装を解説。次数正規化重みc_ijの導出からGATとの比較まで本記事の対比として読むと理解が深まります。

参考文献

  • Veličković, P., et al. “Graph Attention Networks.” ICLR 2018.
  • Kipf, T. N., & Welling, M. “Semi-Supervised Classification with Graph Convolutional Networks.” ICLR 2017.
  • PyTorch Geometric 公式ドキュメント(GATConv)