[Paper Review] 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 Science4 citations

TL;DR

This paper proposes COEFF-KANs, a two-stage deep learning framework that leverages pre-trained molecular embeddings from MoLFormer and fine-tunes them with Kolmogorov-Arnold Networks (KANs) to predict Coulombic Efficiency (CE) of liquid electrolytes. The method achieves state-of-the-art performance with an RMSE of 0.110 on unseen test data, significantly outperforming existing models and demonstrating improved generalization and interpretability.

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.

Motivation & Objective

• To reduce experimental workload in high-energy-density lithium metal battery development by enabling accurate, automated prediction of Coulombic Efficiency (CE) from electrolyte composition.
• To address limitations in existing machine learning and deep learning methods, including error propagation from feature engineering and overfitting due to limited data diversity.
• To improve predictive performance and model interpretability in electrolyte property prediction by integrating KANs into a two-stage fine-tuning paradigm.
• To establish a generalizable framework applicable to other electrolyte properties and materials science domains.

Proposed method

• Pre-train a chemical general model using MoLFormer on a large corpus of unlabeled chemical SMILES strings to obtain molecular-level embeddings for solvents and salts.
• Compute weighted average pooling of token embeddings across all electrolyte components, with weights derived from component molar ratios, to generate formulation-level representations.
• Fine-tune the resulting electrolyte embeddings using either a Multi-layer Perceptron (MLP) or a Kolmogorov-Arnold Network (KAN) as the downstream regressor for CE prediction.
• Apply the CIDO (Curriculum-based Input Data Optimization) method to enhance training efficiency and model performance during fine-tuning.
• Utilize KANs' learnable activation functions to improve model interpretability and generalization, especially with limited training data.
• Train and evaluate models on real-world datasets, including in-domain ($\mathcal{T}_{\text{in}}$) and out-of-domain ($\mathcal{T}_{\text{out}}$) test sets, to assess generalization and robustness.

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

Experimental results

Research questions

• RQ1 Can a two-stage paradigm combining pre-trained molecular representations with fine-tuned KANs improve CE prediction accuracy compared to existing machine learning and deep learning models?
• RQ2 Does replacing traditional MLPs with KANs in the downstream head enhance predictive performance and model interpretability in electrolyte property prediction?
• RQ3 To what extent does the use of weighted average pooling on MoLFormer embeddings improve formulation-level representation learning for complex electrolyte mixtures?
• RQ4 How effective is the CIDO method in accelerating convergence and improving performance with limited fine-tuning data?
• RQ5 Can the proposed framework generalize to unseen electrolyte formulations and outperform state-of-the-art models on real-world test sets?

Key findings

• COEFF-KAN achieves an RMSE of 0.110 on the out-of-domain test set ($\mathcal{T}_{\text{out}}$), representing a 0.085 reduction compared to the next-best baseline, MultiModal-MoLFormer.
• The model outperforms all baselines, including machine learning (min. RMSE 0.170) and deep learning (min. RMSE 0.195) methods, across all test sets.
• COEFF-KAN reduces RMSE by 0.022 compared to COEFF-MLP on $\mathcal{T}_{\text{out}}$, demonstrating the advantage of KANs in learning complex, generalizable mappings.
• The use of CIDO reduces the RMSE of COEFF-MLP from 0.393 to 0.132 and COEFF-KAN from 0.388 to 0.110, confirming its effectiveness in improving training efficiency and performance.
• KANs show superior stability and performance across different network depths and widths, with optimal results at 17.42% better performance than MLPs on $\mathcal{T}_{\text{out}}$.
• Parity plots show that predicted CE values for COEFF-MLP and COEFF-KAN are closely aligned with actual values, indicating strong generalization and robustness on unseen data.

Figure 2: Overview of our proposed COEFF framework.