はじめに
Nishikaさん主催の材料コンペに参加し、2位に入賞することができました。
今回はその解法についての記事になります。
データや評価指標などはコンペのページをご参照ください。
前処理
まず、Formation energy per atom の前処理をしました。下図のように、何も前処理をしていない状態でも比較的きれいな分布ですが、より単峰性のデータに近づけるための変換をしました。
学習データに対し、構造単位のFormation energyに変換します。この値に対し、原子数で線形回帰をしたモデルを作成します。構造単位のFormation energyからモデルでの予測値を引き、原子単位に戻したものを学習のターゲットとしました。ざっくりとしたコードは以下です。
from sklearn.linear_model import LinearRegression
df["formation_energy"] = df["formation_energy_per_atom"] * df["nsites"]
df_elem = df.iloc[:, 1:90]
all_zero = df_elem.sum() == 0
print(f"not exists: {all_zero.index[all_zero]}")
df_elem.drop(columns=all_zero.index[all_zero].tolist(), inplace=True)
elements = df_elem.columns
X = df_elem.values
lin = LinearRegression(fit_intercept=False)
lin.fit(X, df["formation_energy"])
df["noralized_formation_energy_per_atom"] = (df["formation_energy"] - lin.predict(X)) / df["nsites"]
ちなみにAr, He, Neはデータが無かったので平均値で埋めています。
変換前後の分布は以下のようになります。
モデル
GNNのアーキテクチャとしてNequIPを採用しました。これは球面調和関数のテンソル積で特徴量をアップデートするモデルで、角度などの空間情報を持つことができます。モデル全体としてE(3)-equivariantになっており、内部でリッチな特徴量が形成されます。
NequIPをベースにして、いくつかの変更を行いました。ただし、後で述べるように、いくつかはあまり効果が無く、最終的には採用していません。
- 原子のembeddingをmat2vecから得る
- 各レイヤーで出力された特徴量と、最終レイヤーの特徴量を組み合わせる（gatherd feat）
- ハンドクラフトな特徴量（global feat）をEnergy headの入力に追加する
- 原子に働くForceを算出し、事前学習に用いる
全体としては以下のようになります。
InteractionはNequIPのモジュールをそのまま使用しています。
mat2vec
mat2vecは論文データを学習したWord2Vecモデルです。例えば"Mg"を固定長のdenseなベクトルに変換できます。ある程度元素間の関係を学習できていると考えられるため、これを各ノードのembeddingとしました。ちなみに、通常のNequIPでは単純にOne-Hotで変換しています。
学習時はmat2vecのembedding層は固定し、後段のLinear層でGNNに入力するノード特徴量を取得しています。
gatherd feat
Interactionレイヤーは複数重ねますが、前段のレイヤーでは近距離、後段のレイヤーでは長距離の相互作用が支配的であると思い、全レイヤーの出力を結合したモデルを作成しました。
global feat
GNNだけでは格子の体積の情報などを明示的に取り入れていないため、追加で情報を与えた方が良いかと考え、Energy headの入力に任意の特徴量を追加できるアーキテクチャにしました。試した特徴量は体積関連、電荷関連の特徴量です。
forces
今回のコンペはDFTで構造緩和後の構造を対象としているため、各原子に働くForceはほぼ0であるべきです。これを制約条件として（Forceを0に近づけるロス関数を設定して）事前学習を行いました。テストデータも含めた事前学習は許可されているか分からなかったので、訓練データのみで行っています。
なお、energy, forcesのマルチタスクとしても解くことができますが、事前学習しか試していません。
学習
学習はCosineAnnealingのスケジューラで学習率を徐々に下げながら、80-100エポック程度行いました。
最初はoptimizerをAdamに固定していましたが、コンペ中にGoogleから出たLionが良さそうだったので使ってみました。
Lion
LionはGoogleが論文発表したoptimizerで、AdamWをベースにプログラム自動探索で改良が施されています。sign関数で更新幅を固定長にすることがポイントのようです。多くのタスクでAdamWよりも良い汎化性能が得られると報告されています。
このコンペに関してもAdamより汎化性能が向上しました。ただし、Adamよりもパラメータの調整がシビアな印象で、Adamもパラメータをもう少し調整すれば同程度の汎化性能になっていた可能性もあります。
今回のケースでは、バッチサイズを大きく（64以上）、学習率を小さく（Adamの1/20くらい）し、勾配クリッピングをかけることでようやく安定して学習が行えました。warmupの有無はあまり安定性に寄与しませんでした。
精度検証
今回は時間と計算資源の都合上、1つのfoldのvalidationスコアのみで簡易的に精度比較しています。LBでのスコアで比較していませんのでご注意ください。
上記の1-4の変更に関して、2,3はあまり効果が無さそうでした。2はNequIP自身がskip connectionを持っており、特に必要なかったのかなと考えています。3は特徴量をしっかり作りこめば精度は改善する気がしますが、今回作った特徴量では改善は見られませんでした。
4に関しては、効果はありそうだったものの、学習時に1ステップ当たり2回のbackwardが発生し、計算コストが2倍以上になってしまいます。同じ追加時間でも、単純にアンサンブルを多くした方が効果的のように感じたので、実際は使用していません。
結果的に、追加部分は1だけなので、かなりシンプルなアーキテクチャで落ち着きました。
提出モデル
10-foldで学習を行い、それぞれのfoldでのモデルを単純にアンサンブルして予測値を出しました。
最後に
コンペ参加は初めてだったのですが、他の方の解法など大変勉強になりました。Nishikaさん、コンペの開催ありがとうございました。