Hessian Matching for Machine-Learned Coarse-Grained Molecular Dynamics

# Hessian Matching for Machine-Learned Coarse-Grained Molecular Dynamics
Source: [https://arxiv.org/html/2605.12823](https://arxiv.org/html/2605.12823)
\\undefine@key

newfloatplacement\\undefine@keynewfloatname\\undefine@keynewfloatfileext\\undefine@keynewfloatwithin

Sanya Murdeshwar1Sanjit Shashi1,2Kevin Bachelor1William Noid3Ashwin Lokapally2Razvan Marinescu1,2 1University of California, Santa Cruz2GiwoTech Inc\.3Pennsylvania State University \{smurdesh, sashashi, kwbachelor, ramarine\}@ucsc\.edu wgn1@psu\.eduashwin@giwotech\.com

###### Abstract

Coarse\-grained \(CG\) molecular dynamics enables simulations of atomic systems such as biomolecules at timescales inaccessible to all\-atom \(AA\) methods, but existing CG neural potentials trained via force matching capture only the gradient of the free\-energy surface, leaving its curvature unconstrained\. We introduce a framework that augments force matching with stochastic Hessian\-vector product \(HVP\) matching, instilling second\-order curvature information into CG potentials without constructing the full Hessian\. We derive a decomposition of the target CG Hessian into a model\-independent projected AA Hessian, precomputed once before training, and a model\-dependent covariance correction computed online at negligible cost\. We construct an unbiased stochastic estimator of the Hessian\-matching objective by using random probe vectors\. We evaluate our method by comparing against force matching on a benchmark of nine fast\-folding proteins unseen during training\. HVP matching outperforms plain force matching on 8 of 9 proteins on slow\-mode metrics, with reductions of up to 85% in the Kullback–Leibler divergence between the CG and reference distributions along the slowest collective mode of the largest protein\. Our results demonstrate that higher\-order physical supervision is a practical path to more accurate and transferable CG potentials for biomolecular simulation\.

## 1Introduction

Molecular dynamics \(MD\) simulations are one of the primary tools for understanding the dynamical and functional properties of biomolecules\. However, the timescales accessible to all\-atom \(AA\) MD remain fundamentally mismatched with the biological processes researchers aim to study\. Certain meaningful processes, such as the folding of proteins, can evolve over milliseconds or even seconds, far beyond the microsecond\-order dynamics accessible to specialized AA MD hardwareShaw2009;Shaw2021\. This gap motivates coarse\-grained \(CG\) modelingNoid2013;Kmiecik2016\. CG models represent groups of atoms as single interaction sites \(beads\), reducing the number of degrees of freedom and enabling simulations of longer timescales\. The central challenge is learning a CG potential whose free\-energy surface reproduces that of the underlying AA system, including the locations of metastable states, the heights of barriers between them, and the local curvature that governs fluctuations within each basin\.

Traditionally, CG modeling has been done with force matchingErcolessi1994;Izvekov2005;Noid2008, which fits the CG potential to reproduce the mean force of the AA reference\. The original approach used predefined functional forms, but CG modeling was later generalized to neural network approximators such as CGnetWang2019\(a feedforward network\) and CGSchNetHusic2020\(a graph neural network\)\. These models have become the state of the art for learning CG potentials of biomolecules and can be trained to produce stable, reasonably accurate simulations of fast\-folding proteins\.

Despite this progress, a fundamental limitation of force matching remains; it only captures the gradient of the free\-energy surface, not its curvature\. Two potentials can produce identical forces while differing substantially in how those forces respond to perturbations, and force matching cannot distinguish between them\. In practice, this leads to poor recovery of metastable basin populationsThaler\_Stupp\_Zavadlav\_2022;Köhler\_Chen\_Krämer\_Clementi\_Noé\_2023, with the FM validation loss failing to reflect global free\-energy surface qualityThaler\_Stupp\_Zavadlav\_2022\. Consequently, CG neural potentials trained with force matching alone tend to degrade on slow conformational modes with extended training due to overfitting the gradient signal and losing the shape of the energy landscape\. Furthermore, models trained on one set of proteins extrapolate poorly to unseen, out\-of\-distribution sequences, producing unrealistically low energies in configurations not sampled during trainingMajewski2023, a direct result of a free\-energy surface whose local curvature is underconstrained\. These issues fundamentally limit the data efficiency, generalization ability, and stability of CG neural potentials\.

A natural approach to address this is to match the Hessian of the free\-energy surface alongside its gradient \(forces\)\. The Hessian encodes information about local curvature which governs vibrational frequencies and barrier shapes, providing higher\-order physical information missing from first\-order force matching\. However, directly incorporating Hessian supervision into training CG neural potentials imposes major computational challenges with regards to scaling\. For a CG system withd=3​Nd=3Ndegrees of freedom, the full Hessian hasd2d^\{2\}entries and requires𝒪​\(d\)\\mathcal\{O\}\(d\)force evaluations to construct column\-by\-column, in contrast to the gradient, which only hasddentries and requires𝒪​\(1\)\\mathcal\{O\}\(1\)force evaluations\. For small systems, this is feasible but slow\. For larger systems,ddscales into the thousands, and explicit construction becomes intractable both in memory and compute\.

### 1\.1Related work

Several recent works have incorporated second\-order curvature information into atomistic machine\-learned interatomic potentials \(MLIPs\), showing improvements in extrapolation, data efficiency, and transition state optimizationFang2024BeyondNumerical;Yuan2024AnalyticalHessian;Rodriguez2025HessianData\. The most closely related proposal to ours in this vein is Projected Hessian Learning \(PHL\)Rodriguez2026PHL, which replaces explicit Hessian construction with stochastic Hessian\-vector product probes\. PHL produces an unbiased trace\-based loss that recovers most of the benefits of full Hessian supervision at under 8% of the cost\. Hessian Interatomic Potentials \(HIP\)HIP2025takes a different route, directly predicting Hessians with an equivariant readout rather than via autograd\.

However, all of these methods apply to atomistic MLIPs trained on exact quantum\-mechanical Hessians of very small molecules\. CG systems can scale to thousands of degrees of freedom, making stochastic HVP probes essential rather than optional\. Furthermore, the CG Hessian decomposes into a projected AA Hessian plus a covariance correction arising from integrated\-out atomic degrees of freedom, and the latter has no analogue in the atomistic regime\. We also note that the central goal in CG modeling is generalization to unseen protein conformations and sequences, and none of the existing Hessian\-informed methods have been evaluated to that end\.

A parallel line of work pursues transferable CG protein potentials via diverse multi\-protein training\. Majewski et al\.Majewski2023trained CG potentials on twelve proteins jointly, and Charron et al\.Charron2025Navigatingdeveloped a transferable CGSchNet\-based force field demonstrating extrapolative MD on new sequences\. Both rely solely on force matching, achieving transferability through data scale rather than richer physical supervision\.

Theoretical foundations for force matching with general \(linear and nonlinear\) CG maps were developed by Ciccotti et al\.Ciccotti2005via the vectorial Blue Moon ensemble and reformulated as a geometric projection problem by Kalligiannaki et al\.Kalligiannaki2015, who also showed that force matching and relative\-entropy minimization are asymptotically equivalent for general CG maps\. Our work extends this framework from first to second derivatives of the free energy\.

Our work is also distinct from optimizer\-side curvature methods such as L\-BFGS, K\-FAC, and Hessian\-free optimizationMartens2010;MartensGrosse2015;Pearlmutter1994, which approximate the Hessian of the training objective in parameter space\. We instead match the Hessian of the CG potential in configuration space, which encodes physical curvature of the energy landscape rather than optimization dynamics\.

#### Our contribution\.

We introduce a framework for training CG neural potentials with stochastic Hessian matching via Hessian\-vector products \(HVPs\)\. We demonstrate that the CG Hessian identity decomposes cleanly into two terms with fundamentally different computational properties: a projected AA Hessian term that is model\-independent and can be precomputed once before training, and a covariance correction term that is model\-dependent but built from quantities already available during the forward pass\. We match the action of this Hessian, ad×dd\\times dmatrix, onKKrandom probe vectors, avoiding constructing the entire Hessian explicitly, all while retaining unbiased estimation of the full Hessian\-matching objective\. The resulting loss contributes𝒪​\(K​d\)\\mathcal\{O\}\(Kd\)additional work per frame forKKprobes, which is linear in system size, and requires no architectural changes to the underlying model\.

Our experimental results on fast\-folding benchmark proteins show that supplementing force matching with HVP matching provides the following advantages:

- \(i\)improvements to slow\-mode accuracy \(represented by time\-lagged independent components capturing folding/unfolding transitions\) by up to 85%; and
- \(ii\)better generalization to unseen proteins when trained on a separate single\-chain dataset\.

Ultimately, these findings suggest that instilling higher\-order physical information into CG neural potentials is a practical, scalable path to more accurate and transferable biomolecular simulation\.

## 2Deriving the coarse\-grained Hessian identity

We derive the symbolic expression for the CG Hessian of the free\-energy surface, which serves as the training objective for our model\. For proteins, the standard coarse\-graining mechanism places one bead at the Cαcarbon of each amino acid residue\. Note that because each bead is selected from the AA atoms, the mapping is linear, simplifying our derivation significantly\. The general nonlinear expression is found in Appendix[A](https://arxiv.org/html/2605.12823#A1)\.

We make use of the Blue Moon ensembleCarter1989;Kidder2021to compute thermodynamic quantities associated to constrained molecular systems\. Suppose that the AA system containsnnatoms and that the CG system consists ofN<nN<nbeads, both in three dimensions\. The partition function is

Z​\(𝐑\)=∫𝑑𝐫​δ​\(𝚵r​𝐫−𝐑\)​e−β​ℋ​\(𝐫\),β≡1kB​T\.Z\(\\mathbf\{R\}\)=\\int d\\mathbf\{r\}\\,\\delta\(\\bm\{\\Xi\}\_\{r\}\\mathbf\{r\}\-\\mathbf\{R\}\)e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\},\\ \\ \\ \\ \\beta\\equiv\\frac\{1\}\{k\_\{\\text\{B\}\}T\}\.\(1\)Here𝐫∈ℝ3​n\\mathbf\{r\}\\in\\mathbb\{R\}^\{3n\}is an element of the full configuration space,𝚵r\\bm\{\\Xi\}\_\{r\}is the3​N×3​n3N\\times 3ncoarse\-graining matrix mapping AA positions to CG positions, and𝐑∈ℝ3​N\\mathbf\{R\}\\in\\mathbb\{R\}^\{3N\}is an element of the CG configuration space\. The Hamiltonianℋ\\mathcal\{H\}is a scalar function\.β\\betais the thermodynamic parameter computed from the temperatureTTand Boltzmann constantkBk\_\{\\text\{B\}\}\. Given the partition function, the free energy is defined as

ℱ​\(𝐑\)≡−β−1​log⁡Z​\(𝐑\)\.\\mathcal\{F\}\(\\mathbf\{R\}\)\\equiv\-\\beta^\{\-1\}\\log Z\(\\mathbf\{R\}\)\.\(2\)The CG force is the negative of the first derivative∇Rℱ≡−𝐅CG\\nabla\_\{R\}\\mathcal\{F\}\\equiv\-\\mathbf\{F\}\_\{\\text\{CG\}\}, while the CG Hessian is the second derivative\(∇R⊗∇R\)​ℱ≡𝐇CG\(\\nabla\_\{R\}\\otimes\\nabla\_\{R\}\)\\mathcal\{F\}\\equiv\\mathbf\{H\}\_\{\\text\{CG\}\}\. To compute the latter, we need the former, which is

∇Rℱ=−1β​Z​∇RZ=−1β​Z​∫𝑑𝐫​\[∇Rδ​\(𝚵r​𝐫−𝐑\)\]​e−β​ℋ​\(𝐫\),\\nabla\_\{R\}\\mathcal\{F\}=\-\\frac\{1\}\{\\beta Z\}\\nabla\_\{R\}Z=\-\\frac\{1\}\{\\beta Z\}\\int d\\mathbf\{r\}\\big\[\\nabla\_\{R\}\\delta\(\\bm\{\\Xi\}\_\{r\}\\mathbf\{r\}\-\\mathbf\{R\}\)\\big\]e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\},\(3\)However, we want to recast the derivative inside of the integral to be with respect to the AA coordinates𝐫\\mathbf\{r\}\. To do so, we use the identity

∇Rδ​\(𝚵r​𝐫−𝐑\)=−𝚵F​\[∇rδ​\(𝚵r​𝐫−𝐑\)\],\\nabla\_\{R\}\\delta\(\\bm\{\\Xi\}\_\{r\}\\mathbf\{r\}\-\\mathbf\{R\}\)=\-\\bm\{\\Xi\}\_\{F\}\\left\[\\nabla\_\{r\}\\delta\(\\bm\{\\Xi\}\_\{r\}\\mathbf\{r\}\-\\mathbf\{R\}\)\\right\],\(4\)where𝚵F\\bm\{\\Xi\}\_\{F\}is the following3​N×3​n3N\\times 3nforce\-projection matrix:

𝚵F≡\(𝚵r​𝚵rT\)−1​𝚵r\.\\bm\{\\Xi\}\_\{F\}\\equiv\\left\(\\bm\{\\Xi\}\_\{r\}\\bm\{\\Xi\}\_\{r\}^\{T\}\\right\)^\{\-1\}\\bm\{\\Xi\}\_\{r\}\.\(5\)For the Cαcoarse\-graining map,𝚵r​𝚵rT=𝐈3​N⟹𝚵F=𝚵r\\bm\{\\Xi\}\_\{r\}\\bm\{\\Xi\}\_\{r\}^\{T\}=\\mathbf\{I\}\_\{3N\}\\implies\\bm\{\\Xi\}\_\{F\}=\\bm\{\\Xi\}\_\{r\}\(𝐈3​N\\mathbf\{I\}\_\{3N\}is the3​N×3​N3N\\times 3Nidentity matrix\), but we will continue to assume more general linear mappings here\. Plugging \([4](https://arxiv.org/html/2605.12823#S2.E4)\) into \([3](https://arxiv.org/html/2605.12823#S2.E3)\) and integrating by parts yields

∇Rℱ=1Z​∫𝑑𝐫​δ​\(𝚵r​𝐫−𝐑\)​e−β​ℋ​\(𝒓\)​\(𝚵F​∇rℋ\)\.\\nabla\_\{R\}\\mathcal\{F\}=\\frac\{1\}\{Z\}\\int d\\mathbf\{r\}\\,\\delta\(\\bm\{\\Xi\}\_\{r\}\\mathbf\{r\}\-\\mathbf\{R\}\)e^\{\-\\beta\\mathcal\{H\}\(\\bm\{r\}\)\}\\big\(\\bm\{\\Xi\}\_\{F\}\\nabla\_\{r\}\\mathcal\{H\}\\big\)\.\(6\)In Hamiltonian mechanics,−∇rℋ\-\\nabla\_\{r\}\\mathcal\{H\}is identified as the AA force𝐅AA\\mathbf\{F\}\_\{\\text\{AA\}\}, and so the expression for the first derivative∇Rℱ\\nabla\_\{R\}\\mathcal\{F\}can be reinterpreted as an ensemble average conditioned on𝐑\\mathbf\{R\}\(written as⟨⋅⟩R\\expectationvalue\{\\cdot\}\_\{R\}\):

∇Rℱ=−⟨𝚵F​𝐅AA⟩R⟹𝐅CG=⟨𝚵F​𝐅AA⟩R\.\\nabla\_\{R\}\\mathcal\{F\}=\-\\expectationvalue\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\}\_\{R\}\\implies\\mathbf\{F\}\_\{\\text\{CG\}\}=\\expectationvalue\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\}\_\{R\}\.\(7\)As defined by the force matching framework, the CG force is the ensemble average of the projected AA forces\. The linear identity here, and its nonlinear generalization in Appendix[A](https://arxiv.org/html/2605.12823#A1), recover known mean\-force results from the Blue Moon literatureCiccotti2005;Kalligiannaki2015\.

The CG Hessian is computed by differentiating the ensemble average \([7](https://arxiv.org/html/2605.12823#S2.E7)\)\. Through further path\-integral manipulations \(detailed in Appendix[A](https://arxiv.org/html/2605.12823#A1)\), we obtain

𝐇CG=−⟨𝚵F​\(∇r𝐅AA\)​𝚵FT⟩R−β​𝚺​\(𝚵F​𝐅AA,𝚵F​𝐅AA\),\\mathbf\{H\}\_\{\\text\{CG\}\}=\-\\langle\{\\bm\{\\Xi\}\_\{F\}\\left\(\\nabla\_\{r\}\\mathbf\{F\}\_\{\\text\{AA\}\}\\right\)\\bm\{\\Xi\}\_\{F\}^\{T\}\}\\rangle\_\{R\}\-\\beta\\bm\{\\Sigma\}\(\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\},\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\),\(8\)where𝚺​\(⋅,⋅\)\\bm\{\\Sigma\}\(\\cdot,\\cdot\)is the covariance matrix\. We further note that the AA Hessian𝐇AA\\mathbf\{H\}\_\{\\text\{AA\}\}is defined as the second derivative of the Hamiltonian,i\.e\.𝐇AA≡\(∇r⊗∇r\)​ℋ=−∇r𝐅AA\\mathbf\{H\}\_\{\\text\{AA\}\}\\equiv\(\\nabla\_\{r\}\\otimes\\nabla\_\{r\}\)\\mathcal\{H\}=\-\\nabla\_\{r\}\\mathbf\{F\}\_\{\\text\{AA\}\}\. We then obtain our primary expression for the CG Hessian:

𝐇CG=⟨𝚵F𝐇AA𝚵FT⟩R−β𝚺\(𝚵F𝐅AA,𝚵F𝐅AA\)\.\\boxed\{\\mathbf\{H\}\_\{\\text\{CG\}\}=\\langle\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{H\}\_\{\\text\{AA\}\}\\bm\{\\Xi\}\_\{F\}^\{T\}\}\\rangle\_\{R\}\-\\beta\\bm\{\\Sigma\}\(\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\},\\mathbf\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\)\.\}\(9\)The first term is the ensemble\-averaged projected Hessian, which captures the average curvature of the AA potential energy surface as seen through the CG mapping\. The second term is the force covariance correction, which captures the softening of the effective CG potential due to thermal fluctuations of the integrated\-out AA degrees of freedom\. At higher temperatures, the force covariance𝚺\\bm\{\\Sigma\}grows faster thanβ\\betashrinks, so Term 2 contributes more to the effective CG curvature\.

These two terms carry distinct physical information that goes beyond what the force alone provides, while also containing fundamentally different computational properties\. Term 1 depends only on the AA potential and the CG mapping, making it model\-independent, while Term 2 depends on the CG model’s current force predictions through the force residual\. We exploit this decomposition in the methodology that follows\.

## 3Methodology

We now describe how to turn the CG Hessian identity \([9](https://arxiv.org/html/2605.12823#S2.E9)\) into a practical training objective via computationally\-feasible Hessian\-vector product \(HVP\) calculations\. We formulate the stochastic HVP matching loss \(Section[3\.1](https://arxiv.org/html/2605.12823#S3.SS1)\) and present the training pipeline \(Section[3\.2](https://arxiv.org/html/2605.12823#S3.SS2)\)\.

### 3\.1Stochastic Hessian matching via Hessian\-vector products

#### Intractability of full Hessian matching\.

For a CG system withd=3​Nd=3Ndegrees of freedom, the full Hessian𝐇CG∈ℝd×d\\mathbf\{H\}\_\{\\text\{CG\}\}\\in\\mathbb\{R\}^\{d\\times d\}requires𝒪​\(d2\)\\mathcal\{O\}\(d^\{2\}\)storage and𝒪​\(d\)\\mathcal\{O\}\(d\)force evaluations to fully construct\. For the Chignolin protein \(d=30d=30\) this is feasible though slow, but for larger proteinsddcan scale to the thousands, making explicit Hessian construction intractable\.

#### The HVP reformulation\.

Instead of matching the entire Hessian, we match its action onKKrandom probe vectors, each providing curvature information along one particular direction of the energy landscape\. For each training frame, we sample unit vectors\{𝐯k\}k=1K\\\{\\mathbf\{v\}\_\{k\}\\\}\_\{k=1\}^\{K\}by drawing𝐳k∼𝒩​\(0,𝐈d\)\\mathbf\{z\}\_\{k\}\\sim\\mathcal\{N\}\(0,\\mathbf\{I\}\_\{d\}\)and normalizing to unit length\. Multiplying both sides of \([9](https://arxiv.org/html/2605.12823#S2.E9)\) by𝐯k\\mathbf\{v\}\_\{k\}gives the target HVP:

𝐇CG​𝐯k=𝚵F​\(𝐇AA​𝚵FT​𝐯k\)⏟Term 1−β​δ​𝐉​\(δ​𝐉T​𝐯k\)⏟Term 2\\mathbf\{H\}\_\{\\text\{CG\}\}\\mathbf\{v\}\_\{k\}=\\underbrace\{\\bm\{\\Xi\}\_\{F\}\\left\(\\mathbf\{H\}\_\{\\text\{AA\}\}\\bm\{\\Xi\}\_\{F\}^\{T\}\\mathbf\{v\}\_\{k\}\\right\)\}\_\{\\text\{Term~1\}\}\-\\underbrace\{\\beta\\,\\delta\\mathbf\{J\}\\left\(\\delta\\mathbf\{J\}^\{T\}\\mathbf\{v\}\_\{k\}\\right\)\}\_\{\\text\{Term~2\}\}\(10\)These two terms play distinct physical roles in the Hessian\-matching objective\. For Term 1, which is the projected AA HVP, we embed the CG probe𝐯k\\mathbf\{v\}\_\{k\}into AA space via𝚵FT\\bm\{\\Xi\}\_\{F\}^\{T\}, acting only on the Cαatoms\. The atomistic Hessian𝐇AA\\mathbf\{H\}\_\{\\text\{AA\}\}then returns the AA curvature along this restricted direction, and𝚵F\\bm\{\\Xi\}\_\{F\}projects this curvature back to CG space\. Term 1 depends only on the AA potential and the CG mapping, so it is model\-independent and can therefore be precomputed once before training\.

Term 2 is the covariance correction, driven by the force residual,

δ​𝐉≡𝚵F​𝐅AA−𝐅NN,\\delta\\mathbf\{J\}\\;\\equiv\\;\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\\;\-\\;\\mathbf\{F\}\_\{\\text\{NN\}\},\(11\)which measures how far the CG model’s current mean\-force prediction𝐅NN\\mathbf\{F\}\_\{\\text\{NN\}\}is from the projected AA force𝚵F​𝐅AA\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\. This term is model\-dependent because the quantities change with every training iteration, but since the forces are already computed during the forward pass, it can be computed on\-the\-fly at negligible additional cost\.

Each HVP is a vector of sizedd, and we compute onlyKKof them per frame\. This gives a total cost of𝒪​\(K​d\)\\mathcal\{O\}\(Kd\), which is linear inddfor fixedKK\. Thus, this reformulation lets us instill second\-order curvature information into the CG potential without ever constructing the fulld×dd\\times dHessian matrix\.

#### Loss function and unbiased estimation\.

Given a set ofTTtraining frames, the standard force\-matching loss is the mean\-squared error between the model’s predicted forces and the AA forces projected to CG spaceIzvekov2005;Noid2008:

ℒFM=1T​∑t=1T13​N\(t\)​‖𝐅NN\(t\)−𝚵F​𝐅AA\(t\)‖2,\\mathcal\{L\}\_\{\\text\{FM\}\}=\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}\\frac\{1\}\{3N^\{\(t\)\}\}\\big\\\|\\mathbf\{F\}\_\{\\text\{NN\}\}^\{\(t\)\}\-\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}^\{\(t\)\}\\big\\\|^\{2\},\(12\)where\(t\)designates quantities at framettand𝚵F\\bm\{\\Xi\}\_\{F\}is the force\-projection matrix\. We supplement this with the following HVP\-matching loss:

ℒHVP=1T​∑t=1T1K​∑k=1K13​N\(t\)​‖\(𝐇NN​𝐯k\)\(t\)−\(𝐇CG​𝐯k\)\(t\)‖2,\\mathcal\{L\}\_\{\\text\{HVP\}\}=\\frac\{1\}\{T\}\\sum\_\{t=1\}^\{T\}\\frac\{1\}\{K\}\\sum\_\{k=1\}^\{K\}\\frac\{1\}\{3N^\{\(t\)\}\}\\big\\\|\(\\mathbf\{H\}\_\{\\text\{NN\}\}\\mathbf\{v\}\_\{k\}\)^\{\(t\)\}\-\(\\mathbf\{H\}\_\{\\text\{CG\}\}\\mathbf\{v\}\_\{k\}\)^\{\(t\)\}\\big\\\|^\{2\},\(13\)where\(𝐇NN​𝐯k\)\(t\)\(\\mathbf\{H\}\_\{\\text\{NN\}\}\\mathbf\{v\}\_\{k\}\)^\{\(t\)\}is obtained via a second autograd call through the GNN\. A set ofKKprobe vectors\{𝐯k\(t\)\}\\\{\\mathbf\{v\}\_\{k\}^\{\(t\)\}\\\}is chosen independently for each framett, though seeded deterministically by the frame index for reproducibility\. For probe vectors uniformly distributed on the unit sphere,𝔼𝐯​\[‖𝐇𝐯‖2\]=‖𝐇‖F2/d\\mathbb\{E\}\_\{\\mathbf\{v\}\}\[\\\|\\mathbf\{H\}\\mathbf\{v\}\\\|^\{2\}\]=\\\|\\mathbf\{H\}\\\|\_\{F\}^\{2\}/dfor any matrix𝐇\\mathbf\{H\}, where\|\|⋅\|\|F\|\|\\cdot\|\|\_\{F\}is the Frobenius norm\. So, each per\-frame term is an unbiased stochastic estimate of the full Hessian\-matching objective, where finiteKKcontributes variance rather than bias\. The total training objective combines force and HVP matching:

ℒ=wFM​ℒFM\+wHVP​ℒHVP,\\mathcal\{L\}=w\_\{\\text\{FM\}\}\\,\\mathcal\{L\}\_\{\\text\{FM\}\}\+w\_\{\\text\{HVP\}\}\\,\\mathcal\{L\}\_\{\\text\{HVP\}\},\(14\)

### 3\.2Training pipeline

The overall pipeline is shown schematically in Figure[1](https://arxiv.org/html/2605.12823#S3.F1)\. We summarize each step here\.

#### Probe selection\.

For each framet=1,…,Tt=1,\\dots,T, we generateKKprobes using a deterministic seeding schemes\(t\)=ts^\{\(t\)\}=t\. Probe consistency for HVP matching is guaranteed by this scheme, in which the same seed produces identical probes per frame on any device and across multi\-GPU setups\. Term 1 targets can therefore be precomputed once and reused across any subsequent model trained with the HVP loss\.

#### Precomputing Term 1 target\.

Since the Term 1 targets in \([10](https://arxiv.org/html/2605.12823#S3.E10)\) are model\-independent, we compute them once before training and store the results\. We use the Amber14 force fieldMaier2015evaluated in OpenMMEastman2017with no nonbonded cutoff\. For each frame, the CG probe is embedded into AA space via𝐯~k≡𝚵FT​𝐯k\\tilde\{\\mathbf\{v\}\}\_\{k\}\\equiv\\bm\{\\Xi\}\_\{F\}^\{T\}\\mathbf\{v\}\_\{k\}, and the AA HVP is approximated via central finite differences:

𝐇AA​𝐯~k≈−𝐅AA​\(𝐫\+ε​𝐯~k\)−𝐅AA​\(𝐫−ε​𝐯~k\)2​ε,\\mathbf\{H\}\_\{\\text\{AA\}\}\\tilde\{\\mathbf\{v\}\}\_\{k\}\\approx\-\\frac\{\\mathbf\{F\}\_\{\\text\{AA\}\}\(\\mathbf\{r\}\+\\varepsilon\\tilde\{\\mathbf\{v\}\}\_\{k\}\)\-\\mathbf\{F\}\_\{\\text\{AA\}\}\(\\mathbf\{r\}\-\\varepsilon\\tilde\{\\mathbf\{v\}\}\_\{k\}\)\}\{2\\varepsilon\},\(15\)whereε\\varepsilonis some small finite\-difference step\. The result is projected back to CG space via𝚵F\\bm\{\\Xi\}\_\{F\}and unit\-converted\.

#### Computing the CG model’s HVP\.

The CG\-side HVP𝐇NN​𝐯k\\mathbf\{H\}\_\{\\text\{NN\}\}\\,\\mathbf\{v\}\_\{k\}is obtained via two sequential applications of automatic differentiation, which is the primary source of computational overhead in our pipeline\. First, the energyWNN​\(𝐑\)W\_\{\\text\{NN\}\}\(\\mathbf\{R\}\)is differentiated with respect to positions to produce the forces𝐅NN=−∇RWNN\\mathbf\{F\}\_\{\\text\{NN\}\}=\-\\nabla\_\{R\}W\_\{\\text\{NN\}\}, with the computation graph retained\. Note that the forces are not an explicit output of the GNN, so this step is what creates them as a differentiable node in the graph in the first place\. Subsequently, the inner product𝐅NN⋅𝐯k\\mathbf\{F\}\_\{\\text\{NN\}\}\\cdot\\mathbf\{v\}\_\{k\}is differentiated with respect to positions, yielding the HVP:

𝐇NN​𝐯k=−∇R\(𝐅NN⋅𝐯k\)\.\\mathbf\{H\}\_\{\\text\{NN\}\}\\,\\mathbf\{v\}\_\{k\}=\-\\nabla\_\{R\}\(\\mathbf\{F\}\_\{\\text\{NN\}\}\\cdot\\mathbf\{v\}\_\{k\}\)\.\(16\)The graph from the first differentiation step is reused across allKKprobes, so the marginal cost of each additional probe is a single backward pass through the network\.

#### Computing the covariance correction \(Term 2\)\.

The covariance correctionβ​δ​𝐉​\(δ​𝐉T​𝐯k\)\\beta\\,\\delta\\mathbf\{J\}\(\\delta\\mathbf\{J\}^\{T\}\\mathbf\{v\}\_\{k\}\)is computed from the force residualδ​𝐉\\delta\\mathbf\{J\}, so it can only be computed online during training\. However, because this quantity is built from forces already produced during the forward pass \(𝐅NN\\mathbf\{F\}\_\{\\text\{NN\}\}and𝚵F​𝐅AA\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\), it adds negligible computational overhead\. The residual is detached from the computation graph, so gradients do not flow through the target\. Term 2 is then subtracted from the precomputed Term 1 target to form the full HVP target for the loss\.

#### Training loop\.

At each iteration, the GNN forward pass produces energy, forces, and HVPs via the two\-pass autograd pipeline\. Probe vectors are regenerated at training time, avoiding the memory overhead of savingKKvectors per frame\. Our pipeline supports two variants ofℒHVP\\mathcal\{L\}\_\{\\text\{HVP\}\}—with the covariance correction enabled or with it disabled\. If the covariance correction is included, then Term 2 is assembled on GPU from the force residual and subtracted from Term 1 before the loss computation\. The combined lossℒ\\mathcal\{L\}is then backpropagated through the network, with no gradients flowing through the HVP target since it is either precomputed or detached\.

Shared probes𝐯1,…,𝐯K\\mathbf\{v\}\_\{1\},\\ldots,\\mathbf\{v\}\_\{K\}Deterministic seeds\(t\)=ts^\{\(t\)\}=tEnergy functionGNNWNN​\(𝐑\)W\_\{\\text\{NN\}\}\(\\mathbf\{R\}\)Autograd HVP𝐅NN=−∇RWNN\\mathbf\{F\}\_\{\\text\{NN\}\}=\-\\nabla\_\{R\}W\_\{\\text\{NN\}\}𝐇NN​𝐯k=−∇R\(𝐅NN​𝐯k\)\\mathbf\{H\}\_\{\\text\{NN\}\}\\mathbf\{v\}\_\{k\}=\-\\nabla\_\{R\}\\big\(\\mathbf\{F\}\_\{\\text\{NN\}\}\\mathbf\{v\}\_\{k\}\\big\)Finite differences𝐇AA​𝐯k=−𝐅​\(𝐫\+ε​𝐯~\)−𝐅​\(𝐫−ε​𝐯~\)2​ε\\mathbf\{H\}\_\{\\text\{AA\}\}\\mathbf\{v\}\_\{k\}=\-\\dfrac\{\\mathbf\{F\}\(\\mathbf\{r\}\+\\varepsilon\\tilde\{\\mathbf\{v\}\}\)\-\\mathbf\{F\}\(\\mathbf\{r\}\-\\varepsilon\\tilde\{\\mathbf\{v\}\}\)\}\{2\\varepsilon\}Term 1 target𝚵F​\(𝑯AA​𝚵FT​𝐯k\)\\bm\{\\Xi\}\_\{F\}\\big\(\\bm\{H\}\_\{\\text\{AA\}\}\\bm\{\\Xi\}\_\{F\}^\{T\}\\mathbf\{v\}\_\{k\}\\big\)Term 2\(online\)β​δ​𝐉​\(δ​𝐉T​𝐯k\)\\beta\\,\\delta\{\\mathbf\{J\}\}\\big\(\\delta\{\\mathbf\{J\}\}^\{T\}\\mathbf\{v\}\_\{k\}\\big\)Combined lossℒ=wFM​ℒFM\+wHVP​ℒHVP\\mathcal\{L\}\\;=\\;w\_\{\\text\{FM\}\}\\,\\mathcal\{L\}\_\{\\text\{FM\}\}\\;\+\\;w\_\{\\text\{HVP\}\}\\,\\mathcal\{L\}\_\{\\text\{HVP\}\}CG spaceAA space\(precomputed\)predictiontarget

Figure 1:Overview of the HVP matching pipeline\. Shared probe vectors𝐯k\\mathbf\{v\}\_\{k\}are generated from deterministic per\-frame seeds, and are used on both sides of the matching\.Left:the GNN produces energy and forces, and two\-pass automatic differentiation yieldsHNN​𝐯kH\_\{\\text\{NN\}\}\\mathbf\{v\}\_\{k\}\(the prediction\)\.Right:central finite differences on the AA force field produce Term 1 targets, precomputed in one shot before training\.Center:the force residualδ​𝐉\\delta\{\\mathbf\{J\}\}is computed from the CG model’s force prediction and the projected AA forces in order to evaluate Term 2 online\. The combined loss matches both forces \(ℒFM\\mathcal\{L\}\_\{\\text\{FM\}\}\) and curvature \(ℒHVP\\mathcal\{L\}\_\{\\text\{HVP\}\}\) factors\.

## 4Experiments

We evaluate exclusively in the out\-of\-distribution setting, as the validation loss curves \(Appendix[B](https://arxiv.org/html/2605.12823#A2)\) confirm that curvature matching does not compromise force prediction accuracy within the training distribution\. The models are trained on a separate single\-chain dataset \(Appendix[C](https://arxiv.org/html/2605.12823#A3)\) and evaluated on unseen benchmark proteins \(Section[4\.1](https://arxiv.org/html/2605.12823#S4.SS1)\)\.

#### Datasets\.

We use the benchmark suite ofAghili2025, comprising 9 fast\-folding proteins spanning a range of sizes from Chignolin \(10 CG beads\) to Lambda repressor \(80 CG beads\)\. For the generalization experiment, we train on a curated dataset of 99 single\-chain proteins \(detailed in Appendix[C](https://arxiv.org/html/2605.12823#A3)\) with 10,000 frames per protein \(∼990,000\{\\sim\}990\{,\}000total training frames\) and no overlap with the benchmark set\.

#### Model variants\.

We compare three training configurations with identical architecture and hyperparameters:

- •FM: force matching only;
- •FM\+AAp\{\}\_\{\\text\{p\}\}: force and HVP matching with Term 1 \(the AA Hessian projection\) only; and
- •FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov: the full objective including the Term 2 covariance correction\.

The loss weights in \([14](https://arxiv.org/html/2605.12823#S3.E14)\) are set towFM=1w\_\{\\text\{FM\}\}=1and eitherwHVP=0w\_\{\\text\{HVP\}\}=0for FM orwHVP=0\.01w\_\{\\text\{HVP\}\}=0\.01for FM\+AAp\{\}\_\{\\text\{p\}\}and FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov\.

To represent the energy functionWNNW\_\{\\text\{NN\}\}, all models used the SchNet\-based graph neural network from TorchMD\-NetDoerr2021;pelaez2024torchmdnet, following the CGSchNet frameworkHusic2020for coarse\-grained protein potentials\. The architecture comprises 4 message\-passing layers, 64\-dimensional embeddings, 12 radial basis functions \(exponential normal\), and SiLU activations, with a cutoff of 3\.0–12\.0 Å\.

All three models were trained on the single\-chain dataset \(Appendix[C](https://arxiv.org/html/2605.12823#A3)\) until convergence, plateauing at 52 \(FM\), 57 \(FM\+AAp\{\}\_\{\\text\{p\}\}\), and 62 \(FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov\) epochs respectively\. Training used AdamWLoshchilov2019at a learning rate of10−410^\{\-4\}with batch size 200 on 4×\\timesA40 GPUs\. All simulations and training data corresponded toT=300T=300K, giving the Term 2 prefactorβ=1/\(kB​T\)=1\.677\\beta=1/\(k\_\{\\text\{B\}\}T\)=1\.677mol/kcal\. The HVPs were computed withK=8K=8probe vectors per frame, withε=10−5\\varepsilon=10^\{\-5\}nm for the finite\-differences computation \([15](https://arxiv.org/html/2605.12823#S3.E15)\)\. By the end of training, our per\-epoch wall\-clock averaged 105 min \(FM\) and 122 min \(FM\+AAp\{\}\_\{\\text\{p\}\}\), a∼\\sim16% overhead, with the covariance correction adding negligible additional per\-step cost\. HVP target precomputation is a one\-time∼\\sim2\.4 hour cost\.

#### Evaluation\.

Following the benchmarking methodology of Aghili et al\.Aghili2025, we run CG molecular dynamics from each trained potential and compare trajectory metric distributions against AA reference simulations\. Time\-lagged independent component analysis \(TICA\)Perez2013identifies and ranks the directions of slowest collective motion in molecular trajectories\. In particular, TIC 0 captures the slowest \(typically folding/unfolding\) transition\. We project our trajectories onto these components and compute distributional metrics on the projections\. The TICA 2D distribution refers to the joint Gaussian\-KDE\-estimated density over the two slowest components, TIC 0 and TIC 1\.

All metrics are computed over 20 independent CG MD replica trajectories per protein, averaging out variance from MD sampling\. Each model variant was trained from a single random seed, and multi\-seed training is left to future work\. We evaluate our models using the following metrics:

- •Wasserstein\-1 \(W1\) and Kullback–Leibler \(KL\) divergence on TICA projections \(TIC 0–3\), capturing slow conformational dynamics; and
- •W1 on bond length, bond angle, and dihedral distributions, and radius of gyration\.

The evaluation metrics are listed in Table[1](https://arxiv.org/html/2605.12823#S4.T1)\.

### 4\.1Generalization to unseen proteins

Table 1:Generalization to unseen proteins\. All models were trained on the single\-chain dataset \(no overlap with benchmark\)\. Best result per protein per metric inbold\.↓\\downarrow= lower is better\.Slow dynamicsStructuralProtein\(No\. of beadsNN\)ModelTICA 2DW1↓\\downarrowTIC 0KL↓\\downarrowTIC 1KL↓\\downarrowTIC 2KL↓\\downarrowTIC 3KL↓\\downarrowDihed\.W1↓\\downarrowAngleW1↓\\downarrowBondW1↓\\downarrowGyr\.W1↓\\downarrowChignolinN=10N=10FM0\.8312\.380\.581\.810\.160\.5850\.0720\.00070\.127FM\+AAp\{\}\_\{\\text\{p\}\}0\.5280\.680\.520\.350\.080\.4020\.0520\.00020\.095FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov0\.7280\.971\.503\.130\.860\.6450\.2300\.00170\.170Trp\-cageN=20N=20FM1\.4204\.454\.660\.880\.680\.1710\.1180\.00090\.175FM\+AAp\{\}\_\{\\text\{p\}\}1\.1854\.973\.991\.411\.110\.1530\.0680\.00110\.129FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov0\.9734\.182\.400\.580\.700\.1960\.0900\.00140\.065BBAN=28N=28FM1\.1593\.731\.441\.460\.960\.2700\.0840\.00270\.127FM\+AAp\{\}\_\{\\text\{p\}\}1\.5564\.442\.241\.301\.730\.4490\.0460\.00180\.290FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov1\.1235\.341\.931\.660\.840\.1480\.1280\.00090\.115WW domainN=34N=34FM1\.3868\.543\.821\.791\.540\.4930\.0690\.00270\.063FM\+AAp\{\}\_\{\\text\{p\}\}1\.3017\.202\.651\.751\.420\.6420\.1160\.00090\.252FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov1\.1057\.142\.701\.391\.650\.4340\.0860\.00200\.173Protein BN=47N=47FM0\.9054\.142\.463\.562\.580\.3340\.1960\.00160\.234FM\+AAp\{\}\_\{\\text\{p\}\}1\.1284\.142\.681\.073\.000\.5210\.1760\.00040\.111FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov0\.6400\.890\.891\.391\.550\.3610\.1180\.00470\.183HomeodomainN=54N=54FM0\.9360\.530\.981\.121\.100\.3890\.1140\.00470\.097FM\+AAp\{\}\_\{\\text\{p\}\}1\.4742\.577\.073\.284\.510\.2800\.0790\.00210\.103FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov0\.8671\.110\.861\.400\.440\.3580\.0790\.00210\.081Protein GN=56N=56FM1\.0713\.271\.590\.871\.670\.3740\.0830\.00930\.220FM\+AAp\{\}\_\{\\text\{p\}\}1\.2245\.532\.311\.012\.240\.3960\.0570\.00310\.194FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov0\.8271\.121\.350\.710\.290\.5000\.0450\.00470\.171α\\alpha3DN=73N=73FM1\.6519\.211\.190\.840\.780\.6050\.2120\.00900\.277FM\+AAp\{\}\_\{\\text\{p\}\}1\.5689\.051\.421\.441\.390\.6790\.2640\.00080\.259FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov2\.83412\.4712\.922\.264\.520\.8360\.3790\.00310\.190Lambda repressorN=80N=80FM1\.34710\.195\.373\.047\.290\.2960\.2030\.00150\.116FM\+AAp\{\}\_\{\\text\{p\}\}1\.0853\.381\.612\.331\.630\.2710\.1520\.00210\.081FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov0\.7941\.490\.780\.680\.580\.2500\.0950\.00110\.129

The clearest result from Table[1](https://arxiv.org/html/2605.12823#S4.T1)is that some form of HVP matching outperforms plain force matching on slow\-mode TICA metrics for 8 of 9 proteins, confirming that second\-order curvature supervision improves generalization to unseen proteins\. The most striking individual result is on Lambda repressor, the largest protein in the benchmark at 80 CG beads\. There, FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov reduces TIC 0 KL from 10\.19 to 1\.49, TIC 1 KL from 5\.37 to 0\.78, and TIC 3 KL from 7\.29 to 0\.58, yielding improvements of 85%, 85%, and 92% relative to force matching alone, achieved entirely out\-of\-distribution\. This result alone establishes that stochastic HVP matching scales to large proteins and provides the strongest generalization gains precisely where force matching struggles most\. Which model variant is optimal, however, depends on system size, and we address this below\.

#### Term 1 alone is sufficient for small systems\.

On Chignolin, FM\+AAp\{\}\_\{\\text\{p\}\}wins every metric, reducing TICA 2D W1 by 36% relative to FM and improving all four TICs, dihedrals, and gyration\. Figure[2](https://arxiv.org/html/2605.12823#S4.F2)illustrates this qualitatively\. Adding the covariance correction degrades performance across the board \(e\.g\. TIC 2 KL increases from 0\.35 to 3\.13\), suggesting that the force residualδ​𝐉\\delta\\mathbf\{J\}for small, fast\-converging systems is dominated by training noise rather than genuine thermal fluctuations\.

![Refer to caption](https://arxiv.org/html/2605.12823v1/figures/tica_contour_comparison_chignolin.png)Figure 2:TICA free\-energy contours for Chignolin\. Blue contours show the ground truth \(AA reference\), and red contours show the CG model\.\(a\) FM:the model’s density concentrates in a single basin \(TIC 0≈−0\.5\\text\{TIC~0\}\\approx\-0\.5\) with no density in the secondary basin \(TIC 0≈1\.5\\text\{TIC~0\}\\approx 1\.5\)\.\(b\) FM\+AAp\{\}\_\{\\text\{p\}\}:the model reproduces both basins more accurately with continuous density bridging the transition region, demonstrating that the curvature signal from HVP matching enables meaningful extrapolation to protein conformations not seen during training\.
#### The full identity is necessary for larger systems\.

For proteins with 20 or more CG beads, FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov consistently produces the best slow\-mode accuracy \(excludingα\\alpha3D, which we discuss below\)\. On Trp\-cage, Cov wins TICA 2D W1 and reduces TIC 1 KL by∼\\sim48%\. On WW domain, Cov achieves the best TICA 2D W1 and also wins TIC 0 and TIC 2\.

The effect is most dramatic for the largest proteins\. In particular, Lambda repressor exhibits a clear improvement progression on every slow\-mode metric, with TIC 0 KL dropping from 10\.19 \(FM\) to 3\.38 \(AAp\{\}\_\{\\text\{p\}\}\) to 1\.49 \(Cov\), yielding an 85% reduction from force matching alone\. Protein G and Protein B show similarly strong Cov results, with TICA 2D W1 reductions of 23% and 29% relative to FM\. Across the benchmark, FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov achieves the best TICA 2D W1 on 7 of 9 proteins\.

In contrast, FM\+AAp\{\}\_\{\\text\{p\}\}can be worse than plain force matching on slow modes for larger proteins\. On Homeodomain, AAp\{\}\_\{\\text\{p\}\}alone increases TIC 1 KL from 0\.98 to 7\.07 and TIC 3 from 1\.10 to 4\.51\. On Protein G, TIC 0 worsens from 3\.27 \(FM\) to 5\.53 \(FM\+AAp\{\}\_\{\\text\{p\}\}\)\. In both cases, adding the covariance correction fully recovers performance: Homeodomain Cov achieves TIC 1 of 0\.86 and TIC 3 of 0\.44, and Protein G Cov achieves TIC 0 of 1\.12\. This observation suggests that for larger systems, Term 1 alone provides an incomplete curvature target, so the model\-dependent correction from Term 2 is necessary to produce a physically consistent Hessian signal\.

#### The𝜶\\bm\{\\alpha\}3D outlier\.

α\\alpha3D \(73 beads\) is the sole benchmark protein where FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov damages slow\-mode performance\. Term 1 alone provides only marginal improvement on TICA 2D W1 and TIC 0, and FM wins TIC 1–3\. Cov is exceptionally damaging\. A plausible explanation is thatα\\alpha3D’s three\-helix bundle topology is underrepresented in the single\-chain training set, so curvature supervision alone cannot compensate for the distributional gap\. Probe coverage may also play a contributing role, sinceK=8K=8probes sample a diminishing fraction of thedd\-dimensional curvature space, though Lambda repressor’s gains \(with largerdd\) suggest that topology coverage matters more\.

#### Structural metrics\.

Local structural properties \(bonds, angles, dihedrals\) show a more mixed pattern\. FM\+AAp\{\}\_\{\\text\{p\}\}tends to improve bond length distributions, likely due to the Hessian directly governing the stiffness of bonded interactions\. Dihedral and angle improvements are protein\-dependent and do not follow a consistent trend across model variants\. These observations are expected, since these local geometric properties are already well\-constrained by force matching\.

## 5Discussion

We have presented a framework for instilling second\-order curvature information into coarse\-grained neural potentials via stochastic HVP matching\. Our key theoretical result is the decomposition of the CG Hessian into a model\-independent projected AA Hessian \(Term 1\) and a model\-dependent covariance correction \(Term 2\)\. Term 1 targets are precomputed once and reused across training runs, while Term 2 is assembled from quantities already available during the forward pass, enabling the addition of curvature supervision to any force\-matching pipeline with no architectural changes and𝒪​\(K​d\)\\mathcal\{O\}\(Kd\)additional cost per frame\.

Our ablation experiments on nine fast\-folding benchmark proteins in the out\-of\-distribution setting confirm that our method comes with a clear empirical payoff\. Some form of HVP matching outperforms plain force matching on slow\-mode TICA metrics for 8 of 9 proteins, establishing that second\-order curvature supervision meaningfully improves generalization to unseen proteins\. Which variant is optimal depends on system size\. Term 1 alone suffices for small systems, while larger systems require the full two\-term objective, yielding reductions in TIC 0 KL of up to 85% relative to force matching alone\. Notably, the covariance correction also recovers the cases where Term 1 alone degrades performance, confirming that the complete Hessian identity is necessary to produce a physically consistent curvature signal at scale\.

Several limitations point toward future work\. Probe coverage may bottleneck the largest systems, sinceK=8K=8probes sample a diminishing fraction of thedd\-dimensional curvature space as protein size grows\. Adaptive probe selection along physically important directions could improve efficiency without increasingKK\. Our evaluation is also limited to Cαlinear coarse\-graining of single\-chain proteins, leaving the extension to diverse complexes, non\-standard CG mappings, and the nonlinear Hessian identity in Appendix[A](https://arxiv.org/html/2605.12823#A1)as natural next steps\. Furthermore, the loss weightwHVP=0\.01w\_\{\\text\{HVP\}\}=0\.01was chosen heuristically, and Term 1 and Term 2 are weighted equally in the target \([10](https://arxiv.org/html/2605.12823#S3.E10)\)\. A systematic study ofwHVPw\_\{\\text\{HVP\}\}, alongside an interpolation coefficientα∈\[0,1\]\\alpha\\in\[0,1\]that continuously scales the covariance correction, could replace the binary FM\+AAp\{\}\_\{\\text\{p\}\}vs\. FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov ablation with a single hyperparameter whose optimal value may correlate with system size\. Finally, theα\\alpha3D outlier suggests that curvature supervision alone cannot fully compensate for gaps in the training distribution, motivating the use of larger, more diverse training sets\.

Ultimately, our results indicate that the accuracy bottleneck in CG neural potentials is addressable not by increasing model capacity or data scale, but by enriching the physical content of the training signal\. The same GNN architecture, when trained with higher\-order physical targets, produces substantially better generalization to unseen proteins without any changes to model design or training data\. We expect that this principle—that higher\-order derivatives of the target potential carry transferable information beyond that provided by first\-order matching—extends naturally beyond the CG MD setting to any domain where learned potentials must extrapolate reliably\.

## References

## Appendix AThe Hessian for nonlinear coarse graining

In Section[2](https://arxiv.org/html/2605.12823#S2)of the main paper, we derived the Hessian for linear coarse\-graining maps\. The first\-derivative results recover the local mean force expression of Kalligiannaki et al\.\[Kalligiannaki2015\], while our novel contribution is the Hessian identity in \([29](https://arxiv.org/html/2605.12823#A1.E29)\)\. Here, we find the more general expression\. We again use the Blue Moon ensemble and assumennatoms andN<nN<nbeads in three dimensions\. This time, the partition function is

Z​\(𝐑\)=∫𝑑𝐫​δ​\(𝝃r​\(𝐫\)−𝐑\)​e−β​ℋ​\(𝐫\),β≡1kB​T\.Z\(\\mathbf\{R\}\)=\\int d\\mathbf\{r\}\\,\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\},\\ \\ \\ \\ \\beta\\equiv\\frac\{1\}\{k\_\{\\text\{B\}\}T\}\.\(17\)Again,𝐫∈ℝ3​n\\mathbf\{r\}\\in\\mathbb\{R\}^\{3n\}and𝐑∈ℝ3​N\\mathbf\{R\}\\in\\mathbb\{R\}^\{3N\}are elements of the AA space and the CG space, respectively\.𝝃r:ℝ3​n→ℝ3​N\\bm\{\\xi\}\_\{r\}:\\mathbb\{R\}^\{3n\}\\to\\mathbb\{R\}^\{3N\}is the coarse\-graining function\.

### A\.1Computing the coarse\-grained force

The computation of the first derivative of the free energyℱ\\mathcal\{F\}is very similar to the linear case\. The starting point is

∇Rℱ=−1β​Z​∇RZ=−1β​Z​∫𝑑𝐫​\[∇Rδ​\(𝝃r​\(𝐫\)−𝐑\)\]​e−β​ℋ​\(𝐫\),\\nabla\_\{R\}\\mathcal\{F\}=\-\\frac\{1\}\{\\beta Z\}\\nabla\_\{R\}Z=\-\\frac\{1\}\{\\beta Z\}\\int d\\mathbf\{r\}\\big\[\\nabla\_\{R\}\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)\\big\]e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\},\(18\)and we again need to rewrite the derivative to be in terms of𝐫\\mathbf\{r\}\. For the general case, the necessary formula is

∇Rδ​\(𝝃r​\(𝐫\)−𝐑\)=−𝚵F​\[∇rδ​\(𝝃r​\(𝐫\)−𝐑\)\],\\nabla\_\{R\}\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)=\-\\bm\{\\Xi\}\_\{F\}\\left\[\\nabla\_\{r\}\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)\\right\],\(19\)where this time we define the3​N×3​n3N\\times 3nmatrix𝚵F\\bm\{\\Xi\}\_\{F\}in terms of the3​N×3​n3N\\times 3nJacobian of𝝃r\\bm\{\\xi\}\_\{r\}:

𝚵F≡\(𝑱ξ​𝑱ξT\)−1​𝑱ξ,𝑱ξ≡∇r𝝃r=\(∂r1𝝃r​⋯​∂r3​n𝝃r\)\.\\bm\{\\Xi\}\_\{F\}\\equiv\\left\(\\bm\{J\}\_\{\\xi\}\\bm\{J\}\_\{\\xi\}^\{T\}\\right\)^\{\-1\}\\bm\{J\}\_\{\\xi\},\\ \\ \\ \\ \\bm\{J\}\_\{\\xi\}\\equiv\\nabla\_\{r\}\\bm\{\\xi\}\_\{r\}=\\begin\{pmatrix\}\\partial\_\{r\_\{1\}\}\\bm\{\\xi\}\_\{r\}\\cdots\\partial\_\{r\_\{3n\}\}\\bm\{\\xi\}\_\{r\}\\end\{pmatrix\}\.\(20\)Then integration by parts yields

∇Rℱ=1Z​∫𝑑𝐫​δ​\(𝝃r​\(𝐫\)−𝐑\)​e−β​ℋ​\(𝒓\)​\(𝚵F​∇rℋ−1β​∇r⋅𝚵F\),\\nabla\_\{R\}\\mathcal\{F\}=\\frac\{1\}\{Z\}\\int d\\mathbf\{r\}\\,\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)e^\{\-\\beta\\mathcal\{H\}\(\\bm\{r\}\)\}\\left\(\\bm\{\\Xi\}\_\{F\}\\nabla\_\{r\}\\mathcal\{H\}\-\\frac\{1\}\{\\beta\}\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\\right\),\(21\)where∇r⋅𝚵F\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}is the divergence taken over the3​n3ncolumn indices and so is a3​N×13N\\times 1vector\.

Unlike in the linear case \([6](https://arxiv.org/html/2605.12823#S2.E6)\), there are now two terms in∇Rℱ\\nabla\_\{R\}\\mathcal\{F\}because the Jacobian and, thus,𝚵F\\bm\{\\Xi\}\_\{F\}are no longer assumed as constant\. This implies that the CG force is no longer obtained by applying a projection to the AA force\. In the language of ensemble averages and the AA force𝐅AA\\mathbf\{F\}\_\{\\text\{AA\}\}, the first derivative and corresponding CG force𝐅CG\\mathbf\{F\}\_\{\\text\{CG\}\}are

∇Rℱ=−⟨𝚵F​𝐅AA\+1β​∇r⋅𝚵F⟩R⟹𝐅CG=⟨𝚵F​𝐅AA\+1β​∇r⋅𝚵F⟩R\.\\nabla\_\{R\}\\mathcal\{F\}=\-\\expectationvalue\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\+\\frac\{1\}\{\\beta\}\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\}\_\{R\}\\implies\\mathbf\{F\}\_\{\\text\{CG\}\}=\\expectationvalue\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\+\\frac\{1\}\{\\beta\}\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\}\_\{R\}\.\(22\)

### A\.2Differentiating the coarse\-grained force

To derive the formula for the CG Hessian, it is helpful to compute a slightly more general derivative—that of the ensemble average of an arbitrary function𝒪​\(𝐫\)\\mathcal\{O\}\(\\mathbf\{r\}\)of the atomic position𝐫\\mathbf\{r\}\. This is because∇Rℱ\\nabla\_\{R\}\\mathcal\{F\}itself is an ensemble average\. By definition, we have that

⟨𝒪⟩R=1Z​∫𝑑𝐫​δ​\(𝝃r​\(𝐫\)−𝐑\)​e−β​ℋ​\(𝐫\)​𝒪​\(𝐫\)\.\\expectationvalue\{\\mathcal\{O\}\}\_\{R\}=\\frac\{1\}\{Z\}\\int d\\mathbf\{r\}\\,\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\}\\mathcal\{O\}\(\\mathbf\{r\}\)\.\(23\)By using the product rule, we write

∇R⟨𝒪⟩R=−∇RZZ2​∫𝑑𝐫​δ​\(𝝃r​\(𝐫\)−𝐑\)​e−β​ℋ​\(𝐫\)​𝒪​\(𝐫\)\+1Z​∫𝑑𝐫​∇Rδ​\(𝝃r​\(𝐫\)−𝐑\)​e−β​ℋ​\(𝐫\)​𝒪​\(𝐫\)\.\\begin\{split\}\\nabla\_\{R\}\\expectationvalue\{\\mathcal\{O\}\}\_\{R\}&=\-\\frac\{\\nabla\_\{R\}Z\}\{Z^\{2\}\}\\int d\\mathbf\{r\}\\,\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\}\\mathcal\{O\}\(\\mathbf\{r\}\)\\\\ &\\qquad\\quad\+\\frac\{1\}\{Z\}\\int d\\mathbf\{r\}\\,\\nabla\_\{R\}\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\}\\mathcal\{O\}\(\\mathbf\{r\}\)\.\\end\{split\}\(24\)The first term can be rewritten in terms of the free energy and⟨𝒪⟩R\\expectationvalue\{\\mathcal\{O\}\}\_\{R\}\. The second term can be further manipulated by invoking \([19](https://arxiv.org/html/2605.12823#A1.E19)\)\. This leads to

∇R⟨𝒪⟩R=β​∇Rℱ​⟨𝒪⟩R−1Z​∫𝑑𝐫​𝚵F​∇rδ​\(𝝃r​\(𝐫\)−𝐑\)​e−β​ℋ​\(𝐫\)​𝒪​\(𝐫\)\.\\nabla\_\{R\}\\expectationvalue\{\\mathcal\{O\}\}\_\{R\}=\\beta\\,\\nabla\_\{R\}\\mathcal\{F\}\\expectationvalue\{\\mathcal\{O\}\}\_\{R\}\-\\frac\{1\}\{Z\}\\int d\\mathbf\{r\}\\,\\bm\{\\Xi\}\_\{F\}\\nabla\_\{r\}\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\}\\mathcal\{O\}\(\\mathbf\{r\}\)\.\(25\)Using \([22](https://arxiv.org/html/2605.12823#A1.E22)\) and the definition𝐅AA≡−∇rℋ\\mathbf\{F\}\_\{\\text\{AA\}\}\\equiv\-\\nabla\_\{r\}\\mathcal\{H\}, we arrive at the following equation:

∇R⟨𝒪⟩R\\displaystyle\\nabla\_\{R\}\\expectationvalue\{\\mathcal\{O\}\}\_\{R\}=⟨𝚵F​∇r𝒪⟩R\+β​\(⟨𝚵F​𝐅AA​𝒪⟩R−⟨𝚵F​𝐅AA⟩R​⟨𝒪⟩R\)\\displaystyle=\\expectationvalue\{\\bm\{\\Xi\}\_\{F\}\\nabla\_\{r\}\\mathcal\{O\}\}\_\{R\}\+\\beta\\big\(\\expectationvalue\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\\mathcal\{O\}\}\_\{R\}\-\\expectationvalue\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\}\_\{R\}\\expectationvalue\{\\mathcal\{O\}\}\_\{R\}\\big\)\+⟨\(∇r⋅𝚵F\)​𝒪⟩R−⟨∇r⋅𝚵F⟩R​⟨𝒪⟩R\.\\displaystyle\\qquad\+\\expectationvalue\{\\left\(\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\\right\)\\mathcal\{O\}\}\_\{R\}\-\\expectationvalue\{\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\}\_\{R\}\\expectationvalue\{\\mathcal\{O\}\}\_\{R\}\.\(26\)We can also formulate a vectorial version of this identity; indeed, this is how we obtain the CG Hessian\. Specifically, if we consider𝓞\\bm\{\\mathcal\{O\}\}to be aDD\-dimensional vector represented as aD×1D\\times 1matrix, then we write

∇R⟨𝓞⟩R\\displaystyle\\nabla\_\{R\}\\expectationvalue\{\\bm\{\\mathcal\{O\}\}\}\_\{R\}=⟨𝚵F​∇r𝓞T⟩R\+β​\(⟨𝚵F​𝐅AA​𝓞T⟩R−⟨𝚵F​𝐅AA⟩R​⟨𝓞T⟩R\)\\displaystyle=\\langle\{\\bm\{\\Xi\}\_\{F\}\\nabla\_\{r\}\\bm\{\\mathcal\{O\}\}^\{T\}\}\\rangle\_\{R\}\+\\beta\\big\(\\langle\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\\bm\{\\mathcal\{O\}\}^\{T\}\}\\rangle\_\{R\}\-\\expectationvalue\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\}\_\{R\}\\langle\\bm\{\\mathcal\{O\}\}^\{T\}\\rangle\_\{R\}\\big\)\+⟨\(∇r⋅𝚵F\)​𝓞T⟩R−⟨∇r⋅𝚵F⟩R​⟨𝓞T⟩R,\\displaystyle\\qquad\+\\langle\{\\left\(\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\\right\)\\bm\{\\mathcal\{O\}\}\}^\{T\}\\rangle\_\{R\}\-\\expectationvalue\{\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\}\_\{R\}\\langle\{\\bm\{\\mathcal\{O\}\}^\{T\}\}\\rangle\_\{R\},\(27\)with the result treated as a3​N×D3N\\times Dobject\. To compute the Hessian, we make the following substitution \(takingD=3​ND=3N\):

𝓞→−\(𝚵F​𝐅AA\+1β​∇r⋅𝚵F\)\.\\bm\{\\mathcal\{O\}\}\\to\-\\left\(\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\+\\frac\{1\}\{\\beta\}\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\\right\)\.\(28\)After a tedious calculation making use of the fact that the AA Hessian𝐇AA≡−∇r𝐅AA\\mathbf\{H\}\_\{\\text\{AA\}\}\\equiv\-\\nabla\_\{r\}\\mathbf\{F\}\_\{\\text\{AA\}\}is symmetric \(𝐇AAT=𝐇AA\\mathbf\{H\}\_\{\\text\{AA\}\}^\{T\}=\\mathbf\{H\}\_\{\\text\{AA\}\}\), we get the following generalization of the formula for linear coarse grainings \([9](https://arxiv.org/html/2605.12823#S2.E9)\):

𝐇CG\\displaystyle\\mathbf\{H\}\_\{\\text\{CG\}\}=⟨𝚵F​𝐇AA​𝚵FT⟩R−β​𝚺​\(𝚵F​𝐅AA,𝚵F​𝐅AA\)\\displaystyle=\\langle\{\\bm\{\\Xi\}\_\{F\}\\mathbf\{H\}\_\{\\text\{AA\}\}\\bm\{\\Xi\}\_\{F\}^\{T\}\}\\rangle\_\{R\}\-\\beta\\bm\{\\Sigma\}\(\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\},\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\)−⟨𝓣⟩R−𝚺​\(𝚵F​𝐅AA,∇r⋅𝚵F\)−𝚺​\(∇r⋅𝚵F,𝚵F​𝐅AA\)\\displaystyle\\qquad\-\\expectationvalue\{\\bm\{\\mathcal\{T\}\}\}\_\{R\}\-\\bm\{\\Sigma\}\(\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\},\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\)\-\\bm\{\\Sigma\}\(\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\},\\bm\{\\Xi\}\_\{F\}\\mathbf\{F\}\_\{\\text\{AA\}\}\)−1β\[⟨𝚵F∇r\(∇r⋅𝚵F\)T⟩R\+𝚺\(∇r⋅𝚵F,∇r⋅𝚵F\)\]\.\\displaystyle\\qquad\-\\frac\{1\}\{\\beta\}\\left\[\\langle\{\\bm\{\\Xi\}\_\{F\}\\nabla\_\{r\}\\left\(\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\\right\)^\{T\}\}\\rangle\_\{R\}\+\\bm\{\\Sigma\}\(\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\},\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\)\\right\]\.\(29\)where we recall that𝚺​\(⋅,⋅\)\\bm\{\\Sigma\}\(\\cdot,\\cdot\)is the covariance matrix, and we define𝓣\\bm\{\\mathcal\{T\}\}as a matrix with components

\(𝓣\)I​J=∑i,j\(𝐅AA\)i​\(𝚵F\)I​j​∂\(𝚵F\)J​i∂rj\.\\left\(\\bm\{\\mathcal\{T\}\}\\right\)\_\{IJ\}=\\sum\_\{i,j\}\\left\(\\mathbf\{F\}\_\{\\text\{AA\}\}\\right\)^\{i\}\\left\(\\bm\{\\Xi\}\_\{F\}\\right\)\_\{Ij\}\\partialderivative\{\\left\(\\bm\{\\Xi\}\_\{F\}\\right\)\_\{Ji\}\}\{r\_\{j\}\}\.\(30\)The first line of \([29](https://arxiv.org/html/2605.12823#A1.E29)\) is the same as in the linear case and reduces to it\. The remaining terms all stem from the nonlinearity of𝝃r\\bm\{\\xi\}\_\{r\}and vanish if𝚵F\\bm\{\\Xi\}\_\{F\}is constant\.

As a sanity check, we prove that the right\-hand side of \([29](https://arxiv.org/html/2605.12823#A1.E29)\) is symmetric\. The entire first line, the sum of covariances on the second line, and the covariance on the third line are all manifestly symmetric, so the only problematic piece is−⟨𝒯⟩R−β−1⟨𝚵F∇r\(∇r⋅𝚵F\)T⟩R\-\\expectationvalue\{\\mathcal\{T\}\}\_\{R\}\-\\beta^\{\-1\}\\langle\{\\bm\{\\Xi\}\_\{F\}\\nabla\_\{r\}\\left\(\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\\right\)^\{T\}\}\\rangle\_\{R\}\.

To manipulate this sum, we rewrite it in integral form and integrate by parts\. Doing so in index notation \(used for clarity\) yields

\(−⟨𝓣⟩R−1β⟨𝚵F∇r\(∇r⋅𝚵F\)T⟩R\)I​J\\displaystyle\\left\(\-\\expectationvalue\{\\bm\{\\mathcal\{T\}\}\}\_\{R\}\-\\frac\{1\}\{\\beta\}\\langle\{\\bm\{\\Xi\}\_\{F\}\\nabla\_\{r\}\\left\(\\nabla\_\{r\}\\cdot\\bm\{\\Xi\}\_\{F\}\\right\)^\{T\}\}\\rangle\_\{R\}\\right\)\_\{IJ\}\(31\)=1β​Z​∑i,j∫𝑑𝐫​\[∂∂ri​δ​\(𝝃r​\(𝐫\)−𝐑\)​\(𝚵F\)I​j​∂\(𝚵F\)J​i∂rj\+δ​\(𝝃r​\(𝐫\)−𝐑\)​∂\(𝚵F\)I​j∂ri​∂\(𝚵F\)J​i∂rj\]​e−β​ℋ​\(𝐫\)\.\\displaystyle\\ \\ =\\frac\{1\}\{\\beta Z\}\\sum\_\{i,j\}\\int d\\mathbf\{r\}\\,\\left\[\\frac\{\\partial\}\{\\partial r\_\{i\}\}\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)\(\\bm\{\\Xi\}\_\{F\}\)\_\{Ij\}\\partialderivative\{\(\\bm\{\\Xi\}\_\{F\}\)\_\{Ji\}\}\{r\_\{j\}\}\+\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)\\partialderivative\{\(\\bm\{\\Xi\}\_\{F\}\)\_\{Ij\}\}\{r\_\{i\}\}\\partialderivative\{\(\\bm\{\\Xi\}\_\{F\}\)\_\{Ji\}\}\{r\_\{j\}\}\\right\]e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\}\.The second term in the brackets is symmetric with respect to the free indicesI​JIJ, since we can freely swap the summation indicesi​jij\. As for the first term, we must use the fact that

∂∂ri​δ​\(𝝃r​\(𝐫\)−𝐑\)=−∑K\(𝐉ξ\)K​i​∂∂RK​δ​\(𝝃r​\(𝐫\)−𝐑\),\\frac\{\\partial\}\{\\partial r\_\{i\}\}\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)=\-\\sum\_\{K\}\(\\mathbf\{J\}\_\{\\xi\}\)\_\{Ki\}\\frac\{\\partial\}\{\\partial R\_\{K\}\}\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\),\(32\)as well as the following identity \(whereδJ​K\\delta\_\{JK\}here is the Kronecker delta\):

∑i\(𝚵F\)J​i​\(𝐉ξ\)K​i=δJ​K\.\\sum\_\{i\}\(\\bm\{\\Xi\}\_\{F\}\)\_\{Ji\}\(\\mathbf\{J\}\_\{\\xi\}\)\_\{Ki\}=\\delta\_\{JK\}\.\(33\)This is simply the equation𝚵F​𝐉ξT=𝐈3​N\\bm\{\\Xi\}\_\{F\}\\mathbf\{J\}\_\{\\xi\}^\{T\}=\\mathbf\{I\}\_\{3N\}in index form\. \([33](https://arxiv.org/html/2605.12823#A1.E33)\) can be differentiated to obtain

∑i∂\(𝚵F\)J​i∂rj​\(𝐉ξ\)K​i\+∑i\(𝚵F\)J​i​∂\(𝐉ξ\)K​i∂rj=0\.\\sum\_\{i\}\\frac\{\\partial\(\\bm\{\\Xi\}\_\{F\}\)\_\{Ji\}\}\{\\partial r\_\{j\}\}\(\\mathbf\{J\}\_\{\\xi\}\)\_\{Ki\}\+\\sum\_\{i\}\(\\bm\{\\Xi\}\_\{F\}\)\_\{Ji\}\\partialderivative\{\(\\mathbf\{J\}\_\{\\xi\}\)\_\{Ki\}\}\{r\_\{j\}\}=0\.\(34\)By making use of \([32](https://arxiv.org/html/2605.12823#A1.E32)\) and \([34](https://arxiv.org/html/2605.12823#A1.E34)\) in \([31](https://arxiv.org/html/2605.12823#A1.E31)\), the first term becomes

1β​Z​∑i,j,K∂∂RK​∫𝑑𝐫​δ​\(𝝃r​\(𝐫\)−𝐑\)​e−β​ℋ​\(𝐫\)​\(∂2∂ri​∂rj​\(𝝃r\)K\)​\(𝚵F\)I​j​\(𝚵F\)J​i,\\frac\{1\}\{\\beta Z\}\\sum\_\{i,j,K\}\\frac\{\\partial\}\{\\partial R\_\{K\}\}\\int d\\mathbf\{r\}\\,\\delta\(\\bm\{\\xi\}\_\{r\}\(\\mathbf\{r\}\)\-\\mathbf\{R\}\)e^\{\-\\beta\\mathcal\{H\}\(\\mathbf\{r\}\)\}\\left\(\\frac\{\\partial^\{2\}\}\{\\partial r\_\{i\}\\,\\partial r\_\{j\}\}\(\\bm\{\\xi\}\_\{r\}\)\_\{K\}\\right\)\(\\bm\{\\Xi\}\_\{F\}\)\_\{Ij\}\(\\bm\{\\Xi\}\_\{F\}\)\_\{Ji\},\(35\)where\(𝝃r\)K\(\\bm\{\\xi\}\_\{r\}\)\_\{K\}is theKKth component of the coarse\-graining map\. Just like the second term in \([31](https://arxiv.org/html/2605.12823#A1.E31)\), this expression is symmetric with respect to the free indices, so the proof is complete\.

## Appendix BTraining and validation loss comparisons

![Refer to caption](https://arxiv.org/html/2605.12823v1/figures/fig_train_loss_curves.png)Figure 3:Training loss decomposition\.Left:FM and FM\+AAp\{\}\_\{\\text\{p\}\}training losses, where the∼\\sim38 kcal2mol\-2Å\-2offset reflects the HVP matching term \(wHVP​‖𝐇CG​𝐯−𝐇θ​𝐯‖2w\_\{\\mathrm\{HVP\}\}\\\|\\mathbf\{H\}\_\{\\mathrm\{CG\}\}\\mathbf\{v\}\-\\mathbf\{H\}\_\{\\theta\}\\mathbf\{v\}\\\|^\{2\}\)\.Right:FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov training loss, dominated by the covariance correction \(Term 2\)\. Despite the order\-of\-magnitude difference in training\-loss scales, all three models achieve equivalent force prediction accuracy on held\-out data \(Figure[4](https://arxiv.org/html/2605.12823#A2.F4)\)\.![Refer to caption](https://arxiv.org/html/2605.12823v1/figures/fig_val_loss_annotated.png)Figure 4:Validation loss \(force\-matching MSE\) during training on a single\-chain dataset \(99 proteins\) atT=300T=300K\. All three objectives—FM, FM\+AAp\{\}\_\{\\text\{p\}\}, and FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov—converge to comparable validation performance \(∼\\sim590\.9 kcal2mol\-2Å\-2\), confirming that the auxiliary curvature\-matching terms do not degrade force prediction accuracy\. Models were trained for 52 \(FM\), 57 \(FM\+AAp\{\}\_\{\\text\{p\}\}\), and 62 \(FM\+AAp\{\}\_\{\\text\{p\}\}\+Cov\) epochs withwHVP=0\.01w\_\{\\text\{HVP\}\}=0\.01and batch size 200 on 4×\\timesA40 GPUs\.
## Appendix CSingle\-chain training set

We use a curated dataset of 99 single\-chain proteins, constructed to include at least 2 proteins containing each bond type, 4 reference proteins \(6MRR, 2VQ4, 5YM7, 3ID4\), and additional proteins to bring the total to 99\. The dataset has been filtered to exclude short \(cis conformation\) bonds; disulfide bonds were excluded at the base dataset level\. The full set contains 15,522 Cαatoms:

> 1AAJ 1ACF 1AMM 1AYD 1BK2 1BZ4 1CM2 1DVN 1EQ6 1EYH 1FNA 1GYV 1I2T 1IO2 1J7X 1JMW 1JW4 1K40 1LPJ 1MJC 1N81 1NG6 1OGW 1P1L 1QTP 1QZM 1UKF 1W8V 1WP5 1WWI 1WY3 1X3O 1XGD 1XT0 1YP5 1YU5 1YW5 1ZEQ 1ZLM 2BFH 2EVB 2H0M 2ICT 2NSC 2OJ4 2RH3 2V14 2VQ4 2W9Q 3HVM 3ID4 3KJE 3LFO 3MX7 3MZZ 3P7K 3PG4 3R6D 3RJP 4BPF 4J5Q 4N6T 4QAJ 4QBE 4RZ9 4U3H 4Y2K 5BTH 5GY3 5IHW 5J2V 5UVR 5XEF 5Y8E 5YM7 6APK 6DR3 6GPM 6KND 6L7Q 6LG3 6MRR 6Q3V 6SYG 6TYY 6WEY 6WIN 7DMF 7DMS 7LIQ 7TGP 7XCD 8A8S 8ERE 8G22 8GQQ 8HQW 8I2D 8JED