Knowledge Distillation of a Protein Language Model Yields a Foundational Implicit Solvent Model

TL;DR

The paper addresses the longstanding accuracy gap in implicit solvent models by distilling evolutionary information from a protein language model into a GNN-based potential. It trains Schake to reproduce SS8 motif likelihoods predicted by ESM3 (teacher) and couples its energy with a GB electrostatics term to yield a transferable implicit solvent model that captures folding free-energy landscapes and IDP ensembles. Key contributions include (i) a data-efficient distillation pipeline that transfers solvent-sensitive structural preferences, (ii) a one-state energy E_GNN^{os} that stabilizes native basins and a multi-state energy E_GNN^{MS} that balances folded and disordered states, and (iii) demonstration of stable ML/MD across 11 proteins totaling 6.8 microseconds and improved IDP behavior compared with conventional ISMs. The approach offers a scalable foundation ISM that can accelerate large-scale protein simulations and guide future data-driven ISM development.

Abstract

Implicit solvent models (ISMs) promise to deliver the accuracy of explicit solvent simulations at a fraction of the computational cost. However, despite decades of development, their accuracy has remained insufficient for many critical applications, particularly for simulating protein folding and the behavior of intrinsically disordered proteins. Developing a transferable, data-driven ISM that overcomes the limitations of traditional analytical formulas remains a central challenge in computational chemistry. Here we address this challenge by introducing a novel strategy that distills the evolutionary information learned by a protein language model, ESM3, into a computationally efficient graph neural network (GNN). We show that this GNN potential, trained on effective energies from ESM3, is robust enough to drive stable, long-timescale molecular dynamics simulations. When combined with a standard electrostatics term, our hybrid model accurately reproduces protein folding free-energy landscapes and predicts the structural ensembles of intrinsically disordered proteins. This approach yields a single, unified model that is transferable across both folded and disordered protein states, resolving a long-standing limitation of conventional ISMs. By successfully distilling evolutionary knowledge into a physical potential, our work delivers a foundational implicit solvent model poised to accelerate the development of predictive, large-scale simulation tools.

Figures (7)

• Figure 1: Knowledge distillation enables training of an efficient, multiscale GNN architecture. (A) Capped Ala-Val dipeptide with the C$\alpha$, C, and N backbone atoms highlighted. These atoms serve as inputs to the GNN. The color of the left half of the circle indicates the backbone atom type, while the color on the right half indicates the amino acid type. (B) Diagram of the Schake multiscale message-passing scheme. The short-ranged SAKE message-passing layer acts on all backbone C$\alpha$, C, and N atoms within the cutoff distance $r_\mathrm{cut}^\mathrm{SAKE} = 1$ nm, while the long-ranged SchNet message-passing layers acts on only C$\alpha$ atoms within $r_\mathrm{cut}^\mathrm{SchNet} = 2.5$ nm. (C) Diagram of the knowledge distillation training methodology. By inputting an amino acid sequence, ESM3 is used to predict a matrix of SS8 motif likelihoods. A GNN that inputs both an amino acid sequence and a structure is then trained to match the ESM3-predicted matrix of SS8 motif likelihoods using the cross entropy loss function. The sequence and structure of chignolin CLN025 is shown as an example, with the likelihoods for the folded state motifs shown in the matrices.
• Figure 2: Schake matches the SS8 motif predictions from ESM3-open. (A) Confusion matrix for ESM3-open SS8 motif predictions. (B) Architecture of the Schake GNN. Backbone atoms are inputted, and SS8 motif logits are outputted. Refer to the Methods section for a detailed description of the Schake architecture. (C) Confusion matrix for GNN SS8 motif predictions. For (A) and (C), the mean predicted likelihood for each SS8 motif conditioned on the true SS8 motif is displayed. The percentage in the upper right corner indicates the mean predicted likelihood of the true SS8 motif (computed by averaging along the diagonal). Each letter corresponds to a particular SS8 motif: G to $3_{10}$-helices, H to $\alpha$-helices, I to $\pi$-helices, T to hydrogen-bonded turns, E to $\beta$-sheets, B to $\beta$-bridges, S to non-hydrogen-bonded bends, and C to non-categorized structures.
• Figure 3: GNN energy changes closely correspond to structural changes.$E_\mathrm{GNN}^\mathrm{os}$ was computed across long time-scale MD trajectories generated from a previous work.lindorff-larsen_how_2011 Each trajectory was subset following the procedure used in a previous GNN study.airas_transferable_2023 Results for the four largest proteins are shown here, while results for the other eight proteins are shown in the Fig. S2. RMSD from the folded state (reference structures are detailed in the Methods section) is shown in red on the left y-axis, while $E_\mathrm{GNN}^\mathrm{os}$ is shown in black on the right y-axis. Here, the scaling constant $\gamma = 2.5$ and the temperature $T = 300$ K. Note that non-standard amino acids and capping groups are excluded from calculation of $E_\mathrm{GNN}^\mathrm{os}$. For all proteins, the lowest-energy structure is that with the lowest RMSD from the folded state.
• Figure 4: ML/MD simulations with Schake maintain structures close to their starting folded conformations. Starting folded structures are shown in yellow; ML/MD structures correspond to the lowest-$E_\mathrm{GNN}^{\mathrm{os}}$ conformations matching the median trajectory RMSD and are shown in teal. Each protein is labeled with its name, length, and RMSD from the starting structure. Flexible regions (Fig. S6, Tab. S1) are excluded from RMSD calculations.
• Figure 5: The Schake GNN energy, when combined with GBn2, accurately reproduces protein folding landscapes. Umbrella sampling simulations were performed using three different ISMs, and TIP3P results were obtained from lindorff-larsen_how_2011.