Mixup-Augmented Meta-Learning for Sample-Efficient Fine-Tuning of Protein Simulators

TL;DR

This paper introduces a novel meta-learning approach with data augmentation for efficient and accurate molecular dynamics simulations, capable of generalizing to unseen conditions with limited data.

Contribution

It proposes a two-stage framework combining data mixing and meta-learning for soft prompt-based molecular simulation, enhancing generalization and sample-efficiency.

Findings

Outperforms existing methods in accuracy on in-domain data

Demonstrates strong generalization to unseen and out-of-distribution scenarios

Improves sample-efficiency in fine-tuning with meta-learning

Abstract

Molecular dynamics simulations have emerged as a fundamental instrument for studying biomolecules. At the same time, it is desirable to perform simulations of a collection of particles under various conditions in which the molecules can fluctuate. In this paper, we explore and adapt the soft prompt-based learning method to molecular dynamics tasks. Our model can remarkably generalize to unseen and out-of-distribution scenarios with limited training data. While our work focuses on temperature as a test case, the versatility of our approach allows for efficient simulation through any continuous dynamic conditions, such as pressure and volumes. Our framework has two stages: 1) Pre-trains with data mixing technique, augments molecular structure data and temperature prompts, then applies a curriculum learning method by increasing the ratio of them smoothly. 2) Meta-learning-based fine-tuning framework improves sample-efficiency of fine-tuning process and gives the soft prompt-tuning better initialization points. Comprehensive experiments reveal that our framework excels in accuracy for in-domain data and demonstrates strong generalization capabilities for unseen and out-of-distribution samples.