[論文レビュー] COEFF-KANs: A Paradigm to Address the Electrolyte Field with KANs

X. Allen Li, Zhuoying Feng+5 more|arXiv (Cornell University)|Jul 24, 2024

Machine Learning in Materials Science被引用数 4

ひとこと要約

本論文は、MoLFormerから得られる事前学習済み分子埋め込みを活用し、コルモゴロフ＝アルノルドネットワーク（KANs）で微調整することで液体電解質のクーロン効率（CE）を予測する二段階のディープラーニングフレームワーク、COEFF-KANsを提案する。この手法は、未観測のテストデータにおいてRMSEが0.110に達し、既存のモデルを著しく上回り、一般化性能と解釈可能性の両方を向上させた。

ABSTRACT

To reduce the experimental validation workload for chemical researchers and accelerate the design and optimization of high-energy-density lithium metal batteries, we aim to leverage models to automatically predict Coulombic Efficiency (CE) based on the composition of liquid electrolytes. There are mainly two representative paradigms in existing methods: machine learning and deep learning. However, the former requires intelligent input feature selection and reliable computational methods, leading to error propagation from feature estimation to model prediction, while the latter (e.g. MultiModal-MoLFormer) faces challenges of poor predictive performance and overfitting due to limited diversity in augmented data. To tackle these issues, we propose a novel method COEFF (COlumbic EFficiency prediction via Fine-tuned models), which consists of two stages: pre-training a chemical general model and fine-tuning on downstream domain data. Firstly, we adopt the publicly available MoLFormer model to obtain feature vectors for each solvent and salt in the electrolyte. Then, we perform a weighted average of embeddings for each token across all molecules, with weights determined by the respective electrolyte component ratios. Finally, we input the obtained electrolyte features into a Multi-layer Perceptron or Kolmogorov-Arnold Network to predict CE. Experimental results on a real-world dataset demonstrate that our method achieves SOTA for predicting CE compared to all baselines. Data and code used in this work will be made publicly available after the paper is published.

研究の動機と目的

• 電解質の組成からクーロン効率（CE）を正確かつ自動的に予測できる仕組みを提供することで、高エネルギー密度リチウム金属電池の開発における実験的負荷を低減すること。
• 特徴工学による誤差伝搬や、データ多様性が限られる場合の過学習といった、既存の機械学習およびディープラーニング手法の限界を是正すること。
• 二段階の微調整パラダイムにKANsを統合することで、電解質物性予測における予測性能とモデルの解釈可能性を向上させること。
• 他の電解質物性および材料科学分野への応用を想定した汎用性の高いフレームワークを確立すること。

提案手法

• 大規模なラベルなし化学スミルス文字列のコーパスを用いてMoLFormerで化学一般モデルを事前学習し、溶媒および塩の分子レベルの埋め込みを取得する。
• 成分のモル比から導出される重みを用いて、全電解質成分のトークン埋め込みの重み付き平均プーリングを計算し、配合物レベルの表現を生成する。
• 得られた電解質埋め込みを、CE予測用の下流回帰器として、マルチレイヤーパーセプトロン（MLP）またはコルモゴロフ＝アルノルドネットワーク（KAN）で微調整する。
• 微調整中の学習効率とモデル性能を向上させるために、CIDO（カリキュラムベースの入力データ最適化）手法を適用する。
• KANsの学習可能な活性化関数を活用し、特に限られた学習データ条件下でもモデルの解釈可能性と一般化性能を向上させる。
• 実世界のデータセット、すなわちドメイン内（$\mathcal{T}_{\text{in}}$）およびドメイン外（$\mathcal{T}_{\text{out}}$）のテストセットを用いてモデルを訓練および評価し、一般化性能とロバストネスを評価する。

Figure 1: Overview of our method’s two-stage paradigm: pre-training a chemical general model and fine-tuning on downstream domain data.

実験結果

リサーチクエスチョン

• RQ1 事前学習済み分子表現と微調整済みKANsを組み合わせた二段階パラダイムは、既存の機械学習およびディープラーニングモデルと比較して、CE予測精度を向上させることができるか？
• RQ2 下流ヘッドにおける従来のMLPの代わりにKANsを採用することで、電解質物性予測における予測性能とモデルの解釈性が向上するか？
• RQ3 MoLFormer埋め込みに重み付き平均プーリングを適用することで、複雑な電解質混合物の配合物レベルの表現学習がどの程度向上するか？
• RQ4 CIDO手法は、限られた微調整データ条件下で収束を加速させ、性能を向上させるのにどの程度効果的か？
• RQ5 提案フレームワークは、未観測の電解質配合物に一般化可能であり、実世界のテストセットにおいて最先端のモデルを上回る性能を示すか？

主な発見

• COEFF-KANはドメイン外テストセット（$\mathcal{T}_{\text{out}}$）においてRMSEが0.110に達し、次に優れたベースラインであるMultiModal-MoLFormerと比較して0.085の低下を示した。
• すべてのテストセットにおいて、機械学習（最小RMSE 0.170）およびディープラーニング（最小RMSE 0.195）手法を含むすべてのベースラインを上回った。
• COEFF-MLPと比較して、$\mathcal{T}_{\text{out}}$ においてRMSEを0.022削減したため、KANsが複雑で一般化可能な写像を学習する利点が裏付けられた。
• CIDOの使用により、COEFF-MLPのRMSEは0.393から0.132に、COEFF-KANのRMSEは0.388から0.110に低下し、学習効率と性能向上の有効性が確認された。
• KANsは、異なるネットワークの深さと幅においても優れた安定性と性能を示し、$\mathcal{T}_{\text{out}}$ においてMLPsに比べて最適な条件下で17.42%優れた性能を示した。
• パリティプロットでは、COEFF-MLPおよびCOEFF-KANの予測CE値が実測値とよく一致しており、未観測データに対する強力な一般化性能とロバストネスが示された。

Figure 2: Overview of our proposed COEFF framework.