[论文解读] 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

一句话总结

该论文提出COEFF-KANs，一种两阶段深度学习框架，利用MoLFormer预训练的分子嵌入，并通过Kolmogorov-Arnold网络（KANs）进行微调，以预测液态电解质的库仑效率（CE）。该方法在未见测试数据上实现了0.110的RMSE，性能达到最先进水平，显著优于现有模型，并展现出改进的泛化能力和可解释性。

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整合到两阶段微调范式中，提升电解质性质预测的预测性能和模型可解释性。
• 建立一个可推广的框架，适用于其他电解质性质和材料科学领域。

提出的方法

• 在大规模未标注化学SMILES字符串语料库上使用MoLFormer预训练一个化学通用模型，以获得溶剂和盐的分子级嵌入。
• 通过基于组分摩尔比的加权平均池化方法，对所有电解质组分的标记嵌入进行池化，生成配方级表示。
• 使用多层感知机（MLP）或Kolmogorov-Arnold网络（KAN）作为下游回归器，对生成的电解质嵌入进行微调，以实现CE预测。
• 在微调过程中应用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 在下游头部用KANs替代传统MLP，是否能提升电解质性质预测的预测性能和模型可解释性？
• RQ3 在MoLFormer嵌入上使用加权平均池化，能在多大程度上改善复杂电解质混合物的配方级表征学习？
• RQ4 CIDO方法在有限微调数据下，是否能有效加速收敛并提升性能？
• RQ5 所提出的框架是否能泛化到未见过的电解质配方，并在真实世界测试集上优于最先进模型？

主要发现

• COEFF-KAN在域外测试集（$\mathcal{T}_{\text{out}}$）上实现了0.110的RMSE，相比表现最佳的基线模型MultiModal-MoLFormer降低了0.085。
• 该模型在所有测试集上均优于所有基线模型，包括机器学习方法（最小RMSE 0.170）和深度学习方法（最小RMSE 0.195）。
• 与COEFF-MLP相比，COEFF-KAN在$\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.