[논문 리뷰] COEFF-KANs: A Paradigm to Address the Electrolyte Field with KANs

X. Allen Li, Zhuoying Feng외 5명|arXiv (Cornell University)|2024. 07. 24.

Machine Learning in Materials Science인용 수 4

한 줄 요약

이 논문은 MoLFormer에서 확보한 사전 학습된 분자 임베딩을 활용하고, 쿨롱 효율성(Coulombic Efficiency, CE)을 예측하기 위해 코르모고로프-아르놀드 네트워크(Kolmogorov-Arnold Networks, KANs)로 보정하는 두 단계의 딥러닝 프레임워크인 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.

연구 동기 및 목표

• 전해질 조성으로부터 정확하고 자동화된 쿨롱 효율성(Coulombic Efficiency, CE) 예측을 가능하게 하여 고에너지 밀도 리튬 메탈 이차전지 개발의 실험적 부담을 줄이기 위해.
• 기존 기계학습 및 딥러닝 방법의 한계, 즉 특성 공학에서 기인하는 오차 전파와 제한된 데이터 다양성으로 인한 과적합 문제를 해결하기 위해.
• 두 단계의 보정 파라다임에 KANs를 통합하여 전해질 성질 예측의 예측 성능과 모델 해석 가능성을 향상시키기 위해.
• 다른 전해질 성질 및 재료 과학 분야에 적용 가능한 일반화 가능한 프레임워크를 구축하기 위해.

제안 방법

• 대규모 비라벨화된 화학적 SMILES 문자열 코퍼스를 기반으로 MoLFormer을 사전 학습하여 용매 및 염류의 분자 수준 임베딩을 확보한다.
• 성분의 몰 농도 비율에서 유도된 가중치를 사용해 모든 전해질 성분의 토큰 임베딩에 대해 가중 평균 풀링을 수행하여 구성 수준의 표현을 생성한다.
• 결과적으로 도출된 전해질 임베딩을 기반으로, CE 예측을 위한 최종 회귀기로 다층퍼셉트론(Multi-layer Perceptron, MLP) 또는 코르모고로프-아르놀드 네트워크(Kolmogorov-Arnold Network, KAN)를 사용하여 보정한다.
• 보정 단계 동안 학습 효율성과 모델 성능 향상을 위해 CIDO(Curriculum-based Input Data Optimization) 방법을 적용한다.
• 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를 조합한 두 단계 파라다임이 기존 기계학습 및 딥러닝 모델 대비 쿨롱 효율성 예측 정확도를 향상시킬 수 있는가?
• RQ2 하류 헤드에서 기존의 MLP를 KANs로 대체할 경우 전해질 성질 예측의 예측 성능과 모델 해석 가능성에 어떤 향상이 이루어지는가?
• RQ3 MoLFormer 임베딩에 대해 가중 평균 풀링을 적용할 경우, 복잡한 전해질 혼합물의 구성 수준 표현 학습에 얼마나 기여하는가?
• RQ4 제한된 보정 데이터에서 CIDO 방법이 수렴 속도를 가속화하고 성능을 향상시키는 데 얼마나 효과적인가?
• RQ5 제안된 프레임워크는 검증되지 않은 전해질 조성에 대해 일반화가 가능하며, 실제 세계 테스트 세트에서 최고 성능을 기록할 수 있는가?

주요 결과

• COEFF-KAN은 도메인 외 테스트 세트($\mathcal{T}_{\text{out}}$)에서 RMSE 0.110을 기록하였으며, 이는 다음으로 우수한 베이스라인인 MultiModal-MoLFormer 대비 0.085의 감소를 의미한다.
• 모든 테스트 세트에서 기계학습(min. RMSE 0.170) 및 딥러닝(min. 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}}$에서 MLP 대비 최적 성능으로 17.42% 높은 성능을 기록했다.
• 등가도표 분석 결과, COEFF-MLP와 COEFF-KAN의 예측된 CE 값이 실제 값과 매우 밀접하게 일치함을 확인하여, 새로운 데이터에 대한 강건성과 뛰어난 일반화 능력을 입증했다.

Figure 2: Overview of our proposed COEFF framework.