Graph Neural Networkの化学分野への応用

みなさんこんにちは。先進技術セクションでアルバイトをしている三瓶です。
今回は Graph Neural Network (GNN) の化学分野への応用について紹介します。
近年 GNN に関連する論文は増加しており、化学分野への応用を念頭に置いた研究も数多く見受けられる中で、私の専門分野である化学のドメイン知識を用いてこれらを捉え直すことを試みました。

本記事では、まず化学物質(化合物)のグラフ表現とその妥当性を整理し、次にドメイン知識を GNN に組み入れた際の精度改善について紹介します。

初めに

この章では化合物について考える前に、前提となる GNN およびグラフについて説明します。

・Graph Neural Networkとは
Graph Neural Network とは、入力としてグラフを受け取るようなニューラルネットワークのことです。詳しくは弊社データ分析基礎知識”GCNs (Graph Convolutional Networks) とは”をご覧ください。

グラフとは
グラフとは頂点 (node, vertex) と辺 (edge) の集合で構成されるデータ構造であり、辺により頂点間の連結関係が示されます。グラフの頂点や辺には何らかの情報を付加することができます。この情報を重みと呼び、頂点や辺に長さ、コスト等を持たせることが出来ます。例えば、電車の路線図は駅を頂点、路線を辺としたグラフということができます。この時、乗り換えアプリでは路線図の頂点に乗り換え時間を、辺に所要時間や運賃を重みとして置き、出発地から目的地までの重みの合計が最小になるような経路を探索していると捉えることができます。
辺は2つの頂点を結びますが、同一の頂点を結ぶループ構造を持つこともでき、複数の辺が同じ2頂点を結ぶこともできます(多重辺)。また、辺に向きが存在するようなグラフを有向グラフ、向きが存在しないようなグラフを無向グラフと呼びます。
図1に無向グラフを、図2に重み付きの有向グラフを一例として示します。

  図1 無向グラフ              図2 有向グラフ


・グラフの行列表現
ニューラルネットワークに入力するために、グラフを何らかの形で表現する必要があります。この表現としてはよく行列が用いられます。
グラフの情報を表す行列としては次数行列と隣接行列が知られています。
ある頂点について辺を通して結ばれている頂点の個数をその頂点の次数といい、この次数を対角成分とした対角行列を次数行列と呼びます。また、ある頂点と別の頂点について辺で結ばれているか否かを表す行列を隣接行列と呼びます。例えば、図1についての次数行列 (D: Degree matrix) と隣接行列 (A: Adjacency matrix) はそれぞれ次のようになります。

D = \left(\begin{array}{rr} 1 & 0 & 0 & 0 & 0 & 0 \\\\ 0 & 3 & 0 & 0 & 0 & 0 \\\\ 0 & 0 & 1 & 0 & 0 & 0 \\\\ 0 & 0 & 0 & 3 & 0 & 0 \\\\ 0 & 0 & 0 & 0 & 1 & 0 \\\\ 0 & 0 & 0 & 0 & 0 & 1 \end{array}\right), A = \left(\begin{array}{rr} 0 & 1 & 0 & 0 & 0 & 0 \\\\ 1 & 0 & 1 & 1 & 0 & 0 \\\\ 0 & 1 & 0 & 0 & 0 & 0 \\\\ 0 & 1 & 0 & 0 & 1 & 1 \\\\ 0 & 0 & 0 & 1 & 0 & 0 \\\\ 0 & 0 & 0 & 1 & 0 & 0 \end{array}\right)

辺に重みの付いているグラフの時は隣接行列の成分は辺の重みになり、有向グラフの際は行が辺の出て行く頂点、列が辺の入る頂点に対応します。この時、図2の次数行列と隣接行列はそれぞれ次のようになります。

D = \left(\begin{array}{rr} 2 & 0 & 0 & 0 & 0 \\\\ 0 & 3 & 0 & 0 & 0 \\\\ 0 & 0 & 2 & 0 & 0 \\\\ 0 & 0 & 0 & 2 & 0 \\\\ 0 & 0 & 0 & 0 & 1 \end{array}\right), A = \left(\begin{array}{rr} 0 & 3 & 0 & 2 & 0 \\\\ 0 & 0 & 5 & 0 & 7 \\\\ 0 & 0 & 0 & 0 & 0 \\\\ 0 & 0 & 4 & 0 & 0 \\\\ 0 & 0 & 0 & 0 & 0 \end{array}\right)


化合物のグラフ表現とその妥当性

前章では GNN とグラフ及びその行列表現について説明しました。
この章では GNN の入力として化合物の構造を考え、化合物をグラフとする際にどの様な特徴が現れるかを考察します。

・化合物のグラフ化
そもそも、化合物をグラフとして扱うためには化合物とグラフの変換ができる必要があります。化合物をコンピュータで扱う際には SMILES 記法により化合物を表し、文字列として扱うことが一般的です。SMILES 記法は化合物の構造を記述するための記法で、1)異性体を区別できる表記方法を持っている 2)芳香環は単結合と二重結合の組み合わせとは区別される 3)水素原子は原子価の足りない原子に暗黙的に結合していると考え明記しない と言った特徴があります。GNN への入力は一般に SMILES 記法をグラフ化することで作成します。

・化合物とグラフの対応関係
化合物は原子を頂点、結合を辺とするグラフで表すことができます。この時、化学的な性質を考慮すると、頂点および辺は以下の様な表現となると考えられます。
まず、頂点(原子)について考えます。この時、原子番号は頂点の種類と捉えるべきだと考えられます。原子の性質は原子番号と対応する一方、比例するものではありません。よって、原子番号を one-hot vector のように扱う方が頂点の重みと捉えるよりも化合物の特徴を反映させやすくなると考えられます。
次に、辺(結合)について考えます。結合は方向性もループも持たないのでグラフはループのない無向グラフであることがわかります。また、結合次数は辺の状態としてグラフに組み入れる必要があります。SMILES 記法では水素原子を省略するため、原子の種類と原子同士の結合の有無のみでは化合物を一意に定めることができません。例として、図3に結合次数を考慮しない窒素とヒドラジンを示します。
図3 結合次数を考慮しない窒素とヒドラジン
辺の状態の表し方としては1)結合次数を多重辺として表す 2)結合次数を辺の重みとして表す 3)結合次数を辺の種類として表す の3種の方法が考えられます。しかし、1)については芳香環の結合を整数本の多重辺として表すことができないという欠点があります。芳香環の結合は単結合と二重結合の中間的な結合であることが知られているため、1.5重結合の様な小数を用いた結合次数とすることが適当と考えられるためです。2)については結合のエネルギーや結合距離等の結合次数と相関のある特徴量を暗に含有できるという利点がある一方、芳香族に特有な置換反応の起こりやすさの様に結合次数に相関のない性質を表せないという欠点があります。3)については結合次数を区別して学習を進めるため特徴量をその傾向によらず無理なく抽出できる一方、辺の種類ごとの演算によりパラメータや計算量が増えることが懸念されます。実用の上ではタスクやモデルに応じて2)と3)を使い分けるべきだと考えられます。
最後に、グラフ全体について考えます。グラフ構造は頂点間の繋がり方を表す構造ですので、結合(=辺)の存在しない部分の情報は保持されません。例えば幾何異性体や立体異性体はグラフ化すると同じグラフとなってしまうので判別不可能となります(立体異性は辺の有無や結合次数は等しく、複数の原子の立体的な配置が異なることにより起こります。頂点の立体的な配置の情報はグラフでは保持されません。) 。これらの区別を行うにはグラフ内で既存の頂点と辺の関係を崩さずに異性体の区別をつけられる機構を導入する必要があると考えられます。それに加えて、原子間に結合のない Zeise 塩のような化合物はグラフ化できないため、通常グラフで表せないような結合を表すことのできる機構も導入する必要があると考えられます。

既存研究やその実装では化学のドメイン知識が一部使われているものの、まだ完全にドメイン知識が駆使されているとは言い難いものが多く存在しています。以上のようなドメイン知識を用いてグラフやモデルの改良を行うことで GNN はより高い精度を出せる可能性があると考えられます。

実験

前章では化合物をグラフとした際のグラフの特徴や辺と頂点の表し方、グラフ化で失われる情報を考えました。
この章では前章の考察を基にドメイン知識の有用性を実験にて確かめます。GNN の1つである Neural Finger Print (NFP) というモデルを実装し、NFP に改良を加えたものと精度の比較を行います。

・Neural Finger Print(NFP) とは
今回、私は Neural Finger Print(NFP)[1] という手法に注目し、公開されている既存実装に改良を加えました。既存実装としては Chainer Chemistry のものを用いました。
NFP は官能基の存在の有無により化合物同士の類似度を判定する手法である Circular fingerprints[2] を微分可能に一般化し、ニューラルネットワークを導入したもので、シンプルで理解しやすいアルゴリズムです。
Circular fingerprints は化合物中の各原子について隣接する原子の情報を集めることを繰り返すことで検索対象の官能基が存在するか否かを判別する手法です。
図4 Circular fingerprintsの模式図(カルボキシ基の炭素を例に)
NFP においても、ある原子について隣接した原子およびその原子自身の情報の足し合わせを繰り返すことで近傍の原子の情報を集めます。
図5 水素を考慮しないジエチルエーテルの酸素におけるNFPの模式図
このような情報の集約を全原子について行い、それらの情報から予測値(特徴量)を数値ベクトルとして得ます。
ベースである Circular fingerprints では官能基の有無を予測しますが、一方で NFP では特徴量を予測するため、回帰や分類などの様々なタスクに用いることができます。

・NFP のアルゴリズム
NFP の具体的なアルゴリズムは図6に示す通りです。f, r, v, i はベクトル、H, W は学習可能な行列です。また、青字は Circular fingerprints との変更点を示します。
図6 NFPのアルゴリズム[1]
ここで、H は (atom features の個数 × atom features の個数 ) という形の行列、Wは (atom features の個数 × fingerprint vectorの次元数 ) という形の行列です。また、H は (radiusの大きさ) × (次数の最大値(図6では5としています。)) 個、W は (radiusの大きさ) 個用意します。

・NFP アルゴリズムの2つの解釈
図6のアルゴリズムの7番目の”原子 a に隣接する原子の数 N”の数え方は1)原子 a と繋がっている原子の個数 2)原子 a が辺を通して繋がっている頂点の数 という2通りの解釈が可能です。この2つの解釈は一見同じに見えますが、例えば二酸化炭素の炭素原子を原子 a として考えたときに1)では2つの酸素原子が炭素原子に繋がっているので N は2つとなりますが、2)では各炭素-酸素間の結合は二重結合で多重辺となっており、炭素原子は各酸素原子と2つの辺を通して繋がっているので、N は4つとなります。
このように多重辺が存在するときに2つの解釈には差が生まれますが、SMILES 記法においては水素原子を明記せず、特徴量の計算でも水素は用いないため1)では単結合の時と区別がつかなくなってしまいます(二酸化炭素に対してメタンジオールは同じ予測値が得られます)。
一方、2)の場合も二酸化炭素とオルト炭酸の区別がつかないように思われます。しかし、層を二層以上にすると異なる予測値が得られます。例として二層の時を考えます。最終層ではどちらの炭素原子にも酸素原子4つの特徴量ベクトルが足されますが、一層手前で1つの酸素原子に足される特徴量ベクトルがメタンジオールでは炭素原子1つ分、二酸化炭素では炭素原子2つ分となるため、最終層で炭素に足される炭素原子の特徴量ベクトルの数は異なります。この違いについて、簡単のために二酸化炭素の片方の酸素のみに対しての違いを模式化した図を図7に示します。
図7 隣接する原子数Nの解釈の違いと区別の有無
このように1)で区別できない多重辺の有無も十分な層の深さを確保した下で2)を使用することにより区別が可能となります。
NFP の論文中でははっきりとどちらの解釈をとっているかは述べられておらず、Chainer Chemistry での NFP の実装(以下、Chainer Chemistry 版と呼びます。)では1)の解釈を採用しています。ですが、上に述べたように1)では多重辺の識別ができなくなってしまいます。そこで、2)を使用した NFP (以下、改良版と呼びます。) を実装し、2つのモデルの精度を比較することでドメイン知識に基づいてアルゴリズムの解釈を変えた際の予測精度への影響を実験、考察しました。

・NFPの実装
実装した NFP での Chainer Chemistry 版と改良版の変更点はアルゴリズムの解釈です。具体的な実装の内容及び違いとしては以下の様になります。Chainer Chemistry における NFP の実装はグラフの行列表現を入力として受け取るため、この行列表現部分を変更しました。
Chainer Chemistry 版の NFP では隣接する原子及び原子自身の特徴量ベクトルの足し合わせを行う vrai (図6のアルゴリズムの8番目)において、化合物を重みのないグラフと見たときの隣接行列の対角成分を全て1とした行列(以下 A’ と置きます。)と全原子の特徴量ベクトルを集めた行列(以下 R’ と置きます。)の積 (R’A’) を取っています。
ここで、例として簡単のために二酸化炭素について考えます。各行列の1,3列目が酸素を、2列目が炭素を表す (隣接行列の行についても同様の対応関係があります。) とすると R’ と A’ 及びそれらの積は次の様になります。

R' = \left(\begin{array}{rr} \bf r_1 & \bf r_2 & \bf r_3 \end{array}\right), A' = \left(\begin{array}{rr} 1 & 1 & 0 \\\\ 1 & 1 & 1 \\\\ 0 & 1 & 1 \end{array}\right)

R'A' = \left(\begin{array}{rr} \bf r_1 + \bf r_2 & \bf r_1 + \bf r_2 + \bf r_3 & \bf r_2 + \bf r_3 \end{array}\right)

上の結果より、R’A’の各原子に対応する列ベクトルに注目すれば、確かに隣接原子及び原子自身について、特徴量ベクトルの和が取れていることが分かります。
一方で改良版では、化合物とグラフの対応関係の辺についての考察の “2) 結合次数を辺の重みとして表す” を採用し、隣接行列の重みに結合次数を用いることでアルゴリズムの解釈の “2) 原子 a が辺を通して繋がっている頂点の数” を実装しました。この様な隣接行列を Cheiner Chemistry 版の隣接行列の代わりに用いました。Chainer Chemistry 版と同様に改良版の R’, A’, R’A’ を考えると以下の様になります。

R' = \left(\begin{array}{rr} \bf r_1 & \bf r_2 & \bf r_3 \end{array}\right), A' = \left(\begin{array}{rr} 1 & 2 & 0 \\\\ 2 & 1 & 2 \\\\ 0 & 2 & 1 \end{array}\right)

R'A' = \left(\begin{array}{rr} \bf r_1 + 2\bf r_2 & \bf 2 \bf r_1 + \bf r_2 + 2\bf r_3 & \bf 2\bf r_2 + \bf r_3 \end{array}\right)

上の結果より、R’A’ の各原子に対応する列ベクトルに注目すれば、確かに隣接原子及び原子自身の特徴量ベクトルについて、結合次数に応じた重み付きの和が取れていることが分かります。

・実験方法
本実験は複数のモデルにおける精度の比較がなされている Graph Warp Module(GWM)[3] の論文と同じ条件で行いました。この論文では NFP を含む GNN モデルの実装として Chainer Chemistry を用いており、この論文の結果と比較することで Chainer Chemistry 版の再現実験の妥当性を保証することができます。
実験条件は次の通りです。
  • データセット: QM9[4,5]
  • 使用ラベル: 計12種類[5]
    • mu; Dipole moment (D)
    • alpha; Isotropic polarizability (a03)
    • homo; Energy of Highest Occupied Molecular Orbital (Ha)
    • lumo; Energy of Lowest Unoccupied Molecular Orbital (Ha)
    • Gap(lumo-homo) (Ha)
    • r2; Electronic spatial extent (a02)
    • zpve; Zero point vibrational energy (Ha)
    • U0; Internal energy at 0 K (Ha)
    • U; Internal energy at 298.15 K (Ha)
    • H; Enthalpy at 298.15 K (Ha)
    • G; Free energy at 298.15 K(Ha)
    • Cv; Heat capacity (cal mol-1 K-1)
  • 損失関数: Mean Squared Error (MSE)
  • Optimizer: Adam (α=0.001, β1=0.9, β2=0.999) [6]
  • 層の数(R): 5
  • 特徴ベクトルの次元: 86
  • train, validation, testの分割: scaffold split (Chainer Chemistry の ScaffoldSplitter 使用)、0.8 : 0.1 : 0.1
GWM の論文中に記載のなかったバッチサイズについては Optuna[7] を使用し、train データと validation データを用いて探索を行いました。その結果、バッチサイズは Chainer Chemistry 版では136、改良版でも136と決まりました。
学習では train データと validation データを合わせたものを train データとして扱い、QM9 データセットの SMILES、A、B、C を除く12種類の使用ラベルの物性値を 50epoch 学習させました。
モデルの評価はテストデータについての Mean Absolute Error (MAE) を用いました。

・結果
GWM 論文における NFP の MAE は 6.16 でした。弊社で訓練した Chainer Chamistry 版の MAE は 6.11 でした。一方で改良版の MAE は 5.09 でした。
この結果より、GWM 論文と Chainer Chamistry 版の結果が一致することを確認した上で、化学のドメイン知識を用いたアルゴリズムの改良によって MAE が改善することを示すことができました。

まとめ

本記事では化学のドメイン知識を用いて、物性予測における GNN の性能向上ができないかを考えました。化合物をグラフで表現する際に元の情報が失われている点に着目し、結合次数の違い(多重辺)を考慮した改良を Neural Finger Print (NFP) に加え、結果としてオリジナルの手法よりも精度の高い予測値を得ることができました。

化学的な観点から考えると、多重辺、すなわち結合次数の識別は物性予測において非常に重要な情報です。例としてエタンとエチレンの違いを考えます。エチレンはエタンの炭素原子同士の結合を単結合から二重結合に変えた化合物ですが、分子の構造はエタンとは大きく異なります。エチレンでは二重結合の存在によって構造の立体性が失われ平面構造になります。平面構造の上下に電子雲の存在する二重結合は単結合に比べて反応性(本実験の homo, lumo, gap に関係があります。) が高くなるため、エチレンとエタンは反応性が異なってきます。他にも、窒素原子同士の間に三重結合の存在する窒素は他の化合物と反応しにくい不活性ガスとして知られており、食品を新鮮に保つためのパッキング等に用いられていますが、一方で窒素原子同士を単結合でつなぎ、各窒素に2つの水素原子が結合したヒドラジンは反応性が高く、ロケットエンジンの推進剤に用いられます。結合次数が変わり水素原子との結合が増減する (結合次数を考慮しないと同じグラフになります。) だけで化合物の反応性が変わることは多々あることです。また、二重結合が単結合と交互に存在する共役長が伸びると化合物の吸収波長 (本実験の物性値には存在しません。) がシフトし可視光吸収を示す場合があり、これは化合物の色と関係します。このように結合次数が変わるだけで変化する物性値は数多くあり、結合次数の識別ができるようになったということはこれらの物性値の違いも識別できるようになったと言えます。

今後の展望

現在、GNN のアーキテクチャや訓練テクニックについては活発な研究がなされていますが、一方で GNN に入力する際の化合物の適切なグラフ表現に関しては目立って研究がなされているようには見えませんでした。

今回紹介した結合次数以外にも、結合としては明示的に書かれない水素結合などの分子間力や立体的な配置など、物性や反応性に大きく影響する要素は数多く存在します。グラフは化合物の構造をシンプルに表しますが、それによって失われてしまう情報もたくさんあるのです。

今後はモデル構造や訓練方法だけでなく、化合物の情報を失わないようにグラフ化する、もしくは失われた情報を再度組み込むなど、よりいっそう化合物の扱いに特化した GNN 手法が出てくるかもしれません。


ALBERTでは、様々な専門的バックグラウンドを持つデータサイエンティストを募集しています。詳しくは採用ページ(https://www.albert2005.co.jp/saiyo/?place=blog_text)をご覧ください。

参考文献
[1] David Duvenaud, Dougal Maclaurin, Jorge Aguilera-Iparraguirre, Rafael Gómez-Bombarelli, Timothy Hirzel, Alán Aspuru-Guzik, and Ryan P. Adams. Convolutional Networks on Graphs for Learning Molecular Fingerprints. https://papers.nips.cc/paper/5954-convolutional-networks-on-graphs-for-learning-molecular-fingerprints.pdf (2019/10/30 参照).
[2] Glen, Robert & Bender, Andreas & Hasselgren, Catrin & Carlsson, L & Boyer, S & Smith, James. (2006). Circular fingerprints: Flexible molecular descriptors with applications from physical chemistry to ADME (vol 9, pg 199, 2006). IDrugs: the investigational drugs journal. 9. 311-311.
[3] Katsuhiko Ishiguro, Shin-ichi Maeda, and Masanori Koyama. Graph Warp Module: an Auxiliary Module for Boosting the Power of Graph Neural Network in Molecular Graph Analysis. https://arxiv.org/pdf/1902.01020.pdf (2019/10/30 参照).
[4] L. Ruddigkeit, R. van Deursen, L. C. Blum, J.-L. Reymond, Enumeration of 166 billion organic small molecules in the chemical universe database GDB-17, J. Chem. Inf. Model. 52, 2864–2875, 2012.
[5] R. Ramakrishnan, P. O. Dral, M. Rupp, O. A. von Lilienfeld, Quantum chemistry structures and properties of 134 kilo molecules, Scientific Data 1, 140022, 2014.
[6] Kingma, Diederik & Ba, Jimmy. (2014). Adam: A Method for Stochastic Optimization. International Conference on Learning Representations.
[7] Preferred Networks. Optuna: A hyperparameter optimization framework. github.com/pfnet/optuna, 2018.
[8] Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How Powerful are Graph Neural Networks? https://arxiv.org/pdf/1810.00826.pdf (2019/11/27参照).