Representing local protein environments with machine learning force fields
Authors: Meital Bojan
Sanketh Vedula
Sai Advaith Maddipatla
Nadav Bojan
Anar Rzayev
Federico Napoli
Paul Schanda
Alexander Bronstein
Venue: ICLR 2026 Poster (2026)
Keywords: Machine learning force fields
structural biology
NMR
representation learning
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We show that embeddings from machine learning force fields provide rich, transferable representations of local protein environments, enabling zero-shot generalization and state-of-the-art downstream performance.

Abstract: The local structure of a protein strongly impacts its function and interactions with other molecules. Representing local biomolecular environments remains a key challenge while applying machine learning approaches over protein structures. The structural and chemical variability of these environments makes them challenging to model, and performing representation learning on these objects remains largely under-explored.  In this work, we propose representations for local protein environments that leverage intermediate features from machine learning force fields (MLFFs). We extensively benchmark state-of-the-art MLFFs—comparing their performance across latent spaces and downstream tasks—and show that their embeddings capture local structural (e.g., secondary motifs) and chemical features (e.g., amino acid identity and protonation state), organizing protein environments into a structured manifold. We show that these representations enable zero-shot generalization and transfer across diverse downstream tasks. As a case study, we build a physics-informed, uncertainty-aware chemical shift predictor that achieves state-of-the-art accuracy in biomolecular NMR spectroscopy. Our results establish MLFFs as general-purpose, reusable representation learners for protein modeling, opening new directions in representation learning for structured physical systems.

Published as a conference paper at ICLR 2026
REPRESENTING LOCAL PROTEIN ENVIRONMENTS WITH MACHINE LEARNING FORCE FIELDS
Meital Bojan". Sanketh Vedula*.4.3, Advaith Maddipatla Nadav Sellam',
Anar Rzayev', Federico Napoli', Paul Schanda', Alex Bronstein!.2
IST Austria Techhion, Israel 3 University of Oxford, UK
*Princeton University BBrodd Institute MIT and Harvard "qqual contribution
ABSTRACT
The local structure of protein strongly impacts its function and interactions with other molecules. Representing local biomolecular environments remains a key challenge while applying machine learning approaches over protein struc- tures, The structural and chemical variability of these environments makes them challenging to model, and performing representation learning on these objects remains largely under-explored. In this work, we propose representations for local protein environments that leverage intermediate features from machine learn- ing force fields (MLFFs). We extensively benchmark state-of-the-art MLFFs, comparing their performance across latent spaces and downstream tasks, and show that their embeddings capture local structural (c.g., secondary motifs) and chemical features (e.g., amino acid identity and protonation state), organizing protein environments into a structured manifold. We show that these represen- tations enable zero-shot generalization and transfer across diverse downstream tasks. As case study, we build a physics-informed, uncertainty-aware chemi- cal shift predictor that achieves state-of-the-ar accuracy in biomolecular NMR spectroscopy. Our results establish MLFFs as general-purposee reusable repre- sentation learners for protein modeling, opening new directions in representa- tion learning for structured physical systems. Code and data are available htttps://qithun: com 2 representation
INTRODUCTION
Proteins are complex three-dimensional structures composed of hundreds to thousands of atoms, arranged in a specific manner in space. Local environments within a protein are highly diverse due to variability in the amino-acid sequence and the folding of the chain into a 3D structure. These factors govern a protein's 's functional mechanisms such as ligand binding, catalysis, and allostery, and make proteins uniquely challenging compared small molecules. Learning compact and transferable representations of local protein environments is a central challenge for performing machine learning (ML) over biomolecules: the difficulty lies in jointly encoding the local chemical context, including atomic identities, bonds, and subtle biochemical properties, into a consistent and generalizable representation that can transfer across diverse protein modeling tasks.
Classical approaches employ hand-crafted descriptors that partially capture this information by embedding features such as dihedral angles, hydrogen bonds, or electrostatic terms, which often limits generalization across proteins and tasks. In computational chemistry. descriptors like Parrinello- Behler symmetry functions (Behler 2011] Behler & Parrinello/ 2007 Behler 2016) are widely used to represent molecular environments (Jaeer et al. (2018). These methods encode atomic interactions and geometry into concise, invariant representations suited for geometric modeling. Modern atomistic MI. techniques (Schütt al.) 2018 2021) implicitly learn similar representations with neural networks, achieving accuracy comparable to density functional theory (DFT) (Deng et al., 2023) and are used for molecular dynamics (MD) simulations. Recent advances have shifted heural-networl interatomic
Proteins are polymers of amino acids (also called residues). Residues share common backbone of four heavy atoms (N, CA, C, O) and differ in their side chains.

---

Published as a conference paper at ICLR 2026
Dalabase DFT energies
Training losses
Energy/torce regressor
pre-training
Local environment
GA
MLFF
Downstream models
properties
Focus residue
transfer earning
representation
environment
Oth, 2nd order descriptors
Figure I: Proposed construction of canonical local protein environment descriptors and their use. Machine learning force field (MLFF) models are pre-trained as energy and force regressors on databases of DFT-calculated energies, enabling them to learn zeroth, first, and second order latent representations of interatomic interactions. We embed local protein environments by extracting latent embeddings from pre-trained MLFFs for a focus residue and all atoms within 5A radius of the residue. These embeddings are then mapped onto the atoms of the focus residue to construct canonical environment descriptors which can be used in downstream models for transfer learning to predict diverse chemical properties.
potentials from element-specific models to universal machine learned force-fields (MLFFs) trained on large datasets of molecules with DFT-simulated energies. Modern MLFFs train large neural networks to achieve DFT-level accuracy on millions of calculations. with representative families including AIMNet (Smith et al., 2018). MACE (Batatia et al. 2022 Kovács et al., 2025), and OrbNet (Qiao
al. 2020).
Contributions and Summary of Results. The key idea of this work is to repurpose MLFF em- beddings from their conventional role in energy and force prediction to representing local chemical environments in proteins. We show that MLFF's learn compact and generalizable representations of atomistic systems that can be reused for variety of downstream tasks. MLFF embeddings, unlike sequence- sequence-based representations (e.g.. ESM (Lin et al. (023)), are trained on quantum mechunncca data and encode physically grounded features such as bond geometry, torsions, and electronic interac- tions. Because these embeddings are defined per atom, they naturally transfer to unseen proteins. The atomic building blocks remain the same even when sequences or folds differ, enabling generalization to rare chemistries and out-of- of-distribution conformations. Analogous to foundation models in vision and language, these representations enable zero-shot and transfer learning, support data-driven priors over biomolecular environments, and position MLFFs as foundation models for structural biology. Outlined below are our major contributions and a summary of results. To assist the reader, we provide succinct summary of the research questions, experimental designs, and key findings in Table B.2
Canonical local environments and MLFF representations. Protein environments vary widely, making their direct comparison challenging. To make MLLF representations comparable across residues and proteins, we introduce canonical environments: regions centered on a focus residue and containing every amino acid whose atom coordinates lie within 5 (Hausdorff distance) of a focus residue. We then construct transferable environment representations from the atom-wise representations of constituent atoms (or subset thereof). To our knowledge, this is the first use of MLFFs to characterize local protein structure.
MLFF representations effectively capture local protein structure and chemistry. By analyzing over 165K protein environments extracted from 1048 non-redundant chains, we demonstrate that MLFF embeddings meaningfully organize biochemical information without any additional training: unsupervised clustering of these embeddings reveal secondary structure features and amino acid identities (see Fig. 2) Furthermore, we use these embeddings for transfer learning by training networks to predict protonation states calculated with Poisson-Boltzmi
solver (Reis et al. 2020). The resulting models based on MLFF embeddings especially our best AIMNet-feature model achieve the lowest errors simulated pK, values, outperforming classical frameworks such as PKa-ANI (Gokcan & Isayev] 2022) and PropKa (Olsson et al.) 201T), as well as GCNs built on learned or ESM-derived embeddings (see Table
2

---

Published as a conference paper at ICLR 2026
MLFF representations allow computation of similarity metrics and calculating likelihoods. We show that MLFF embedding space is well-structured. Specifically, distances between environment embeddings reflect both geometric similarity and chemical context, enabling intuitive clustering of related local motifs. We construct likelihood and conditional likelihoods in the MILF embedding space and show that they effectively capture the distribution of biomolecular environments (see Fig. A.7 A. We show that the defined likelihoods capture subtle structural deformations in protein structure and allow capturing distribution shifts, making them valuable for structural quality assessment, uncertainty estimation, anomaly detection, dataset curation, and potentially in guiding Boltzmann generators (Noé et al., 2019). Furthermore, leveraging the MLFF embedding space, we introduce a similarity metric to compare protein environments (see Appendix H.2
MLFFss enable transfer learning of physics-grounded, uncertainty-aware chemical shift predictor. In the context of biomolecular nuclear magnetic resonance (NMR) spectroscopy, we leverage MLFF embeddings and train downstream networks to predict protein chemical shifts, and demonstrate that this approach outperforms the current ofate-ofthhe-at predictor UCBShift2-X for both backbone and side-chain nuclei (see Fig. 3 We conduct case study examining the ring current effects produced by aromatic side chains on chemical shifts of neighboring nuclei, and demonstrate that our predictor follows physically consistent trends, whereas UCBShift2-X shows unphysical behavior (see Fig. ( We also pair the predictor with MLFF-derived likelihoods to assign confidence scores, yielding reliable uncertainty measure around the predicted shifts (see Fig. A.7
Interpreting MLFF representations. We probe the content of MLFF embeddings in two controlled settings. Firstly, we systematically modify side chain orientations and track how the embeddings respond (see Fig. And 6). Secondly, we follow an unfolding simulation in which a seven-residue a-helix extends into a strand (see Fig. A.9 A.8 and A... In both cases, prineipal-componett analysis uncovers clear latent directions that mirror the underlying structural changes. Finally, to understand the information contained within the MLFF embeddings, we ask the following question: can we invert the MLFF embedding to recover the underlying local protein environment? Casting this question as an inverse prohlem, we guide (Maddipatla et al., 2025 2024 Levy et al., 2025)
AlphaFold3 (Abramson et al.) 2024) and reasonably recover protein conformations underlying a target MLFF embedding. This partial recovery suggests that MLFF embeddings encode necessary information required to invert a local protein conformation (see Appendix. (.2).
Benchmarking different MLFFs. We benchmark three model families MACE, OrbNet, and AIMNet on tasks defined on protein environment that include secondary structure assignment, amino acid identification, acid dissociation constant (pK,) and chemical shift prediction (see Ta-
B.4 B.3 B.5 B.6) MACE--base embeddings perform best on every task except pka prediction, where A Met exhibits superior performance (see Table [T.
Together, these steps recast MLFFs as general-purposs biochemical feature extractors for proteins, providing canonical representations, principled similarity/likelihodd tools, and state-of-the-art down- stream performance with calibrated uncertainty. The remainder of the paper is organized as follows. Section Section[jintroduee intrrduuess MLFFs and reviews closely related work, and Section describes the construc- tion of canonical local environments. Sections 6 demonstrate that the proposed representations capture protein structure and chemistry and enable state-of-the-arr predictors. Section 7 interprets the learned representations, and Section 8 outlines limitations and future directions. Due to space con- straints, the Appendix contains a significant amount of supplementary material, including additional figures, tables, experiments, an extended background, and full experimental details. We encourage the reader to consult the table of contents in the Appendix for a complete overview.
2 BACKGROUND & RELATED WORK
MLFFFs. Machine learning firee- (MLFFs) are neural networks trained to reproduce quantum- mechanical potential energy surfaces. Given atomic coordinates and elements as inputs, MLFFs apply conmetry-preserving encodings (e.g.. local environments, equivariant basis functions), and output per-atom energy contributions whose sum yields the total potential energy. Consequently, forces are obtained as the negative energy gradient, enabling molecular dynamics (MD) simulations with near-quantum accuracy at orders of magnitude lower cost than direct DFT calculations. In this work, we focus on three representative MLFF families with distinct design principles: MACE (Batatia et al. 2022) and Egret (Wagen et al. 2025) combine equivariant message passing with high-order
3

---

Published as a conference paper at ICLR 2026
Table 1: Acid Dissociation constant (pKa) prediction accuracy. Reported are mean absolute errors (MAE) with respect to pK, values estimated using the Poisson-Bolzzaan solver PypKa (Reis et al.)
2020) on protonating / deprotonating amino acids. The table compares (i) classical pKa predictors
(Ossson et al. 2011 Gokcan & Isayev 2.22), (ii) a GCN with learned embeddings ("Learned"), (iii)
ESM-bbaed features GCN (Ouyang- Zhang et al. [025 Hayes et al. 2025] Lin et al.) 2023). and (iv)
MLFF-bbaed features + GCN (ours). PropKa and pKa-ANI were originally trained on experimental
pK, values and are therefore not specifically optimized for the computational reference used here.
The best performing model for each metric is indicated in bold.
ESM-Features GCN
MLE Feetturss + GCN (Ours)
Residues
Samples
PropKa
pKa-ANI
Leurned
ESM3 (Seq)
ISM
ESM (Struc)
ESMFuld
MACE
OrbNet
ATMNet
Egret
Glutamic acid
13592
0.551
0.465
0.518
0.159
0.382
0.360
0.351
0.306
0.306
0.265
0.301
Asgurtic acid
084]
0:160
0.582
0.528
0.368
0.280
Lysine
12009
0.393
0.339
0.255
0.278
0.320
listidine
1182
0.561
0.426
0.815
0.488
0.428
0.408
0.383
0.421
0.441
0.380
0.440
many-body interatomic potentials: OrbNet (Qiao et al., 2020) augments graph neural networks with semi-empirical orbital features to achieve DFT-level accuracy with reduced data requirements; and AIMNet (Zubatyuk et al. (019) extends ANI with learned atomic embeddings and multitask training (energies, charges, spin), improving transferability to charged and open-shell systems. While these models achieve remarkable accuracy for small molecules, materials, and bipeptides, MLFFs continue face challenges in capturing long-range interactions in larger biomolecular systems like proteins. Further details on these models and a short background on MD and DFT is provided in Appendix D
Related work. Recent work has highlighted the versatility of pretrained MLFFs as learned represen- tations for molecular tasks. [Shiota et al. (2024) use MACE and M3GNet features to predict chemical shifts in small molecules. while ElijoSius et al. (2025) employ MACE embeddings for zero-shot molecular generation via evolutionary search. However these works are limited to small-molecules and not scale to large, structurally heterogeneous biomolecules like proteins. Similar our idea, Gokcan & Isayev (2022) extract ionizable residues and use ANI-2x construct per-residue atomic environment vectors (AEVs), which are then used as input to train a regressor for pK, prediction. However, these AEVs are fixed, hand-crafted descriptors derived from symmetry functions that do not use message passing over the local environment. As a result, they lack the ability to adaptively capture the context-dependent interactions present in protein systems. In contrast, our descriptors encode a richer biochemical context. Moreover, we train a GCN over this many-body embedding space, enabling context-aware model. As shown in Table [ our method achieves lower pk, prediction error and generalizes to additional biochemical tasks beyond pK.-
3 REPRESENTING PROTEIN LOCAL ENVIRONMENTS WITH MLFFs
Notation. We denote a residue in a protein as a, and the all-atom representation of the protein structure as X {(x1,=1) ... (Xm, }. where Xj R& and zj are the location and atomic number of jth atom in X', respectively. Let fo X + be the MLFF with parameters 0, and y = {(y1),.... (ym)) the atom-wise embeddings of the MLFF, where y E Rdd. The local enviromment of a focus residue is denoted by X₂ C X. The subset of atoms in residue a used for predicting biochemical properties denoted as Ag. Given a subset of atom indices Ag, we denote by YA. = { (yx) | E Ao} y the restriction of y to those atoms.
We seek representation of local protein environment to be (i) sensitive to local changes, (ii) insensitive to global variations, (iii) fast to compute, (iv) canonical, enabling direct comparison across diverse environments, and (v) effective and generalizable to unseen environments. Encoding full proteins with several thousand atoms with MLFFs is computationally inefficient and redundant since we are solely interested in locally sensitive features. To construct such representations, we propose to construct local environment around focus residue to be encoded later by the MLFF.
Constructing local environments. To balance computational efficiency while retaining local struc- tural context, given a focus residue a, we construct the environment X,, as the union of all residues whose atoms are at most 5 away from the atoms of a. The procedure is described in App. Alg. and an ablation over the choice of radius is provided in Appendix
MLFF representations. Given an environment X,, MLFFs produce atom-level features over the layers of the network. These learned atomistic features are contextual, influenced not only by the
4

---

Published as a conference paper at ICLR 2026
atom alone but also by the other atoms in the environment due to the message-passinn passing operations within MLFFs. Given Xa. atom-wise feature representations are extracted from the final layer of the MLFFF Each MLFF produces representations of shape N x d, where N = |x,l and d is the dimension of the representation. To obtain canonical representations that are comparable across residues, we retain only V'A., the representations corresponding to atoms in A. Given protein structure, this process is applied to each residue in the sequence to obtain residue-wise features. A visual depiction of this procedure is presented in Fig.
Data, To ensure a consistent dataset across all experiments, we sourced data from RefDB, curated subset of the Biomolecular Magnetic Resonance Databank (BMRB). The final dataset includes 1048 BMRB entries, filtered for sequence redundancy and split into 823 training and 225 test proteins. From these structures, a total of 165K local protein environments were extracted using a 5A radius threshold based on Hausdorff distance. We use AlphaFold2 (Jumper et al. 2021) predicted 3D structures for all sequences, followed by hydrogen addition and Amber99 fiele-field relaxation (Hornak et al., (006). We refer to Sec. E for a detailed justification of our choice, and for further details on data curation.
Baselines To evaluate the utility of MLFF embeddings, we compare the performance of MLLF embeddings GCN against different embedding baselines: (i) sequence-model embeddings (ESM3
(Hayes et al. 2025) ) and sequence models with enhanced structural representations (Ouyang-Zhane
et al. 2025 + GCN, (ii) ESMFold (Lin et al. 2023) structure-mode embeddings + GCN, (ii)
hand-crafted statistics-based descriptors (LOCO-HI) (Fazekas et al. (024) + GCN, and (iv) simple learned embeddings GCN. In the leamnddembeddinng baseline, the GCN is trained end-to-end from atom types and 3D Cartesian coordinates from AlphaFold2 [
Jumper et al. (202I). The coordinates are first mapped to a sinusoidal positional encoding, which is then concatenated with the embedded atom-type identifiers to produce the node features for the GCN.
4 MLFF REPRESENTATIONS CAPTURE LOCAL PROTEIN STRUCTURE AND CHEMISTRY
In what follows, we first present a zero-shot analysis of MILF representations to assess whether they inherently capture biochemically meaningful information without task-specific training. We then apply transfer learning by training neural networks on top of the frozen embeddings to predict structural and chemical properties. We first consider amino acid identity and secondary structure prediction as motivating examples, and then progress to more challenging downstream regression tasks such as pK, and NMR chemical shift prediction.
Amino acid identity & Secondary structure. A protein's 's primary structure is its linear sequence of amino acids linked by covalent peptide bonds and implicitly encode the 3D protein structure. Its secondary structure is a local motif of backbone conformation, stabilized by hydrogen bonds between non-adjacent residues (Kabsch & Sander 1983), with the most common motif's being a-helices (H) and B-sheets (E).
To examine the information encoded by the MLFF representations in zero-shot setting, the atomistic features from MACE were projected into a two dimensional space using Uniform Manifold Approxi- mation and Projection (UMAP) (Mcinnes al. (0118). The resulting embeddings, shown in Fig. 2 are annotated with amino acid labels, secondary structure, and dihedral angles. Notably, features corresponding to a-heliess and f-sheets form distinct clusters in the UMAP space. A similar pattern of separation is observed when embeddings are annotated by amino acid identity and backbone dihedral angles, indicating that chemically and structurally distinct features are well-represented in the feature space.
To quantitatively assess this capability, we perform lightweight transfer learning by training classifiers to predict the secondary structures and the amino acid identities from the frozen embeddings. Further details on the loss function, target label space, and the definition of A, can be found in Appendix G.I while a comprehensive description of the model architecture is provided in Appendix G.3 Classifier accuracy was evaluated using precision, recall, and F1 scores. Results for secondary structure prediction are summarized in Table B.3 while amino acid classification results are presented in Table B.4 For secondary structure, models trained with MACE and Egret features consistently
"See Appendix F2 for ablation study comparing representations extracted from different MI.FF layers.
5

---

Published as a conference paper at ICLR 2026
achieve superior prediction average. For amino acid prediction, models trained using Egret features consistently outperform those trained with alternative representations.
UMAP
Secondary structure
Amino acid identity
M
W
Dihecral angles
CA secondary shift [ppm]
180
2
12
+4
-00
-100"
-90"
+90' -180"
Figure 2: MACE embedding space reveals meaningful structural and chemical information. Depicted two-dimensional UMAP coordinates of 165, 913 protein environments from 1327 non- redundant chains predicted by AlphaFold2 (Jumper et al. 2021), labeled left-to-right, top-to-bottom according to the DSSP secondary structure class (Table B.12), amino acid chemical identity, the pair of backbone dihedral angles (ø, 1). and CA secondary chemical shift (relative to a random coil).
Acid dissociation constants (pK.). рK, quantifies the proton-donation propensity of a titratable group, i.e., its acidity. It is the pH at which the group is half protonated across the ensemble. In proteins, pK, values often deviate from their canonical solution values due to local electrostatics and solvent accessibility. These context-dependent shifts make pK, prediction strong test of

---

Published as a conference paper at ICLR 2026
Chemical shift pred prediction error [ppm]
Mean (SS)
UCB Shift
Mean (AA)
MACE
10
10
10-1
10
HA CA
C
CD
CD2
CG2
CE
Atom
Figure 3: Chemical shift prediction errors for different atom types evaluated test set of 132, 228 environments from 203 non-redundant BMRB records with experimentally determined chemical shifts used as the reference. The median prediction error in ppm and the 25% 75% (boxes) and 5% - 95% (whiskers) confidence intervals are depicted.
whether atomistic representations capture local chemical environments (see App. D for biochemical background).
For quantitative assessment, regression model was trained amino acids with ionizable, proton- transferring side chains to predict their pK, values. Ground-trut pk, values were estimated using a Messon-Bolttzaan solver (Reis al. 2020). The evaluation focused on glutamic acid (GLU) and aspartic acid (ASP), which tend to deprotonate, and lysine (LYS) and histidine (HIS), which tend protonate. As shown Table models trained with AIMNet features achieve the best overall performance among all feature types because is trained to predict multiple molecular properties simultaneously. For additional comparisons with PropKA (Olsson al., 2011) and pKa-ANI (Gokean & Isayev 2022) on experimental pK, values and PDB structures, as well as training details, see Appendix
LIKELIHOODS AND SIMILARITY OF LOCAL PROTEIN ENVIRONMENTS
The MLFF embedding space encodes local biomolecular structure, capturing structural and chemical context. We leverage the richness of MLFF embedding spaces to define a distribution over biomolec- ular environments. Given of reference environments & {X1,- -....... Xr,) and reference atom set A, we define the likelihood of an environment X, at focus residue a as
(.fo(aa)bz) - fo(Ka )bo
p(Xa)=
Σ
exp
Everl
202
N..
where a refers to the bandwidth, and it controls the influence of each reference environment on Xa, This is equivalent to performing kernel-density estimation in the MLFF embedding space, with a radial basis function kernel. Intuitively, the likelihood measures how typical X. is among the reference environments Egt. Subsequently, the conditional likelihoods can defined similarly to Eq. bby curating an appropriate set of reference environments that satisfy the chosen conditioning.
Capturing distribution shifts. To evaluate the quality of unconditional likelihood estimation and its ability to detect subtle distribution shifts in the local structure, we curated a set of reference environments randomly sampling 10, 000 environments from 1, 100 protein structures that were relaxed using the Amber99 force field (Hornak et al.) 2006). We then measured the likelihoods of each environment sampled from test set of proteins, before and after performing Amber99 relaxation. The results shown in Fig. A.7 demonstrate that the likelihood function is sensitive to subtle conformational changes. Relaxed structures consistently receive higher likelihoods, and the distribution of the paired differences captures fine-grained structural variations. This makes the approach well-suited for detecting out-of-distribution conformations and assessing local structural quality.
Secondary structure conditioned likelihoods. To further evaluate the sensitivity of conditional likelihoods, randomly sampled 1000 environments from secondary structures annotated

---

Published as a conference paper at ICLR 2026
10
MAGE
UCB S
58.0
MAGE
CM9
UCB Sm
CA chemmaal shitt [opm)
GA63
CA GAS1
47.5
Ring orientation
Figure 4: Synthetic example showing the influence of a phenylalanine sidechain aromatic ring surrounding chemical shifts. A. The magnitude of change of backbone CA chemical shifts over different ring orientations as predicted using the proposed MACE-based predictor (left) and UCBShifi-X (right). A 7À sphere indicates the radius from the ring center which the influence of the ring current is expected to become negligible. Note that UCBShift-X predicts much longer-range. albeit small, ring influence extending beyond 20A. B. Locations of three nearby CA atoms; and C. their predicted chemical shifts vs. the ring orientation. Note the smooth 180% -p0--oeio behavior of the MACE shift prediction and the decay the effect scale with the distance from the ring.
helix (H), B-strand (E), and turn (T) by DSSP, and constructed conditional likelihoods p( p(H)H), EE) and (T, respectively. Fig. A.11 depicts the projections of conditional likelihood triplets (log ((X{E)), P(X H)), log (T) of 8, 163 environments from the test set. The results highlight clear separation by true secondary structure of the environment. The one-dimensional marginal distributions further emphasize the distinct statistical profiles of each structural class.
We further direct the reader to Fig. A.13 embeddings to evaluate chemical environment similarity.
I.Jd and Appendix H.2 for for detailed analysis of the use of MLFF
PHYSICS-GROUNDED AND UNCERTAINTY-AWARE CHEMICAL SHIFT PREDICTION FOR PROTEINS
Background and prior art We refer an uninitiated reader to Appendix D for a short background on chemical shifts and biomolecular NMR. Several computational methods have been developed
to predict NMR chemical shifts from molecular structures (Shen & Bax 2008 Neal et al. 2003 Han al. 2011 Li Brüschweiler 2012 2015 Kohlhoff al. 2009 [ [203 Moon Case 2007). with recent approaches increasingly relying machine learning. The current state-of-the-art, UCBShift (Zhou et al. 2023), combines sequence and structure alignment with a random forest model, using reference proteins and structural descriptors to guide predictions. Its latest iteration, UCBShift 2.0 (Ptaszek et al.) 2024), extends this framework to include side-chain atoms alongside backbone predictions. Despite their effectiveness, these methods depend reference-based similarity measures, which limits their generalizability.
MLFF-based shift predictor. Chemical shift prediction is formulated transfer learning task over frozen MLFF embeddings, where graph neural network takes the pretrained atomistic representations input and predicts the chemical shift of target atom. Separate models trained for the N, CA,
8

---

Published as a conference paper at ICLR 2026
CA cremical shitt prediction entrr togin
1.0
2.5
chemmica predictoo
2.0
2.5
0.4
1.0
lkolioodd
Hetlhood
Figure 5: Lower likelihood environments result in larger chemical shift prediction error. Chem- ical shift prediction accuracy of CA (left) and N (right) atoms stratified by the KDE-estimated likelihood of the corresponding MACE descriptors. Depicted are the median and 25%-500% confi- dence intervals. Higher-iicelihodd environments correspond to lower prediction error and can be used an uncertainty measure.
C, H. and HA backbone atoms, well as the CB, CG, CD, CD2, CG2, and CE side-chain atoms. Dataset curation, model architecture, and training details are provided in Appendix
G.3
As evident from Fig. the proposed predictor outperforms UCBShift2 for both backbone and side- chain heavy atoms, with the exception the Hydrogen alpha atom. A comprehensive quantitative comparison of chemical shift prediction accuracy across different MLFF embeddings, ESM variants, and learned embeddings is provided in Tables B.5 and consistently achieve the best performance.
B.6 We observe that MACF-based models
To further probe whether the shift predictor correctly captures local structural effects, we design three case studies: one is shown here, and the other two provided in Appendix [ for completeness.
Effects of ring currents. An external magnetic field induces a flow of delocalized æ electrons in an aromatic ring, which in turn produces its own magnetic field. Because chemical shifts depend on the net magnetic field produced by the local structural environment, they are strongly influenced by such ring currents. As the phenylalanine ring has C2 symmetry, states separated by rotation of the ring by 180 are indistinguishable. To probe whether our shift predictor captures ring-current effects accurately, we synthetically modified the side chain of Phe (residue 50 PDB ID: 12V6) by rotating its aromatic ring about the X2 dihedral angle from - 180 to 180% Fig. 4 depicts the studied environments and the predicted chemical shifts. Our MLFF-based shift predictor shows the expected 180 periodicity (Haigh & Mallion 1979) in the chemical shifts of backbone CA atoms in the vicinity of the ring. with the ring current influence decaying smoothly as the distance from the ring increases and becoming negligible after 7A. In contrast, UCBShift extends this influence beyond the expected range and fails to reproduce the smoothness and periodicity expected in theory.
Next, investigate whether the aforementioned likelihood computation for environment reliably anticipate the resulting chemical shift prediction error.
Uncertainty estimation. We first build the unconditional likelihood model p(2') for biomolecular environments on the training set used for shift prediction. We then evaluate this likelihood for every test environment and record the corresponding chemical shift prediction error. Fig. 5 reports the distribution of errors, stratified by likelihood. Lower likelihood environments yield higher errors, making likelihood a practical confidence measure representing. the best of our knowledge, the first such score reported for chemical shift prediction.
We summarize by concluding that MLFF-based shift prediction not only delivers better accuracy, but is also physics-gae -rossidegrou and allows us to provide confidence estimates for each prediction.
7 INTERPRETING MLFF REPRESENTATIONS
While MLFF atomistic features excel at representing local protein structure and chemistry, it remains unclear how they respond to structural perturbations or encode meaningful physical properties. To address this, the behavior of MIFF representations under specific conformational changes is

---

Published as a conference paper at ICLR 2026
DMO
Glot
187 195 Sa0
PC
C407
Figure 6: Structure of the MACE embedding space for the rotating phenylalanine aromatic ring from Fig. Left: a fragment of protein with simulated rotation of the phenylalanine sidechain ring. Right: Mace embeddings of backbone CA atoms projected onto the first two principal components. The vertical axis separates the different rotation angles for visual clarity. Note that residues close to the ring (50 and 51) exhibit one-dimensional structure with nearly perfect 180-degree periodicity. The farther residue 49 shows some breakdown of this structure, while residue 40, uninfluenced by the ring current, manifests a lack of relation between ring rotation and the embedding space structure.
examined. Such changes induce smooth, interpretable trajectories in the embedding space, suggesting that the features capture physical motions and respect locality, continuity, and symmetry.
Rotated phenylalanine sidechain aromatic ring. Using the synthetically modified phenylalanine, we perform PCA on MACE embeddings obtained every X2 angle for the nearby CA atoms (residues 50 and 51 of PDB ID: 12V6). We observe that the first two principal components trace smooth, one-dimensional curve, capturing the expected 180 periodicity of the rotation. However, the effect diminishes with distance: residue 49 begins to deviate from the periodic pattern, while distant residue 40 exhibits no discernible structure, illustrating the spatial locality of the descriptor's sensitivity.
Additional case studies further analyzing the embedding space for tllxxtto-ttrad unfolding MD trajectory and experiments related to inverting MLFF embeddings, are presented in Section and Section respectively, in the Appendix.
DISCUSSION
In this work, introduced a novel way represent local protein environments by repurposing embeddings from MLFFs. To our knowledge, this is the first demonstration that MLFF latent spaces- trained exclusively on small-molecule quantum data organize according meaningful biochemical factors such as secondary structure, residue identity, protonation state, and chemical shift. We further show that these embeddings can be directly reused for diverse protein modeling tasks without retraining. This establishes MLFFs as -prppos,, reusable representation learners for structural biology and opens new paradigm for leveraging pretrained molecular models in this domain.
While our downstream models used frozen MLFF embeddings, task-specific fine-tuning of MLFFs is a promising direction that may yield further gains. Moreover, our chemical shift prediction for HA, while accurate, still slightly lags behind UCBShift, pointing to opportunities for future improvement. Notwithstanding these limitations, we believe that the proposed fully differentiable chemical shift predictor can be used guide AlphaFold and similar generative models (Maddipatla et al.) 2025) for determining structure from experimentally measured chemical shifts an important task in protein NMR that has far been hindered by the complexity of -oohetercat- shift relation.
10

---

Published as a conference paper at ICLR 2026
ETHICS STATEMENT
All authors have read and complied with the ICLR Code of Ethics. This work does not involve human subjects, personally identifiable information, sensitive data. All datasets mentioned in Appendix E are publicly available, and we adhered to best practices to ensure privacy, fairness, and transparency. We do foresee any direct ethical risks associated with the methods results presented.
10 REPRODUCIBILITY STATEMENT
The code, trained models, and data- processing scripts publicly available at httpts:/pgggtt.t thub com/
ccm/mb012/ mmmmm/1112////1mmmm 2/MFFF "tt_rrrrreeettaaee aton In addition, complete details of the models and opti-
mization parameters are provided in Appendix G.3 The hardware resources used to produce the results are specified Appendix [.4] The loss functions, evaluation metrics, and details regarding ablation studies are specified in Appendix G These details ensure that all results reported in the paper
can be independently verified.
ACKNOWLEDGMENTS
This work was supported by the Institute of Science and Technology Austria (ISTA) through the IPC grant "Generative Protein NMR" and by the Israeli Science Foundation (ISF) under grant number 1834/24. This research used resources of the Institute of Science and Technology Austria's scientific computing cluster. S.V. was supported in part by funding from the Eric and Wendy Schmidt Center the Broad Institute of MIT and Harvard.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493-500, 2024.
Gustaf Ahdritz, Nazim Bouatta, Christina Floristean, Sachin Kadyan, Qinghui Xia, William Gerecke, Timothy J O' D'Donnell, Daniel Berenberg, Ian Fisk, Niccolò Zanichelli, et al. Openfold: Retraining alphafold2 yields new insights into its learning mechanisms and capacity for generalization. Nature Methods, (8):1514--1523 2024.
Nicolas Ancona, Ananta Bastola, and Emil Alexov. Pkad-2: new entries and expansion of function- alities of the database of experimentally measured pka's 's of proteins. Journal of computational biophysics and chemistry, 22(05):515-522 2023.
Ilyes Batatia, David Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in neural information processing systems, 35:1 1423-11436, 2022.
Jörg Behler. Atom-centered symmetry functions for constructing high-dimensional neural network potentials. The Journal of chemical physics, 134(7), 2011.
Jörg Behler. Perspective: Machine learning potentials for atomistic simulations. The Journal of chemical physics, 145(17), 2016.
Jörg Behler and Michele Parrinello. Generalized neural-network representation of high-dimensional potential-energy surfaces. Physical review letters, 98( 14)::444441,1 2007.
Wolfgang Bermel, Ivano Bertini, Isabella C Felli, and Roberta Pierattelli. Nmr spectroscopy for the structural characterization natural products. Natural Product Reports, 32(4):356-36) 2015.
Amit Boyarski and Alex Bronstein. Multidimensional scaling. In Computer Vision: Reference Guide, pp. 236-849. Springer, 2021.
Stephen K Burley, Helen M Berman, Gerard J Kleywegt, John L Markley, Haruki Nakamura, and Sameer Velankar. Protein data bank (pdb): the single global macromolecular structure archive. Protein crystallography: methods and protocols, pp. 627-641, 2017.
11

---

Published as a conference paper at ICLR 2026
David A Case. Chemical shift prediction in biomolecular nmr. Current Opinion in Structural Biology, 23(2):122-176, 2013.
Lixue Cheng, Matthew Welborn, Anders $ Christensen, and Thomas F Miller. A universal density matrix functional from molecular orbital-based machine learning: Transferability across organic molecules. The Journal of chemical physics, 150(13), 2019.
Stefan Chmiela, Valentin Vassilev- Galindo, Oliver T Unke, Adil Kabylda, Huziel E. Sauceda, Alexan- dre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9())eadd0887, 2023.
Timothy D. D.W. Claridge. High- Righ-Resooution NMR Techniques in Organic Chemistry. Elsevier, 3rd edition, 2016.
Bowen Deng, Peichen Zhong, KyuJung Jun, Janosh Riebesell, Kevin Han, Christopher J Bartel, and Gerbrand Ceder. Chgnet as a pretrained universal neural network potential for charge-informed atomistic modelling. Nature Machine Intelligence,
2023.
Jacob Devlin, Ming-Was Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume (long and short papers), pp. 4171-4186, 2019.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3-11, 2018.
Rokas Elijošius, Fabian Zills, Ilyes Batatia, Sam Walton Norwood, Dávid Péter Kovács, Chris- tian Holm, and Gábor Csányi. Zero shot molecular generation via similarity kemnlss. Nature Communications, 2025.
Zsolt Fazekas, Dóra K. Menyhárd, and András Perczel. Locohd: metric for comparing local environments of proteins. Nature Communications, 15(1):0029 2024.
Dmitrij Frishman and Patrick Argos. Knowledge-based protein secondary structure assignment. Proteins: Structure, Function, and Bioinformatics, 23(4):566-579, 1995.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263-1272. PMLR, 2017.
Hatice Gokcan and Olexandr Isayev. Prediction of protein p k with representation learning. Chemical Science, 13(8):2462-2411 2022.
Hans Günther. NMR Spectroscopy: Basic Principles, Concepts, and Applications in Chemistry. Wiley, 3rd edition, 2013.
CW Haigh and RB Mallion. Ring current theories in nuclear magnetic resonance. Progress in nuclear magnetic resonance spectroscopy, 13(4):003344, 1979.
B. Han, Y. Liu, S.W. Ginzinger, and D.S. Wishart. Shiftx2: chemical shift prediction. Journal of Biomolecular NMR, 50(1):43-57, 2011. doi: 10.1007/ s10858-01 101-0011--9--44-
significantly improved protein
Thomas Hayes, Roshan Rao, Halil Akin, Nicholas J. Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Q. Tran, Jonathan Deaton, Marius Wiggert, Rohil Badkundri, Irhum Shafkat, Jun Gong, Alexander Derry. Raul S. Molina, Neil Thomas, Yousuf A. Khan, Chetan Mishra, Carolyn Kim, Liam J. Bartie, Matthew Nemeth, Patrick Hsu, Tom Sercu, Salvatore Candido, and Alexander Rives. Simulating 500 million years of evolution with language model. Science, 2025. doi: 10. 10.1126/ssiimmcc.as00011
URL http: dx. doi.
ads 0018
Jeffrey C Hoch, Kumaran Baskaran, Harrison Burr, John Chin, Hamid R Eghbalnia, Toshimichi Fujiwara, Michael Gryk, Takeshi Iwata, Chojiro Kojima, Genji Kurisu, et al. Biological magnetic resonance data bank. Nucleic acids research, (1(D)): (D)):6366-D336, 2023.
12

---

Published as a conference paper at ICLR 2026
Viktor Horaak, Robert Abel, Asim Okur, Bentley Strockbine, Adrian Roitberg, and Carlos Simmer- ling. Comparison of multiple amber force fields and development of improved protein backbone parameters. Proteins: Structure, Function, and Bioinformatics, 65(3):712-725, 2006.
Mare OJ Jäger, Eiaki Morooka, Filippo Federici Canova, Lauri Himanen, and Adam S Fos- ter. Machine learning hydrogen adsorption on nanoclusters through structural descriptors. upj Computational Materials, 4(1):37, 2018.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, (966777)):583--88 2021.
Wolfgang Kabsch and Christian Sander. Dictionary of protein secondary structure: pattern recognition of hydrogen--oonded and geometrical features. Biopolymers: Original Research on Biomolecules, 22(12):5577-2637 1983.
Shuichi Kawashima and Minoru Kanehisa. Aaindex: amino acid index database. Nucleic acids research, 28(1):374-3744 2000.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
Thomas N Kipf and Max Welling. Supi-super classification with graph convolutional networks. 2016.
arXiv preprint 1609.
Kai Kohlhoff, Paul Robustelli, Andrea Cavalli, Xavier Salvatella, and Michele Vendruscolo. Fast and accurate predictions of protein nmr chemical shifts from interatomic distances. Journal of the American Chemical Society, 13 1314439:18944-188999 2009.
Dávid Péter Kovács, Harry Moore, Nicholas J Browning, Ilyes Batatia, Joshua T Horton, Venkat Kapil, William C Witt, Ioan-Bogdan Magdău, Daniel J Cole, and Gábor Csányi. Mace-off23; Transferable machine learning force fields for organic molecules. arXiv preprint 1rX1v:2312.15211, 2023.
Dávid Péter Kovács, Harry Moore, Nicholas Browning. Ilyes Batatia, Joshua T Horton, Yixuan Pu, Venkat Kapil, William C Witt, loan-Bogdan Magdau, Daniel J Cole, et al. Mace-off: Short-range transferable machine learning force fields for organic molecules. Journal of the American Chemical Society, 147(21):17598- 17611, 2025.
Ilya Kuprov, Nicola Hagner-Rundeld and P. J. Hore. Nmr of large proteins: Overcoming the relaxation hurdle. Progress in Nuclear Magnetic Resonance Spectroscopy, 50(1):1-19) 2007.
Vivian Law, Craig Knox, Yannick Djoumbou, Tim Jewison, An Chi Guo, Yifeng Liu, Adam Ma- ciejewski, David Arndt, Michael Wilson, Vanessa Neveu, al. Drugbank 4.0: shedding new light on drug metabolism. Nucleic acids research, 42D1))D1091-D1099 2014.
Axel Levy, Eric R. Chan, Sara Fridovich- Keil, Frédéric Poitevin, Ellen D. Zhong, and Gordon Wetzstein, Solving inverse problems in protein space using diffusion- -bsee priors, 2025. URL https: httbss://arxiv..oo http:/xwww.orgaab/////oo/a 04239
Da-Wei Li and Rafael Brüschweiler. PPM: a side-chain and backbone chemical shift predictor for the assessment of protein conformational ensembles. J. Biomol. NMR, 54(3):257-265, nov 2012.
Dawei Li and Rafael Brüschweiler. Ppm_one: a static protein structure based chemical shift predictor. Biomol. NMR, 6203))403-409, 2015.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, al. Evolutionary--saae prediction of atomic-level protein structure with a language model. Science, 37966637) 2023.
Advaith Maddipatla, Nadav Bojan Sellam, Meital Bojan, Sanketh Vedula, Paul Schanda, Ailie Marx, and Alex M. Bronstein. Inverse problems with experiment-guided alphafold. In Forry-second International Conference Machine Learning, 2025.
13

---

Published as a conference paper at ICLR 2026
Sai Advaith Maddipatla, Nadav Bojan Sellam, Sanketh Vedula, Ailie Marx, and Alex Bronstein. Generative modeling of protein ensembles guided by crystallographic electron densities. In NeurlPS 2024 Workshop Machine Learning for Structural Biology, 2024.
Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint
2018.
Jens Meiler. Proshift: protein chemical shift prediction using artificial neural networks. Journal of Biomolecular NMR, 26:25-37, 2003.
Seongho Moon and David A Case. A new model for chemical shifts of amide hydrogens in proteins. Journal of biomolecular NMR, 38:139-150, 2007.
Stephen Neal, Amy M Nip, Hongwei Zhang, and David S Wishart. Shiftx: chemical shift prediction from protein structures. Journal of Biomolecular NMR, 26(3):215- 240, 2003.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 1147, 2019.
Mats HM Olsson, Chresten R Søndergaard, Michal Rostkowski, and Jan H Jensen. Propka3: consistent treatment internal and surface residues in empirical p a predictions. Journal of chemical theory and computation, 7(2):525-537, 2011.
Jeffrey Ouyang-Zhang. Chengyue Gong, Yue Zhao, Philipp Krachenbuchl, Adam Klivans, and Daniel Jesus Diaz. Distilling structural representations into protein sequence models. In The Thirteenth International Conference Learning Representations, 2025.
Aleksandra Laura Ptaszek, Jie Li, Robert Konrat, Gerald Platzer, and Teresa Head-Gordon. Ucbshift 2.0: Bridging the gap from backbone to side chain protein chemical shift prediction for protein structures. Journal of the American Chemical Society, 146(46):1234419255 2024. doi: 10.1021/ jacs,4c10477
Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R Manby, and Thomas F Miller. Orbnet: Deep learning for quantum chemistry using symmetry-adapleer atomic-orbital features. The Journal of chemical physics, 153(12), 2020.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):177, 2014.
Pedro BPS Reis, Diogo Vicosa, Walter Rocchia, and Miguel Machuqueiro. Pypka: A flexible python module for poisson-bollzman--baeee and modeling, 60(10):4442 4448, 2020.
k calculations. Journal of chemical information
A. A. Rosenberg, Marx, and A. M. Bronstein. Codon-specific ramachandran plots show amino acid backbone conformation depends identity of the translated codon. Nature Communications, 13(1): 23215, 2022. doi: 10. 00.1038/441447-022--0000 641467-022- 30390-
URL https:: http://oo.......1000///1
Aviv A. Rosenberg, Sanketh Vedula, Alex M. Bronstein, and Ailie Marx. Sceing double: Molecular dynamics simulations reveal the stability of certain alternate protein conformations in crystal structures. bioRxiv, 2024. doi: 10. 08 brc/conttet/earllyy/ 00//224/////////// ////////2222222 08 31, 610605
URL http::::////// www.xiibiirxivvvooo
Lars Ruddigkeit, Rund Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of chemical information and modeling, 52(1 ():21)))-27777 2012.
Kristof Schütr, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus- -louerRRbert Müller. Schnet: continuous-file convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377-9388. PMLR, 2021.
14

---

Published as a conference paper at ICLR 2026
Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24), 2018.
Yang Shen and Ad Bax. Sparta+: a modest improvement empirical nmr chemical shift prediction by means of an artificial neural network. Journal of Biomolecular NMR, 40(4):353-362, 2008.
Tomoya Shiota, Kenji Ishihara, and Wataru Mizukami. Universal neural network potentials as descriptors: towards scalable chemical property prediction using quantum and classical computers. Digital Discovery, 171--778, 2024.
Justin S Smith, Ben Nebgen, Nicholas Lubbers, Olexandr Isayev, and Adrian E. Roitherg. Less is more: Sampling chemical space with active learning. The Journal of chemical physics, 148(24), 2018.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting The journal of machine learning research, 15(1):1929-1955) 2014.
Corin C Wagen, Elias L Mann, Jonathon E Vandezande, Arien M Wagen, and Spencer C Schneider. Egret-1: Pretrained neural network potentials for efficient and accurate bioorganic simulation. arXiv preprint arXi::2504.2995. 2025.
David S Wishart. Advances in nmr-based metabolomics: applications in drug discovery and develop- ment. Current Opinion in Drug Discovery & Development, 14(1):25 33, 2011.
Haiyan Zhang, Stephen Neal, and David S Wishart. Refdb: database of uniformly referenced protein chemical shifts. Journal of biomolecular NMR, 15:173-,95, 2003.
Y. Zhou, Y. Wang, H. Li, Y. Zhang, and D. Xu. Ucbshift: Protein chemical shift prediction using transfer learning and machine learning. htttps///ggthubb.ccc//////cabaap Accessed: 2025-05-13.
2023.
Roman Zubatyuk, Justin S Smith, Jerzy Leszczynski, and Olexandr Isayev. Accurate and transferable multitask prediction of chemical properties with an atoms- moloeecuees neural network. Science advances, 5(aa694a, 2019.
15

---

Published as a conference paper at ICLR 2026
APPENDIX
Summary of experiments
17
Additional Figures
19
Quantitative Results & Tables
25
Extended Background
30
Data curation
Environment construction and MLFF embedding
33
Radius ablation
34
F2
MLFF Layer Comparison
34
Predicting amino acid, secondary structure, and μ
34
Secondary Structure and Amino Acid Prediction
................................
Acid dissociation constant (pK,) prediction
35
Model Architecture
36
H
Using likelihoods to define environment similarity measure
37
Computation details
37
H.2
Similarity metric
37
Chemical shift prediction
37
Extended physics-grounded chemical shift results
................................
37
1.2
Details of the chemical shift prediction experiments
37
Model architecture
38
Computation details
38
Extended MLFF interpretations
38
U.1 Helix uniolding
38
0.2 Inverting MACE embeddings
39
Use of Large Language Models (LLMs)
40
16

---

Published as a conference paper at ICLR 2026
A SUMMARY OF EXPERIMENTS
Table B.2: Experimental overview. Summary of key research questions, experimental design, baselines, main takeaways, and supporting evidence. "ESM family" includes embeddings from ESM2/ESMFold (Lin et al. 2023), ISM (Ouyang-Zhang et al. 2025), , and ESM3 (Hayes et al. 2025 On the other hand, "learned" refers the leamed-embedding baseline described Section
Question
Experiment Design
Baselines
Takeaway
Evidence
Do MLLF representa- tions capture structural and chemical proper- ties?
Train classifiers pre- dict amino acid iden- tity (20 types) and sec- ondary structure (a-helix,
MACE,
Egret best for amino acid classification;
Fig.
Egret,
OrbNet, AIMNet, ESM family
MACE/Egret best for secondary struc- ture. UMAP shows clear clustering by chemical identity and structural motifs.
3-strand, other) from MLFF embeddings. Vi- sualize embedding space using UMAP.
2. Can MIFF representations predict pK, values?
Train regression models predict pK, for ion- izable residues (GLU, ASP, LYS, HIS) using MLFF embeddings. Target pk, values are ap- proximated by PypKa's Bolszzmnn solver.
PropKa, PKa-ANI, ESM fam- ily, learned, and MLFF variants
AIMNet features achieve best perfor- mance.
Table
Can MLFF embeddings de- tect distribution shifts and sep- arate structural classes?
(a) Measure likelihood of environments be- fore/after Amber99 relaxation using kernel density estimation. (b) Construct conditional likelihoods for different secondary structures.
Compare distri- butions before/after relaxation
(a) Relaxed struc-
Figs. A.7
tures receive higher likelihoods. (b)
A.ll
Conditional likeli-
hoods clearly sepa- rate secondary struc- ture classes
4. Can MLFF- based predictors outperform state- of-the-art chem- ical shift predic- tion?
Train GNNs MLFF UCBShift2- embeddings to predict NMR chemical shifts for backbone (N, CA, C, H, HA) and side-chain atoms.
MLFF-based pre- dictors outperform UCBShift2-X for backbone and side-chain heavy atoms (except
Fig. Table B.5
X (SOTA), ESM fam- ily, learned
Table B.6
HA). MACE-based models perform best overall.
Are MLFF- based chemical shift predictions physically consis- tent?
Three case studies: (1)
UCBShift2-
MLFF predictor shows expected π r-periodicity, cor- rect Co/C3 shifts during unfolding, and distinguishes conformational states. UCBShifi2- X shows unphysical behavior.
Figs.
Systematic phenylala- nine ring rotation (0-2m)
X, exper-
A.8 A.12
imental
probe ring-current ef- fects. (2) MD simulation of tali--to-ttrand tran- sition. (3) Alternative conformations in protein
chemical
shift trends
40LE:B.
17

---

Published as a conference paper at ICLR 2026
Table B.2 (continued)
Question
Experiment Design
Baselines
Takeaway
Evidence
6. Can likelihood serve as an uncer- tainty estimate?
Build unconditional like- lihood model on train- ing embeddings; corre- late environment likeli- hood with chemical shift prediction error on test
No direct baseline (novel
Low-likelihood environments con- sistently exhibit higher prediction errors, supporting use for uncertainty estimation.
Fig.
metric)
How do MLFF embed- dings respond to structural perturbations?
Perform PCA on MACE embeddings during: (a) systematic phenylalanine ring rotation, (b) helix- to-strand MD simulation. Analyze embedding tra- jectories and spatial local- ity.
Visual anal- ysis of em- bedding tra- jectories
Embeddings trace smooth trajectories reflecting structural changes, show expected period- icity, and reveal interpretable latent directions.
Figs.
A.9 A.10
Can MIFF embeddings be inverted to re- cover structure?
Optimize structures using only MACE matching as loss function. Test on alter- nate conformations in:
AlphaFold3
Direct com- parison to target con- formations
Backbone geometry recovered with high accuracy; side- chain orientations less precise. MLFF embeddings encode most but not all information needed for reconstruction.
Fig. A.14
A.15
descriptor
A.16
7EC8.
9. Compara- tive benchmark- ing and ablation studies
(1) Benchmark all MLFF families on all tasks. (2)
MACE, Egret,
MACE performs best on most tasks (AIMNet best for pk,)) Full atom set improves per- formance; MLFF similarity outper- forms LoCoHD.
Tables B.3 B.11
Ablate atom set selection.
OrbNet,
Fig. A.13
(3) Compare MLFF simi-
AIMNet
larity metric to LoCoHD.
variants;
LoCoHD
metric
18

---

Published as a conference paper at ICLR 2026
B ADDITIONAL FIGURES
Relaxed
Unrelaxed
-60
-50
-40
-30
-20
-10
0
10
Likelihood
Figure A.7: MACE-basee likelihood captures structurally subtle distribution shifts in protein environments. The histogram of likelihoods of CA MACE representations of 26000 protein environ- ments curated from 225 protein structures in the test set, before (red) and after (blue) relaxing the protein structures with an Amber99 force-field (Hornak et 2006).
74
CA MACE
CA UCB Snit
CB UCB Shilt
Chemical shifts [opm]
70
66
64
62
60
a-helix
B-strand
Figure A.8: Simulated molecular dynamics trajectory a helix unfolding into a strand. Depicted the chemical shifts of the CA and CB atoms in the middle of the helix predicted by Mace and chemical shift distributions the
UCBShift2-X. X. Marginal plots show the anperimentally-determmmi atoms in helices (left) and strands (right). The known tendency of CA's have smaller shifts in strands than in helices and the opposite CB tendency is clearly visible.
19

---

Published as a conference paper at ICLR 2026
CA6
CA7
PC
PC
Figure A.9: Structure of the MACE embedding space for the unfolding helix from Fig. A.8 Depicted are Mace embeddings of the backbone CA atoms of residues and projected onto the first two principal components. Observe essentially one-dimensional structure of the trajectory in the PCA space.
Pairvise deser distar ce
Pairwise RMSD (A)
Pearsore PeprsooeSpeammnn 0.92/0.96
Phiraite disecripttrr destonos
0mc
time
08
Painwise RMSD (A)
20
30
Figure A.10: Pairwise descriptor distance is correlated with structural deviations. (Left) Pairwise descriptor distances matrix for all frames along the unfolding MD trajectory. (Middle) Pairwise Co RMSD (A) for the same set of structures. (Right) Relationship between descriptor distance and Co RMSD with Pearson/Spearman correlation coefficients of p = 0.92/0.98. The descriptor distance is the Euclidean distance between their vectorized Ca feature representations.
20

---

Published as a conference paper at ICLR 2026
6
4
2
0
-2
2
4
log p(-IT)
2
log (og
-2
T
-6
-4
-2
0
2
4
log ((-(H)
Figure A.11: Estimated likelihoods of MACE embeddings of backhone CA atoms in differ- ent secondary structures. The two-dimensional plots show projections of the log-likelihoods (log pxx|EE, log p(x;((()x)))))) log pxxpp((TT))))) of 8163 MACE embeddings in strands (E), helices (H), and turns (T). The one mee-imensioonl plots depict the KDE-estimated marginal distributions of each the lieeliiils...ds
21

---

Published as a conference paper at ICLR 2026
CA63
B
[A]
LYS63
CA secondary shift [ppm]
CB63 CB63 (B)
-5
-3
-2
-1
secondary [ppm)
Figure A.12: Distribution of the secondary chemical shifts of the two stable conformers of 40LE. Left: 100 ns MD) simulations (Rosenberg et al. (024) of the two conformers (marked as (containing a helix) and (containing a linearly structured loop) according to the original PDB annotation). Lysine 63 is highlighted. Right: secondary shift distributions of lysine 63 CA (top) and CB (bottom) atoms over the MD trajectory. Note that the helical conformer exhibits an expected excess in the CA secondary shift and defect in the CB secondary shift.
22

---

Published as a conference paper at ICLR 2026
MACE
LoooHD
07
OR
OR
cc
%5r
OW
ON
Figure A.13: Similarity of local protein environments using MACE and LoCoHD. In each plot, environments belonging to different class labels are represented as points while pairwise Euclidean distances approximate the dissimilarities measured using the LoCoHD metric (Fazekas et al. 2024) and metric between MACE likelihoods. Circles indicate each class variability as seen by the metric. Shown conditioning by secondary structure (top row) and amino acid identity (bottom row). We conclude that MACE captures much better the similarities between related chemical and structural classes, even if LoCoHD has been explicitly designed for these tasks.
23

---

Published as a conference paper at ICLR 2026
AppaFarri
0.75
2008
0.25
ANTU
0.02
5051
-0.50
7608
-1.0
3051
Figure 1414: AlphaFold3 guided structures recovered by inverting MACE descriptors. (A) The unguided AlphaFold3 structures (orange) of four proteins (3V3s, 4NPU, 5G51, and 7EC8) with experimentally determined alternate backbone conformations are shown alongside guided predictions (green), obtained by inverting the corresponding conformation's MACE descriptor. All predictions are overlaid on the experimentally determined conformer (white) from the PDB (Burley al. 2017). (B) The violin plots the signed normalized distance distributions relative to conformations A (-1) and (1) in generated structures for the same cases, comparing unguided (orange) and descriptor-guided predictions (green / blue).
A-B
B-A
3V3S
4NPU
5G51
7EC8
Figure A.15: Inverting MACE descriptors by optimizing atomic coordinates. Shown are the outcomes of optimizing atomic coordinates to match the local environment descriptors of alternate conformation (altloc). For each protein, two optimization directions are displayed: starting from con- formation A and matching the descriptors of conformation B (A-+B), and starting from conformation B and matching the descriptors of conformation A (B (B+A)). The proteins and altloc pairs used in this analysis are the same as those used for guidance in Fig, A.14 The reference conformation is shown in white, and the optimized conformation in green.
24

---

Published as a conference paper at ICLR 2026
ALA 59
Target
AF3 prediction
Optimization result
VAL-68
4OLE
Figure A.16: Inverting MACE descriptors by optimizing atomic coordinates Shown in green is the structure of the protein 40LE:B predicted by AlphaFold3. The structure has alternative conformation in the region 60 68, (shown in red is the crystallographic structure from the PDB), whose MACE environment descriptors were used to optimize the AlphaFold prediction. The solution of the optimization problem is depicted in pink.
C QUANTITATIVE RESULTS TABLES
Table B.3: Secondary structure prediction accuracy. Reported are precision, recall, and F1 scores, each calculated with respect the ground truth DSSP secondary structure classes, The best-performing model for each metric is indicated in bold.
ESM-Features + GCN
MLFF-Features GCN (Ours)
Metric
Class
LocoFID
ESM3 (Seg)
ISM
ESM3 (Struc)
ESMFold
Egret
MACE
OrbNet
AIMNet
Weighted Mean
76.917
88.175
92,
94.881
93.533
694
92.300
91.716
Mean
76.424
87.740
02.279
94.981
93.304
95. 55.5540
96.528
92.233
04.522
Precision
R-strand
75.623
91.124
05.555
05.252
95.757
96. 248
95.960
91.749
95.494
Other
73.158
81.753
88.216
92.419
88.941
93.905
933 958
89.049
92.300
Weighted
92.574
Mean
75-532
87.637
92.071
95.709
92.961
98. 508
95.155
91.977
94.335
Recall [%6]
cr-helix
87 227
91.259
94. 562
96.844
96 0223
97. 5-46)
17. 5222
95.548
96. 806
3-strand
Other
50.333
78.753
85.833
89.759
87.521
90.512
597
$4.106
88.720
Weighted Meun
76.
$8.210
92.479
95.433
93.311
98. 633
96.699
92.509
94.720
Mean
87.675
90.3889
5.52
96.4227
92.073
94.412
FL [%]]
re-belix
83. 725
90.798
93 507
96.1733
91.726
97. (103
97. 087
91.717
96. 281
F-strand
81.783
91.983
95.687
96.016
96.047
97.
96.948
94.996
96.477
59.636
87.008
91.009
88.225
92.177
247
86.107
90.479
25

---

Published as a conference paper at ICLR 2026
Table B.4: Amino acid prediction accuracy. Reported are precision, recall, and Fl scores, each calculated with respect to the ground truth amino acid class. The best pest-perforning model for each metric is indicated in bold. Unlike secondary structure prediction, per-class performance is not reported here due to the large number (20) of amino acid types.
ESM-Features GCN
MLFF-Features GCN (Ours)
Metric
Class
LocofID
ESM3 (Seq)
ISM
ESMB (Struc)
ESMFold
Egret
MACE
OrhNet
AIMNet
Precision [%]
Weighted Mcan
07. 734
97:919
96.854
08.884
98-448
95.967
98. 865
Meau
98. 08.012
89.923
96,891
96.828
97.821
98.758
96.031
95,050
98.107
Recall 196]
Weighted Mean
07. 905
91.089
96.920
96.964
98.642
99. 133
98.431
95,955
98, 664
97. 689
95.022
98
FI (88]
Weighted Mean
97. 785
01.079
96.914
08. 703
99.
18 48-433
95.859
98 663
Moan
97.539
89.686
96,101
96.815
97.822
98.8442
96.157
95.078
98. 181
Table B.5: Backbone chemical shift prediction accuracy. Reported is the mean absolute error (MAE), in ppm, relative to the experimentally measured chemical shift values from RefDB (Zhang
al. 2003 for different atoms. The best-performing model for each metric is indicated in bold.
FSM features GCN
MI.FF-frature GCN
Atoen
LocoHD
Learned
ESM2
(5eq)
(Struc)
Egret
AIMNet
HA
0.165
NIA
0.405
0.247
0.235
0.220
0.193
0.203
0.150
0.176
0.200
0.201
H
0.321
NIA
0.528
0.402
0.383
0.317
0.000
0.365
0.300
0.295
0.333
0.325
1.753
N 1.891
2.877
3.174
2.455
2.349
2.275
1 1.845
1.586
1.642
1.667
2,015
1,875
CB
0.758
1.140
1.326
0.923
1,082
0.894
0.758
0.775
0.716
0.705
0.792
0.780
Table B.6: Side chain chemical shift prediction accuracy. Reported is the mean absolute error (MAE), in ppm, relative the experimentally measured chemical shift values from RefDB (Zhang et al. 2003) for different atoms. The best-performing model for each metric is indicated in bold.
MLFF-features + GCN (Ours)
Atom
UCB Shift
LocoHD
МACE
Egret
Orb
AIMNet
CG
0.595
0.647
0.567
0.563
0.638
0.611
CD
0.615
0.512
0.456
0.448
0.521
0.482
CD2
1.116
1.195
1.035
1.085
1.201
1.749
CG2
0.773
0.833
0.713
0.720
0.760
0.742
CE
0.519
2.295
0.447
0.436
0.451
0.446
Table B.7: PropKa vs. MACE pK, comparison. Per-residue pK, prediction errors for Glutamic acid and Aspartic acid from PropKa's training set. PropKa is evaluated in its native setting, using experimental pK, values derived from crystallographie PDB structures, whereas our MACE-bbase model was trained using simulated pK, values from AlphaFold2 structures. This comparison in- herently favors PropKa, yet our model achieves lower error on more than half of the structures in PropKa's training set.
Glutamic acid
Aspartic acid
PDB ID
PropKa
Ours (MACE)
PropKa
Ours (MACE)
1PGA:A
0.413
0.294
0.362
0.292
LIGD:A
0.489
0.182
0.352
0.274
1A2P:A
0.525
1.046
0.426
0.403
4ICB:A
0.579
0.546
0.299
1.018
1BE0:A
0.730
0.511
ALZT:A
0.629
0.665
0.488
0.404
3RN3:A :A
0.548
0.462
0.423
0.412
2RN2:A 1A
0.274
0.185
0,683
0.585
20V0:A
0.377
0.376
0.180
1.571
1XNB:A
1.001
1.270
0.436
1.020
135L:A
0.470
1.700
0.480
1.043
26

---

Published as a conference paper at ICLR 2026
Table B.8: -Na-ANI vs. MACE pk, performance. Mean absolute pK, prediction enror (MAE) for four amino acids using -K,-ANI and our MACE-based model, evaluated on the PKAD-R dataset (Ancona al. 2023), curated version of the dataset on which - was originally trained and tested. While pK, ANI is evaluated in its native setting, model trained simulated pK, values derived from Alphafold2 structures, making this comparison inherently favorable to -NNI Nevertheless, our model outperforms ANa on three of the four amino acids; for Histidine, we hypothesize that the poor performance due to an insufficient number of training examples for that residue type.
Amino acid
pKa-ANI
Ours (MACE)
Glutamic Acid
0.696
0.695
Aspartic Acid
0.978
0.831
Lysine Histidine
0.648
0.360
0.899
0.925
Table B.9; Environment radius ablation. Prediction error (ppm; lower is better) for backbone chemical shift prediction function of the local environment radius We compare 4, 5, A using MACE embeddings as input to the GCN. The best-performing radius (5 A) is highlighted in bold.
Atom
4Â
5A
GÅ
HA
0.194
0.180
0.191
0.320
0.300
0.318
CA
0.620
0.584
0.627
1.779
1.642
1.768
CB
0.758
0.716
0.754
0.798
0.744
0.804
Table 10: MLLF layer ablation. Per-atom mean absolute error (MAE) backbone chemical-shift prediction for different choices of descriptor layer. For MACE, LI and L2 denote descriptors taken from the first and second message-passinng layers, and L1+2 denotes their concatenation. For Orb, L5, L10, and L15 denote descriptors taken from the 5th, 10th, and 15th message- layers, respectively. For each atom type and model, the best- pess-performin layer configuration is highlighted in bold.
Atom
MACE L1
MACE L2
MACE L1+2
Orb L5
Orb L10
Orb L15
C
0.797
0.755
0.744
0.836
0.836
0.847
CA
0,619
0.593
0.584
0,678
0.660
0,683
CB
0.756
0.705
0.716
0.778
0.777
0.792
H
0.318
0.300
0.300
0.330
0.331
0.333
HA 0.191
0.178
0.180
0.198
0.196
0.200
N
1.771
1.666
1.642
1.929
1.889
2.015
Table B.11: Machine Learning Force Field (MLFF) metadata. Listed are the model type, dimen- sionality, geometric properties of the representations, and training dataset for MLFFs in Section 2
Model Name
Dimension
Properties
Dataset
MACE large
S, invariant
QM9 (Ramakriskuan
2014).
ANI-1s
OrhNet vž
256
S. & E(3) invariant
QM9, GDB
Cheng at al 2019
DrugBark (Law et al. 2014
AIMNet2
S,
AN Smith
2018
Egretl
384
S, 0(3) invariant
QM9, MD22
Chmicla et
2025
ANI-IX. QM24
Ruddigkeit et al. 2012
27

---

Published as a conference paper at ICLR 2026
Table B.12: Protein Secondary Structure Codes. Summary of the standard single-letter codes used represent protein secondary structure elements, defined by the DSSP classification (Frishman & Argos 1995).
Symbol
Name
H
a-helix
Turn
G
310-helix
I
æ-helix
E
3-strand (extended)
Isolated B-hridge
S
Bend
Coil
Table B.13: Secondary structure prediction ablation. Evaluation of secondary structure prediction accuracy when using a single Egret atom descriptor compared to using multiple Egret atom deseriptors. The best-performing model for each metric is indicated bold.
Metric
Class
Aa = (CA)
Ao {C, CA, CB, N, H, HA}
Weighted mean
94.822
95.685
Mean
94.679
95.554
Precision [%]
a-helix
95.809
96.466
8-strand
95.233
96. 248
Other
92.997
93.905
Weighted mean
94.853
98.636
Mean
94.552
98.508
Recall [%]
a-helix
96.759
97.546
3-strand
97.971
99.467
Other
88.925
90.512
Weighted mean
94.817
98. 08.633
Mean
94.593
98.558 08.552
F1 [%]
a-helix
96.282
97.003
3-strand
96.582
97.103
Other
90.915
92.177
Table B.14: Amino acid prediction ablation. Evaluation of amino acid prediction accuracy when using a single Egret atom descriptor compared using multiple Egret atom descriptors. The best- performing model for each metric is indicated in bold.
Metric
Class
A, = {CA}
A, = {C, CA, CB, N, H. HA}
Weighted mean Mean
95.897
99.129
Precision [%]
94.327
98.758
Weighted mean
95.912
99.133
Recall [%]
Mean
94.069
98.932
Weighted mean
95.629
99.155 99.135
F1 [%]
Mean
93.765
98.842
28

---

Published as a conference paper at ICLR 2026
Table B.15: Set of protein structures with multi-modal backbone conformations (altlocs) from Rosenberg et al. (2024) guided using MLFF descriptors.
PDB ID
Residue range
Sequence
Resolution [Á]
3
245 250
KAQERD
1.90
7EC8:A A
187 190
DGGI
1.35
5G51:A
290 295
GSASDQ
1.45
4NPU:B
133 136
FEEI
1.50
29

---

Published as a conference paper at ICLR 2026
D EXTENDED BACKGROUND
Chemical shifts. The resonance frequency of nuclear spin, v. depends directly on the splitting of its energy levels a magnetic field. This is given as v x yB, where B is the strength of the magnetic field that the nucleus experiences and y is the gyromagnetic ratio of the nucleus, a characteristic of an atom's 's nucleus (e.g.. 'H. I).. In Nuclear Magnetic Resonance (NMR) experiments, is primarily determined by the externally applied magnetic field Bo. but at the level of an observed atom it is slightly perturbed by local magnetic fields AB induced by the surrounding electron cloud. These perturbations, which are typically several orders of magnitude weaker than Bo, cause small shifts in the resonance frequency known as chemical shifts. The chemical shift reflects subtle variations in the local electronic environment and is therefore highly informative about molecular structure and composition. The chemical change of an atom is often reported in relative terms as 8 = 100 "-Pe where v is the observed resonance frequency and Vo is the resonance frequency of the same type of nucleus in a reference compound commonly the 'H signal of 1e-dimmethyl---illeeettaa sulfonate (DSS) is used for the "H chemical shift scale. Although dimensionless, chemical shifts o are conventionally expressed in parts per million (ppm), hence the 10 scaling factor in the definition.
The electronic environment that gives rise to chemical shifts is influenced by factors such as the nature of chemical bonds, the identities of neighboring atoms, and their ptrre-dimensional spatial arrangement. Additionally, dihedral angles] play significant role in modulating electron density around nuclei and thereby affect chemical shifts. Other factors that influence the chemical shift include hydrogen bonding interactions, ligand binding, or proximity to solvent molecules. Therefore, chemical shifts provide site-specific information about the local environment at the atomic-level. The chemical shift reports the time-averagee environment over all motions on time scales shorter than milliseconds. Chemical shifts are particularly valuable for inferring atomic connectivity, identifying functional groups, and detecting conformational changes, due their sensitivity to fine variations in electronic structure (Claridge 2016;; Günther 2013). Accurate chemical shift prediction plays a critical role in applications like automated NMR resonance assignment, molecular structure determi- nation and validation, and the analysis of complex chemical systems (Wishart 2011 Bermel et al. 2015). Despite this. reliable chemical shift prediction remains a challenge due the sensitivity of chemical shifts to nuanced changes in molecular environment (Kuprov et al. 2007 Case 2013).
Acid dissociation constants (pKg) pK, a thermodynamic quantity describing the tendency of titratable group to donate proton. When the solvent pH is lower than the pKa, the group will be predominantly protonated; conversely, becomes increasingly deprotonated at pH values above the pK,, Each amino acid has canonical pK, in isolation (e.g., Aspartate ~3.9), but within proteins these values can shift significantly. Nearby charged or polar residues, hydrogen-bone partners, and desolvation effects (especially in buried regions) modulate proton affinity and can shift pKa by several pH units. Accurate prediction of these shifts is crucial for understanding enzymatic catalysis, pH-dependent conformational changes, and protein stability.
Molecular Dynamics and Density Functional Theory. Molecular Dynamics (MD) simulations numerically integrate Newton's 's equations of motion to evolve the atomic positions of system over time, requiring potential energies and forces at every timestep. When computed from quantum mechanics, these quantities can be obtained using Density Functional Theory (DFT), first-principles electronic-structure method that solves the Kohn-Sham equations to deliver accurate total energies, forces, and electronic properties. However, the steep computational cost of DFT limits its direct application to small systems and short timescales. Machine learning force fields (MLFFs) address this limitation by providing near-DFT quality energies and forces at onders-of-magniuue cost, enabling quantum quantum-accurate accurate MD simulations of systems with tens of thousands of atoms over nanosecond timescales.
lower
ESM-based descriptors (Lin et al., 2023: Ouyang-Zhang et al., 2025 Hayes et al., 2025).
Evolutionary Scale Modeling (ESM) is a family of transformer-basee protein language models trained on a large-scale dataset of protein amino acid sequences using masked language modeling, similar to BERT (Devlin al. 2019) for natural language. Unlike traditional methods that rely on multiple sequence alignments (MSAs) or co-evolutionary profiles, ESM models learn context- aware representations directly from amino acid sequences, enabling them capture biochemical,
'The torsion angle about the hond hetween atoms and C in four-atom segment A-B-C-D
30

---

Published as a conference paper at ICLR 2026
structural, and evolutionary patterns within a protein sequence without explicit alignment. These representations have been used in two main ways. First, they underpin single-sequence 3D structure prediction in the ESMFold family of models (Lin et al. 2023). Second, later work augments these sequence-only models with structural supervision to obtain more structure-awant representations, for example in multimodal sequence-structus models such as ESM3 (Hayes al. 2025) and in sequence-
based protein language models that explicitly distill structural information into their embeddings (Ouyang- -Zhang et al. 2025).
Given an input amino acid sequence a of length L, a pre-trained ESM model produces a d-dimensional representation for each amino acid. These embeddings capture the identity of each amino acid and its contextual dependencies within the sequence. In this work, we use per-residue representations from the ESM-2 model (output from the 33rd layer of the 2B parameter model) as local environment descriptors, leveraging the model's 's ability to capture the evolutionary and biochemical context for each residue.
MACE and Egret descriptors (Batatia et al., 2022; Wagen al. 2025). Message passing neural networks (Gilmer et al. 2017 Schüitt et al. 2017) typically exchange messages between pairs of nodes in graph. In contrast, the MACE famil of models (MACE, Egret) generalizes this framework by learning equivariant messages involving a system of n (order parameter) nodes (or atoms) at a time. This enables richer geometric and chemical modeling by operating directly over higher-order interactions. Formally, for a given node at layer i of the graph, MACE first learns the following S,, Sn,iinvariant and (33)-equivariant features between pairs of atoms:
Σ
Clams limiddmma
Σ JEN()
R( Riiii (T (f.)) Σ wrO (O) kklz J.klgma
i.klgma
Here, R is a learnable radial function that encodes the distance between atoms and j: Y is the spher- ical harmonic function that encodes the unit vector from atom i to atom j. The summation over the neighbors (i) guarantees permutation invariance in the local aggregation. Additionally, A, i,klams in- corporates information from learned node embeddings at the previous layer using Σx w:tt) h (1)
kkil 3,kggmn
Clama Lastly, Et.ms ,aam Lims (Clebsch- Garddn coefficients) ensure equivvariree is preserved when radial and spherical features are combined. The indices 4 12, L3, m1, m2, my generalize equiv- ariant messages across eigher-dimensoons irreducible representations of 0(3). It must be noted that
A (t) i.klmm captures pairwise directional and distance-barc based interactions, forming the building blocks for higher-order representations. Formally,
B i.M
CLM Σ Im
IIΣ 8=1
A "RRE
Im
(3)
i.klems'
The above equation lifts the pairwise S,, (permutation-) invariant and -eeu-qquivariant features in A i,klgmg to an n-body representation while preserving the geometric properties. The high-order tensor (t) &LM forms the core MACE representation. The update and readout steps are similar to regular message passing frameworks. As special case, we set L. = 0 (scalar representation of O(3)) in Equation 3 we project the equivariant features to a invariant representation of ((3) via scalar coupling, Hence,
B k,00
0.00 mm
IE E=1
(e) A Ekig
.l ....1....4.
i,klgme"
Here 0000 Im would enforce scalar coupling of lower-order equivariant components. In this work, will use descriptors (448 dimensional) that are invariant to both S, and 0(3) transformations by setting L 0.
AlMNet2-based descriptors (Zubatyuk et al. 2019.. AIMNet constructs Atomic Feature Vectors (AFVs) by iteratively embedding the local atomic environments and updating atomic representations using geometric information. Given the atomic coordinates R and atomic number Z, initial atomic feature vector A, (0) is derived for atom i. The initial embeddings are parametrized as learnable vectors that depend atomic number 7, of each atom.
31

---

Published as a conference paper at ICLR 2026
The local environment of atom i is encoded using symmetry functions that capture the spatial arrangement and types of neighboring atoms. This produces an atomic environment vector (AEV).
G (E, R,, 23)EEN())-
An AEV constructed to be invariant to permutations (S,,) and 3D translations and rotations (E(3))
This is typically via wessaag-passiin framework. Next, the geometry and atom-type information are fused by computing an interaction descriptor
MLP
Σ JEN,
G R A
This vector encodes the influence of neighboring atoms on atom in way that is independent of the number of chemical species. Lastly, the atomistic feature vectors A and f, are combined to form an updated atomistic feature vector
A = U P()))
This process is repeated over multiple iterations, enabling information to propagate beyond the local neighborhood and capture non-local effects such charge redistribution and polarization. After several refinement steps, AIMNet produces AFVs of fixed dimensionality (256), which are used predict various molecular and atomic properties. Training is performed in multitask setting, where the network jointly learns to predict variety of molecular and atomic properties including energies, forces, partial charges, and dipole moments using shared representations to improve generalization. The model is trained large datasets of quantum mechanical calculations, with regularization techniques to ensure stability and prevent overfitting.
OrbNet-based descriptors (Qiao et al. 2020). In similar spirit to semi-empirical molecular force- fields, OrbNet constructs feature vectors by leveraging cunmetry-adapter atomic orbitals (AO), which capture local electronic structure information obtained from quantum calculations while respecting the symmetries inherent in molecular systems. Using mean-field quantum mechanics, OrbNet derives overlap matrices S, Fock matrix F, density matrix D. and the Hamiltonian H. Each matrix element O" corresponds to interactions between atomic orbitals v and μl, and blocks Oi represent interactions between atoms and j.
These matrices produce features that are S, invariant and E(3) invariant by projecting these orbital block matrices into symmetry-adapted basis. The atom-wise feature vector is then constructed by aggregating contributions from its neighboring orbital interactions
Σ MLP(O..)-
In OrbNet, the atomic structure is represented as a graph where nodes correspond to an atom (initialized using f,) and edges model interatomic relationships. The model employs the following message passing scheme
(h (O..... 0.)))
h (t+1) = =(0
(1)
Σ JEN(()
(9)
where h," is hidden feature of atom i at layer f, Oij encodes the orbital-based edge features derived from orbital feature matrices, M()) is the message function, and U is the update function for node embeddings. After multiple layers of message passing, the final atomic features f, E 256 serves as a descriptor of atom's 's local environment. This formulation enables OrbNet to encode both local chemical environments and long-range quantum effects efficiently. In contrast to models like MACE, which explicitly encode geometric equivariance and higher-order correlations, OrbNet uses pairwise orbital- -aased features and rely on message passing to capture more complex interactions.
LoCoHD (aazekas et al. 2024 is method for quantifying chemical and structural differences between protein environments. Unlike aligmment baannent-baeed or purely geometric measures, LoCoHD characterizes each local environment as a distribution of chemical "primitive types" such atom
32

---

Published as a conference paper at ICLR 2026
types residue centroids within specified radius. The similarity between two environments is then measured using the Hellinger distance between their respective distributions. this paper, to test the effectiveness of LoCoHD descriptors in representing local environments, LocoHD embeddings using its representation a local environment while computing the similarity metric. Each embedding computed by aggregating statistics of the primitive types, weighted by their distance from the central residue. This results in purely structural and chemical descriptor of the local environment.
construct
E DATA CURATION
The dataset used within the scope of this paper was sourced from RefDB (Zhang et al. 2003 which
subset of the Biomolecular Magnetic Resonance Databank (BMRB) (Hoch al. 2023). RefDB a curated list of BMRB entries with calibrated chemical shifts. We select monomers from the RefDB and perform sequence redundancy filtering with mmseqs via,
mmseqs easy-cluster sequences. fasta --min-seq-id 0.5 -c 0.8
--cov-mode
which results in sequence clusters. We then split the clusters to define train and test sets. Our final dataset consists of 1048 BMRB entries in RefDB (225 test 823 train), from where collected the amino acid sequences, and the experimental chemical shifts.
Why use AlphaFold structures. Instead of using Protein Data Bank (PDB) structures, predict 3D structures for all entries with AlphaFold2 (Jumper et al., (021) via OpenFold (Ahdritz et al. Using experimental crystal structures would introduce several limitations:
2024
Poor generalization. X-ray crystallography covers only a small fraction of the proteome, meaning a model trained on PDB-derived environments would not generalize to most sequences without available structures.
Different experimental conditions. Crystal structures and NMR represent different physical states (solid vs. solution), leading to systematic differences in local geometry that can affect chemical shifts.
Sequence mismatches and ambiguity. We observed ~5..8% sequence misalignment between BMRB entries and corresponding PDB structures. single BMRB ID map multiple PDB entries (e.g., BMRB 15064 - 16 PDBs), making it unclear which structure should be used as reference.
AlphaFold alleviates these issues by providing systematic and consistent structural predictions for every sequence, including those with no experimental structure, enabling uniform training and evaluation.
Local protein environments are extracted from each AlphaFold structure as described in Algorithm Random-coil chemical shifts used in our predictor are obtained from UCBShift2 (Ptaszek al. 2024).
Due to the limited availability of high-quality experimental training data, our pK, prediction models are regressed over synthetic values determined by solving the Poisson-Bollee Boltzmann equation via the method of Reis et (2020)
F ENVIRONMENT CONSTRUCTION AND MLFF EMBEDDING
This section elaborates on the amino acid centric environment construction algorithm discussed in Section
33

---

Published as a conference paper at ICLR 2026
Algorithm 1 Amino Acid Centric Environment Extraction
Require: Protein dataset D radius T'mx 5A
1: & <- 0
2: for all structure E D do
D R RR-an
3: for all amino acid X' do
4:
for all atom x; E a do
D D E E
5:
Xa < {a}
6:
for all amino acid a' E X' do
7:
if xx a' s.t. Xx;
xill
Ymax then
8:
N + Na U {a'}
All atoms in a'
9:
end if end for - E U Xa
10:
12: end for
13:
end for
14:
end for
15:
return
F.1 RADIUS ABLATION
To select the radius used for environment construction, we performed an ablation study over e (4, 5, 6) A using MACE (Kovács et al. 2023) embeddings for backbone chemical shift prediction. MACE is trained with a cut-off of A, so larger radii are meaningful. The results (Table B.9 show that radius of 5 Å performs best. We hypothesize that this radius is large enough to capture the relevant NMR chemical-shin environment while still keeping most residue neighborhoods as complete as possible under the 6 Å cut-off, thereby providing good balance between coverage and truncation.
F.2 MLFF LAYER COMPARISON
This ablation study examines how selecting different hidden layers MLFF embeddings affects performance for MACE and Orb. MACE has two hidden message-passinng layers; throughout this work, the MACE descriptor is defined as the concatenation of the invariant components of their hidden states, matching the default behavior of the mace. get dest.desrr ptoos function. For all other MLFF models, the descriptor corresponds to the hidden state of the final message passing layer. Orb has fifteen hidden message-passing layers, so in this case the descriptor is given by the hidden state of the 15th layer.
Table B. 10 summarizes the per-atom mean absolute errors for all evaluated layers on the chemical-shift prediction task. For MACE, the concatenation of its two hidden layers yields the best performance for most atom types, while for Orb the 10th hidden layer performs best in most cases. Nevertheless, throughout this work we use the 15th Orb layer to maintain consistent configuration across MLFFs with many layers, since concatenating all hidden states would result in impractically high-dimensional descriptors.
G PREDICTING AMINO ACID, SECONDARY STRUCTURE, AND PK
In what follows, we provide additional details for experiments and results described in Section
G.1 SECONDARY STRUCTURE AND AMINO ACID PREDICTION
Target variables. In these tasks, the target space of the target is categorical variable.
In amino acid prediction, the target variable holds of the following categories: {A, C, D, E, F, G, K, N, P.Q. R, S, V. W, Y), each representing standard twenty amino acids (cf. (Kawashima Kanehisa 2000) for further explanation).
34

---

Published as a conference paper at ICLR 2026
In secondary structure prediction, the target variable
of the following
categories: H, E, O), where denotes an a-helix, E denotes 8-strand, and O corre- sponds to all other secondary structure types (see Table Although more secondary structure types exist, we are primarily interested a-hellces and f3strands due to their structural prominence.
Loss function.
To train models for these tasks, we employ cross entropy loss written as
exp(iic) (Urie)
C(v,,
Σ Vio log
(10)
Σ exp((tt)
where 9; the predicted logit for each class and y; the one-hot encoding of the class label for sample i.
Evaluation metrics.
To evaluate the model's performance, we use the following metrics.
Precision: For a given class i, precision is defined as
True Positive;
Precision; True Positive, False Positive;
Recall: For a given class recall is defined
True Positive;
Recall;
True Positive; + False Negative,
score: For a given class i, the Fl score is the harmonic mean of its Precision and Recall:
Fl; = Precision, Recall;
Precision, Recall,
For secondary structure prediction report the precision, recall, and F1 score for each class in In addition, we report an unweighted average and weighted average (mean weighted by number of instances in each class) for the respective tasks.
Ablation study. We conduct an ablation study on the set of atoms A, to evaluate the effect of including descriptors from multiple atoms within amino acid secondary structure and amino acid prediction performance. Specifically, two configurations are compared: one where A₂ {C, CA, CB, N, H. HA}, and reduced version Aa (CA). For this ablation, identical models are trained using Egret descriptors, varying only the atom set. Egret descriptors are used because as shown in Tables B.3 and [B.4] they consistently outperform other atom-level descriptor types. The ablation results for secondary structure prediction are presented Table B. 13 and for amino acid prediction in Table B.14 Clearly, incorporating larger of atoms in A, leads substantial improvement in prediction accuracy.
G.2 ACID DISSOCIATION CONSTANT (PK,) PREDICTION
Loss function. We treat pKa prediction as regression problem. To train pKa regressors, we train graph neural network minimize the mean absolute error (MAE) as follows,
N y:
C,
(11)
where N is the number of samples in the dataset, y is the groundrruth pK, value of protonated site Poisson-Boltzmamn Boltzmann equation from Reis et al. (2020), and Ø is
in protein obtained by solving the Poisson- the predicted pKg.
35

---

Published as a conference paper at ICLR 2026
Comparison experimental data. To quantitatively assess the practical performance of
approach, we compare against ANN Gokcan & Isayev (022) and PROPKA ArepKoonooon Olsson et al. (2011)
on the benchmark summarized in Table B.7 Both reference methods were originally trained experimental pK, values derived from PDB structures, in contrast our models were trained on simulated values and predominantly AlphaFold2 structures.
For pK,-ANI, we evaluate on the PKAD-R (Ancona et al. 2023) dataset, a curated and updated version of the original PKAD dataset used for their training. Because no explicit train-fest split was reported, we benchmark both models on all entries in PKAD-R. For PropKa, we evaluate on their training set, for which exact predictions were provided. The authors of PropKa also mention evaluation on set of 201 proteins assembled from three previously published datasets, but they do not specify the corresponding PDB identifiers make their predictions available, preventing direct comparison on that benchmark.
These comparisons are conservative with respect to our method, because the baselines are trained on experimental PDB-based data whereas models are trained on simulated labels and AlphaFold structures. Nevertheless, our model outperforms PropKa on more than half of the entries and matches or surpasses pK,-AN on 3 out of residue types. For the remaining residue type (Histidine), we hypothesize that the weaker performance is due to the substantially smaller number available training examples.
G.3 MODEL ARCHITECTURE
Across all three tasks, we used consistent model architecture with little changes in hyperparameters. To incorporate the spatial and relation information between atomistic features, we constructed a fully connected graph between the atoms of the same residue. We propose to use a graph convolution network (GCN) (Kipf & Welling 2016) to predict the specific class type. The update rule at layer I E is defined as,
( or (AX (W))
(12)
Here, A (I A) is the normalized adjacency matrix that serves low-pass filter when aggregating information over the neighbors. A is the graph's adjacency matrix, and is an identity matrix. Also, (!!) is the learnable weight matrix to transform the features at layer 1, X the node feature matrix defined for layer I of the GNN. We used SiLU (Sigmoid Linear Unit) (Elfwing et al., 2018) as the non- non-linear activation function and we applied layer normalization after each graph convolution operation stabilize training. The operation in Eq. T2 is repeated for L 10 layers. The initial input X (0) consists of the atomistic features for each atom. At the final layer (1 = L 1), the case of classification, the network outputs a vector shape ly representing class-wise logits, and the network outputs scalar for regression tasks. Lastly, use dropout (Srivastava al.) 2014) for every intermediate layer to prevent overfitting. Below is the table of ayper-parameters used for these experiments.
Aa CA, H,
Batch size: 256
Number of Epochs: 800
Optimizer
Type:: Adam (Kingma & Ba 2015)
.Learning Rate: 10-4
Learning Rate Scheduler
Type: Step Learning Rate
Decay Factor (y): 0.5
Step Size: 10
Number of GCN layers: 10
Dropout probability: 0.3
Hidden dimension of GCN (per layer): 128
36

---

Published as a conference paper at ICLR 2026
In addition, employed gradient clipping with maximum norm of 5, exponential moving average (EMA) weighted model with a decay rate of 0.999, and a StepLR learning rate scheduler incorporating 7500 warm-up steps and step size of 10.
H USING LIKELIHOODS TO DEFINE ENVIRONMENT SIMILARITY MEASURE
We employ the conditional likelihoods defined in Section | define similarity metric compare protein environments,
H.1 COMPUTATION DETAILS
All experiments were performed on single NVIDIA H100 GPU. Individual runs required one hour to complete, with maximum GPU memory usage of approximately 15 GB.
H.2 SIMILARITY METRIC
We analyzed large collection of protein chemical environments by comparing them using MACE embeddings. For each attribute such as amino acid type or secondary structure we estimated a likelihood around each embedding via kernel density estimation with fixed bandwidth as described in Section 5] These likelihoods were converted into distances using a scaled negative exponential, where higher similarity corresponded smaller distances. Since each pairwise distance arises from distribution based on multiple environments belonging to the same attribute we applied stochastic multidimensional scaling (MDS) (Rosenberg et al. 2022; [Boyarkii & Bronstein 220TT, which takes into account both the mean and standard deviation of these distances. The resulting embeddings revealed meaningful clustering that not only separated different amino acid types and secondary structures but also captured finer relationships within them. Chemically structurally related amino acids and secondary structures exhibited higher similarity, demonstrating that the MACE-based comparison effectively reflects the underlying chemical and structural properties. Additional details and visualizations are provided Figure A.13
I CHEMICAL SHIFT PREDICTION
1.1 EXTENDED SEXSICS-GROUNDER CHEMICAL SHIFT RESULTS
In Section 6 we presented one representative case study to examine whether the shift predictor correctly captures local structural influences on chemical shifts. Here, we provide two additional case studies further validate these findings. Both probe distinct types of conformational variation and demonstrate that the proposed predictor consistently produces physically meaningful chemical shift trends.
Helix unfolding into a strand. Chemical shifts of CA and CB have distinct behaviors in helices and strands. We perform an MD simulation of a 7-residue peptide unfolding from a helix to strand. Predicted chemical shifts along the trajectory are depicted in Fig. A.8 We note that our method correctly reproduces the experimentally observed chemical shift distributions for CA and CB atoms. It clearly captures the characteristic pattern where CA shifts decrease in strands compared to helices, while CB shifts exhibit the opposite trend.
Alternative conformations. We evaluate shift predictor over 100 MD simulation of the protein :CLE: 3 from Rosenberg et al.] (2024), which features two stable conformers: conformer A containing helix and conformer with linearly structured loop (Fig. A.12). Our method captures distinct secondary chemical shift distributions for lysine 63 CA and CB atoms. The helical conformer
exhibits the expected increase in CA secondary shift and decrease in CB secondary shift.
1.2 DETAILS OF THE CHEMICAL SHIFT PREDICTION EXPERIMENTS
Target variable. The goal of the chemical shift prediction is to predict the calibrated chemical shift y of an atom of interest given the environment. We train our models to predict the secondary shift, i.e. the difference in the experimental and random-coil chemical shifts, given the sequence.
37

---

Published as a conference paper at ICLR 2026
Loss function. error, defined as
We train shift prediction models minimize the mean absolute shift prediction
L1
(13)
where is the number of data points, yi and ÿ are groundtruth and predicted chemical shifts, respectively.
Evaluation metries. To quantitatively evaluate the accuracy of the model, we report the mean absolute eтror measured on the test set.
1.3 MODEL ARCHITECTURE
We used the same architecture before (Appendix G.3 with slightly different hyperparameters as described below.
A. (b), where b is the atom type for which we are predicting the chemical shift value.
Batch size: 500
Number of Epochs: 2000
Optimizer
Type: Adam (Kingma & Ba 2015)
-Learning Rate: x 10
Learning Rate Scheduler
Typpe: Step Learning Rate
Decay Factor (y): 0.5
Step Size: 10
Number of GCN layers: 5
Dropout probability: 0.3
Hidden dimension of GCN (per layer): 256
In addition, employed gradient clipping with maximum norm of 5, exponential moving average (EMA) with a decay rate of 0.999, and a StepLR learning rate scheduler incorporating 7500
warm-up steps and step size of 10.
1.4 COMPUTATION DETAILS
All experiments were performed on single NVIDIA H100 GPU. Individual runs required three hours to complete, with maximum GPU memory usage of approximately 15 GB.
J EXTENDED MLFF INTERPRETATIONS
J.1 HELIX UNFOLDING
As a complement the main-text case study in Section 7 we analyze the behavior of MLFF embeddings during a helix-to-strand transition. Using MD simulations of 7-residue peptide (Fig. A.9. we perform PCA on the MACE embeddings of CA atoms in residues 6 and 7, and observe that they follow smooth, non-intersecting trajectories. These trajectories reflect the gradual unfolding process, indicating that the descriptor space encodes continuous and interpretable structural transitions.
To further quantify how the learned descriptors relate to structural changes, we compare pairwise dis- tances in descriptor space with backbone Ca RMSDs along the same trajectory (Fig. A.10). For each pair of MD frames, we compute the Euclidean distance between the vectorized MACE embeddings of all Co atoms in the protein and plot this distance matrix alongside the corresponding RMSD)
38

---

Published as a conference paper at ICLR 2026
matrix, well as scatter plot of descriptor distance versus RMSD. The strong Pearson/Spearman correlations (p = 0.92/0.99 /0.98) and the similar block structure of the two matrices show that confor- mations that close in 3D space also close in descriptor space, while structurally dissimilar states are well separated. This approximate isometry implies that the MLFF descriptors preserve the geometry of the unfolding pathway, thereby supporting their interpretability: distances and trajectories in embedding space can be directly related to physically meaningful structural deviations.
J.2 INVERTING MACE EMBEDDINGS
In this section, we explore whether MACE descriptors can be inverted to recover protein structures and distinguish between alternate conformations (altlocs). We treat MACE features not only as structural representations but also potential tools to decode conformational states. We present two complementary approaches: (i) direct optimization of atom positions to fit MACE embeddings, and (it) in line with |Mddiiptla al. (2025). guiding AlphaFold3's reverse diffusion process using MACE descriptors.
J.2.1 DIRECT OPTIMIZATION OF ATOMIC COORDINATES
We first investigate whether MACE features can serve latent representations from which protein conformers can be reconstructed. As a test case for descriptor-based structure inversion, we examined the protein 40LE: : B, which contains distinct alternate conformations. In addition to this example, we applied the same analysis to the proteins listed in Table 15 3v35: R, 7EC8: A, 5G51 : A, and :NPU: B. Each of these structures contains a short region annotated with alternate conformations, typically involving subtle backbone displacements rather than large-scale structural rearrangements. For every protein, atomic coordinates in the altloc region so that the local MACE descriptors of one conformation match those of the other, enabling us to assess how well the descriptors capture and recover experimentally observed conformational differences. The resulting optimized structures are shown in Fig. A.16 and Fig. A.15
mine-residue segment (residues 60-68) modeled as two
optimized the
The optimization loss function quantifies the discrepancy between the descriptors of the current and target structures and is computed for each atom within the optimized region. In particular, we define each local atomic environment as all atoms within a 5 Á radius around the atom of interest. The resulting descriptors capture both geometric and chemical context at multiple equivariant orders. The loss is composed of three terms:
N
a B B
B target
8,00
+
B Blnck.
B Blarget. kn LM
Σ B 2.
B targer,i
2.
where Bi B. and Barget K.M denote the L-th order MACE embeddings of the optimized and target conformations, computed as described in Equation 3 For a given L,, the Clebsch-Gordan coefficients C.L.M. select the valid combinations of input angular momenta. While the second and third terms are not ensures validity of the comparison.
global alignment of the source configuration with the target
This combined loss (weighted by a. (.7,7) drives the structural optimization by progressively aligning the descriptor representations of the initial and target altloc conformations. As illustrated in Fig. A. the backbone atom positions are recovered with high accuracy, but the side chains of glutamic acid (GL.U 62) and lysine (LYS 63) adopt incorrect orientations. We hypothesize that this discrepancy arises because MACE features primarily capture the local chemical environment within given radius, which cannot always resolve distinct side-chain rotamers.
These results suggest that while MACE features reliably capture backbone geometry, they are less sensitive to side-chain orientations. Recovering both backbone and side-chain configurations would require further exploration, but success in this direction could enable the of MACE features (and other atomistic embeddings) for coarse caase-grained protein modeling with minimal loss of structural fidelity.
39

---

Published as a conference paper at ICLR 2026
J.2.2 GUIDING ALPHAFOLD3 REVERSE DIFFUSION
We next investigate whether MACE descriptors can guide AlphaFold3's diffusion-based structure generation process. Specifically, we focus on proteins with alternate conformations (altlocs)- structural heterogeneity intrinsic to the protein within the unit cell. We analyze subset of four proteins (Table B.15) exhibiting distinct backbone-separate altlocs, from Rosenberg et al. Rosenberg et al. 2024).
Starting from Gaussian noise, we guide AlphaFold3's reverse diffusion process to recover target protein structure that minimizes a descriptor-basse loss L. The guided reverse SDE Song Ermon (2019) is given by
dx = - -f log a) 7 B,dt B,dt VB,,
where X are atomic coordinates, is the amino acid sequence, 3e is the noise schedule, and N -V I). The term Vx log Px( a) is the unconditional score function, while Vx provides descriptor-based guidance. The hyperparameter scales the guidance term.
To invert MACE embeddings, we define L as a discrepancy between the descriptors of the di diffusion variable X and those of the target conformation Xaaaget in the altloc region. The atomic environment specified as all atoms within radius:
(X,)
larget
X target)
)|
if Bo
iarget (ttargtt) S 8,
0
2
(14)
B
targer
target
|||
26 otherwise.
Here, Bu is the zeroth-order invariant descriptor derived from X,, and target is the corresponding descriptor the target conformation. The threshold switches the loss from 1.2 (small errors) to (large errors). For all structures, we set r 0.5 and 8 = 0.25.
As shown in Fig. A.14 A, unguided AlphaFold3 typically captures only one of the two alternate conformers. By guiding the reverse diffusion process with MACE descriptors, we can recover the alternate conformation that AlphaFold3 misses or preserve its original prediction. To quantitatively assess the distribution, we compute the normalized distance of each generated structure X relative to conformations XA and X modeled in the PDB (Burley et al. 2017
dA
X
X
dg
X
X
2
dy dp )]
dB dA
Normalized Distance
min
sign(ds dn).
(15)
Negative values indicate proximity to conformation A, while positive values indicate proximity to conformation B. The resulting distributions (Fig. A.14 B) again show that guiding the diffusion process with descriptors shifts the distribution toward the mode AlphaFold3 misses or preserves the original distribution. Both results indicate that MACE descriptor-guide diffusion reliably steers AlphaFold3 toward accurate backbone geometries of both conformations in most cases.
K
USE OF LARGE LANGUAGE MODELS (LLMs)
The use of L.LMs was restricted to minor proof-reading, stylistic polishing, and correcting typos in the text. All intellectual contributions, experiments, data analyses, and conclusions were executed and verified exclusively by the authors. Responsibility for the scientific content rests with the authors, and part of the research process delegated to LLMs.
40


Reviews:
[{"review_id": "PcedQpzSMg", "reviewer": "Reviewer_AnCg", "summary": "This work proposes to leverage pretrained ML interatomic potential embeddings as canonical descriptors of local protein environments. Embeddings are extracted from a residue-centered neighborhood. Various applications including pKa prediction of titrable residues and protein NMR chemical shift prediction are explored, with MLFF-based embeddings show good representative power for downstream tasks.", "questions": "- How sensitive are results to the MLFF layer chosen for embeddings?\n- MACE is an architecture, but there are several variants according to the dataset it is trained on. It seems like the MACE-OFF23 is used in this work, could authors elaborate?\n- For AF3-guided structure selection in Fig A14, could authors report the optimization statistics similarly to Fig A13? (What fraction of optimization result in the target structure, or whether it always lead to the target structure)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Repurposing latent features of MLFFs as canonical protein descriptors is a timely, well-motivated idea that links quantum-level atomistic modeling with biomolecular representation learning.\n- The paper covers a solid range of downstream tasks tied to experimental observables and provides thorough analysis of the physical plausibility of its predictions.", "weaknesses": "- Dataset and baselines: For the pKa and NMR shift tasks, the baselines are evaluated in conditions that differ from their intended use, whereas the MLFF-feature models introduced here are trained directly for the target objective. For instance, the pKa baselines are designed to predict experimental values, while the proposed methods are trained to reproduce a cheaper computational reference. This creates a benchmark mismatch, since the baselines are not optimized for the reference chosen in this work. Because the core question is whether MLFF features are useful, a more appropriate primary comparison would be a standard GCN with simple learned embeddings rather than prior task-specific baselines. This issue is compounded by the use of AFDB structures as inputs.\n- Experimental design: Some experiments are not very indicative of how the embeddings would be used in practice. Inferring amino acid type or secondary structure from full atomic coordinates (Section 4) is a trivial task under those inputs, and the distribution-shift analysis in Section 5 (Fig. 2) mainly shows that embeddings from energy-relaxed structures are similar, which is expected since both MLFFs and classical force fields approximate the same physical energy landscape. Several experiments currently placed in the appendix appear more compelling and better aligned with the paper’s motivation. I recommend reorganizing the manuscript so that the most informative use cases are in the main text, especially given that all experiments are listed as contributions but not all are presented in the main body.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cfD0Tm3rgx", "reviewer": "Reviewer_dp5R", "summary": "This paper proposes repurposing intermediate embeddings from pre-trained Machine Learning Force Fields (MLFFs) as \"physics-grounded\" feature representations for local protein environments. The authors benchmark embeddings from MACE, AIMNet, and others, claiming state-of-the-art (SOTA) results in $pK_{a}$ prediction (reportedly outperforming pKa-ANI) and NMR chemical shift prediction (reportedly outperforming UCBShift2-X). The method's physical realism is validated via case studies, such as aromatic ring current effects.", "questions": "see weakness", "limitations": "", "rating": 4, "confidence": 3, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "1. Repurposing MLFFs as \"foundation models\" for structural biology is in nnovative and valuable.   \n2. The design of the validation experiments (e.g., ring current effect , helix unfolding ) is a commendable standard for physical realism.   \n3.The paper effectively demonstrates the general-purpose nature of the embeddings for zero-shot clustering and generative guidance.", "weaknesses": "1.The paper's central claims rest on comparisons against pKa-ANI and UCBShift2-X that appear to be factually incorrect or based on flawed implementations.\n\n2.The $pK_{a}$ evaluation is missing the entire 2024/2025 SOTA, invalidating its performance claims.\n\n3.The paper compares its structural embeddings against sequence (ESM) embeddings. It critically fails to benchmark against the most obvious and relevant competitors: other structural embeddings, namely those from the AlphaFold2 or ESMFold structure modules.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lBpwnpZL5c", "reviewer": "Reviewer_RxQj", "summary": "Local protein environments contain highly relevant chemistry that affects its function and interactions. This paper benchmarks existing machine learning force field (MLFF) methods on their ability to understand local protein structures. They evaluate a suite of MLFFs on their ability to predict the protonation state, secondary structure, and amino acid types.", "questions": "Are there any other interesting local protein structure properties to evaluate on?\n\nHow well do non-FF based local structure models, such as ESM3, perform on these benchmarks?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper rigorously evaluates multiple popular MLFF methods on across benchmarks and connect them to some protein structure related tasks.\nThey perform an interesting analysis on the chemical shift prediction.", "weaknesses": "The evaluation benchmarks such as secondary structure, amino acid type prediction are a bit straightforward.\n\nWhile MLFF naturally learn the physics of local structural environments, there are other machine learning based approaches that reason over the local structure and are suitable to predict protonation state, secondary structure, and amino acid types. This work does not benchmark MLFF versus these methods on representing local structure.\n\n[1] Simulating 500 million years of evolution with a language model. Hayes et al.\n[2] 3D deep convolutional neural networks for amino acid environment similarity analysis. Torng et al.\n[3] Distilling Structural Representations into Protein Sequence Models. Ouyang-Zhang et al.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HWRYLMQkQf", "reviewer": "Reviewer_fbDS", "summary": "This paper explores the use of MLFFs as general-purpose representation learners for local protein environments. Instead of relying on sequence-based or handcrafted descriptors, the authors repurpose latent embeddings from pretrained MLFFs (AIMNet, MACE, OrbNet, Egret) as compact, physics-grounded descriptors for atomic neighborhoods in proteins. The study benchmarks MLFF embeddings on diverse downstream tasks—secondary structure and amino acid classification, pKa prediction, and NMR chemical shift regression—and shows that these representations outperform or rival specialized baselines such as PropKa and pKa-ANI. Overall, the paper establishes MLFFs as reusable, physics-informed foundation models for structural biology.", "questions": "1. Could the same MLFF-based representation transfer to nucleic acids or protein–ligand complexes, where local environments include non-canonical atoms and charges?\n\n2. MLFFs have multiple internal layers encoding different orders of interaction (0th, 1st, 2nd). Did the authors investigate which layer yields the most informative embeddings for downstream tasks?\n\n3. Since MLFF embeddings originate from networks with different scales and symmetries, how are they aligned or normalized across model families?\n\n4. The authors mention uncertainty-aware predictions. Is uncertainty derived from ensemble variance, likelihood width, or another calibration technique?\n\n5. Could MLFF embeddings serve as complementary features to pretrained protein language models, bridging physics-based and sequence-based representations?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The central insight, repurposing MLFF embeddings as general-purpose, transferable protein environment representations, is both novel and timely. Most prior MLFF applications focus on energy or force prediction for small molecules; extending them to protein representation learning is original and valuable. The paper effectively bridges quantum-chemistry-based potentials and protein machine learning.\n\n2. The experimental setup is comprehensive. The authors curate 165 k environments from 1048 proteins and evaluate four MLFF families on four biologically relevant tasks. Comparisons include classical and ML baselines (PropKa, pKa-ANI, UCBShift2-X). Results demonstrate meaningful improvements in both accuracy and interpretability. The inclusion of uncertainty quantification and physical consistency tests (e.g., ring-current effects) adds rigor.\n\n3. The methodology that extracting embeddings from pretrained MLFFs and mapping them to canonical residue-centered environments is sound and clearly motivated. Statistical reporting (mean absolute errors, standard deviations) is adequate, and all experiments appear well-controlled.\n\n4. The paper is well written and pedagogically organized. The introduction clearly motivates the challenge of representing local protein environments; figures (e.g., Fig. 1 and 2) effectively illustrate how embeddings are constructed and used. Terminology (canonical environment, focus residue, MLFF feature extraction) is consistent and accessible even to readers outside computational chemistry.\n\n5. This work could substantially influence both computational biology and machine-learning communities by providing a physics-consistent alternative to sequence-only protein language models. MLFF embeddings encode quantum-derived information unavailable in existing representations and show transferability to tasks requiring local chemical precision. The idea of using pretrained MLFFs as “foundation models for atoms” could be broadly significant.", "weaknesses": "1. The study restricts environments to 5 Å radius regions; while suitable for local chemistry, it omits long-range electrostatic or conformational effects. For tasks like folding or binding prediction, this locality may be insufficient. A discussion or experiment extending to multi-scale contexts would strengthen generality claims.\n\n2. MLFFs such as MACE and OrbNet require quantum-level pretraining on millions of molecules. Although embeddings are reused, the computational barrier to obtaining them limits accessibility compared to pretrained sequence models (e.g., ESM, ProtT5). The paper could better address scalability and efficiency trade-offs.\n\n3. Qualitative case studies (e.g., helix → strand unfolding) are insightful but anecdotal. Quantitative metrics (e.g., correlation between embedding distances and RMSD/chemical similarity) would solidify interpretability claims.\n\n4. While the authors compare to physics-based predictors, they do not benchmark against modern structure-aware geometric encoders (e.g., GVP-GNN, ProteinNeRF, FrameFold). Including such baselines would clarify whether MLFF embeddings provide advantages beyond standard geometric message-passing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)