グラフニューラルネット物性予測での外挿性問題についての論文を読む
2022年10月末にarXivにて公開された、"A critical examination of robustness and generalizability of machine learning prediction of materials properties"について読みます。MI界最強と謳われたグラフニューラルネット(GNN)ベースに外挿性難あり？という内容です。メインの主張はAbstractを読めば拾えますが、ちょっとした細かいポイントを下に書いていきます。
ざっとまとめ
State of the artな物性予測モデルALIGNNについて、学習データを2018年までのmaterials project(mp)掲載のデータ、テストデータを2021年の最新データとしたところ、下図のように大きく予測を外してしまいます。これは学習、テストデータの性質に違いがあるためで、UMAPを使ってクラスタリングをするとよくわかります。この外挿性を改善するために、少数(データ全体の~1%程度)でもいいから性質が違うデータを、学習データに追加をすることを著者らは提案しています。また加えるべき性質の違うデータの探し方は、UMAPで他と離れていること、３つのMLモデル（XGB, RandomForest, LinearForest)による予想が互いに大きく異なっていること、を目安に行えると主張しています。
（感想）
データの分布に起因する問題についての論文でした。MIの文脈で別クラスタにある物質の予測困難性については以前から問題視されています(BCS超伝導データからCupratesのTcは予測できないとか、Leave One Cluster Out CVなどがその文脈で提案されてたりとか)。結晶構造や原子といった極めて一般性の高そうなものから作り上げるGNNであっても、データ分布の問題から逃げられないということでしょう。ただ今回のモデルが性能の高い（＝表現力の高い）ALIGNNであったために、より強くこの問題が現れたのみで、GNNの仲間の中でもっと汎化性のあるモデルもあるかもしれません。
この論文から得られる実用的なアイデアとして、なにか新しい物質について予測したいときは、UMAPなどでデータ分布の違いを確認し、分布が大きく違う条件について予め数点実験してデータを得ておく、というひと手順を加えるということでしょう。
実験方法
検証に用いているデータは、LiからBaまでの34種類の元素から構成される合金（Alloys of Interest, AOIと表現）についてのもの。mp2018データにはおよそ2000のAOIに該当するデータがあり、mp2021には7800程度存在する。予測する目的変数は生成エネルギー。
メインの検証モデルは、materials project2018にてpre-trainingをしたALIGNN(Atomistic LIne Graph Neural Network)。ALIGNNは結晶構造における原子間の距離と結合角を考慮したモデル。mp2018のAOIに対する予測誤差はRMSEで0.3 eV/atom。
比較としてRandom Forest(scikit-learn)、Linear Forest(linear-forest)、XGBoostの3モデルを使用。mp2018データに対してMatminerを使って特徴量を274生成し、そのうちPearson's Rが0.7を超えて相関の高いものを削除することで、結果、90の特徴量をもつデータを作成、上記3つのモデルの学習データとしている。
結果と考察
どのモデルも生成エネルギーが1 eV/atomを超えるデータに対して過小評価をしている。また、XGB,LinearForestはALIGNNよりよい結果を示している。
この予測性能低下は、学習データとテストデータの間の、目的変数の範囲の差や、構成する元素、構造のデータ数の差（データの代表性）によってのみ決まっているわけではない。より詳細なデータの性質の差から起因しているようだ。Matminerで作成した90の特徴量についてUMAPでさらに次元削減を行うと、予測性能が高いデータと低いデータがきれいにクラスタリングできている。個々の特徴量で見れば（下図）、特定の特徴量において、学習とテストデータで分布が異なる場合があり、このデータの性質の差が性能低下の原因になっていそうである。
ここでまじでIt is worth noting thatなのが、このきれいなクラスタ分けがPCAではできなかったこと。PCAが線形アルゴリズムであり、特徴量間の潜在的な非線形の関係を解読することが苦手であること、PCAはデータの分散を最大化するような新しい次元を探すが、UMAPのようにデータの局所的なトポロジーを保存しないことが理由ではないかと考察されている。
右図は、値の範囲がほぼ同じなmean CV_VoronoiNNと大きくことなるmean..を軸にデータの分布をプロットしている。縦の値が大きく異なっても、学習データに含まれている範囲のため予測がうまく行っており、また逆はそうではないことから、データ範囲の差が予測性能を大きく左右していることがよくわかる。
提案手法
外挿領域での予測性能を向上させるため、性質の異なるデータを学習データに加えることを提案してる。下図は各種手法でデータを加えることによる性能向上を示しており、全テストデータ5539中の50データ程度を加えるだけでも、極めて性能向上に効果的であることがわかる。加えるデータを選ぶ手法に関しては、もとの学習データから外れたUMAPクラスターに存在するデータから、ランダムに選択するUMAP+random、上記３つのモデルの予測値が互いに大きく異なるものから選択するQuery by Comittie(QBC)、UMAPクラスタリングとQBCを組み合わせたOptimalを提案している。
追記として、もしデータを追加できない状況であるならば、LinearModelのように汎化性の高いモデルを使うことを提案している。また今回のXGBとLinearForestの性能について、AoIのデータ以外からも学習していることによってある程度の汎化性能が得られており、もし学習データをmp2018のうちAoIのデータのみに限定すると、RFの性能並みに低下してしまうことも指摘している。
以上