Unifying Force Prediction and Molecular Conformation Generation Through Representation Alignment

Lucas Pinede, Soojung Yang, Juno Nam, Rafael Gomez-Bombarelli

Keywords: Boltzmann emulator, MLIP, Machine Learning Interatomic Potential, Representation Alignment

Abstract: Sampling equilibrium conformational ensemble is essential for understanding biomolecular functions. Although i.i.d. ensemble samplers (i.e., Boltzmann emulators) hold great promise, their training strategies have been underexplored. We hypothesize that machine learning interatomic potentials (MLIPs), trained on abundant non-equilibrium data, already capture geometric and energetic information needed for ensemble generation. Inspired by Yu et al. (ICLR'25), we introduce a representation alignment loss that regularizes the emulator's hidden states to be similar to those of pretrained MLIP. This simple addition only requires a one-time inference cost of samples yet shortens the training time by 1.5$\times$ and yields better sampling performances.

Submission Number: 137