Total Energy Alignment (TEA)

• Total Energy Alignment (TEA) is a methodology that uses analytic and data-driven post-processing to reconcile energy discrepancies in both ML models and atomistic simulations.
• In machine learning, TEA computes per-class energy shifts to debias energy-based models trained on imbalanced data, improving calibration without retraining.
• In quantum chemistry, TEA employs a two-step protocol (ICEA and AEC) to align energy scales across diverse computational methods, enhancing cross-dataset integration.

Total Energy Alignment (TEA) encompasses a class of methodologies developed to resolve misalignment and bias in energy-based models and atomistic simulation datasets. The term is instantiated in two major but distinct contexts: (i) statistical energy-based machine learning as Energy Aligning, and (ii) quantum chemistry/atomistic-scale simulations as a protocol for unified data integration across heterogeneous computational methods. Both employ principled post-hoc transformations to reconcile discrepancies, permitting robust model performance and union of disparate datasets. The following provides a comprehensive review of TEA, addressing its theoretical basis, algorithmic workflow, mathematical underpinnings, empirical results, and practical guidance for deployment (Zhao et al., 2021, Shiota et al., 2024).

1. Theoretical Motivation and Context

In supervised deep learning and statistical machine learning, TEA (as Energy Aligning) targets systematic class imbalance in energy-based models. Training on class-imbalanced data induces an energy bias, with free energies $E_\theta(y)$ tracking empirical label frequencies, resulting in under-representation of rare classes. Similarly, in computational quantum chemistry and materials science, TEA resolves the challenge that datasets generated using diverse electronic structure methods (differing in basis set, functional, pseudopotential, etc.) exhibit systematic offsets and scaling discrepancies in total energies and forces. These discrepancies historically mandated computationally prohibitive recalculation of molecular and crystalline configurations under a single reference method.

In both domains, TEA is motivated by the principle that systematic biases or incommensurabilities can be identified and eliminated by analytic or data-driven post-processing, without retraining or expensive recomputation.

2. Mathematical Formulation

2.1 Energy Aligning in Machine Learning

Consider a deep classifier with outputs $f_\theta(x) \in \mathbb{R}^C$ for input $x$, inducing softmax probabilities $p_\theta(y=i|x)$. Introducing an energy function $E_\theta(x, y) = -f_\theta(x)[y]$, the marginal free energy of class $i$ is defined as

$$E_\theta(y=i) = -\log \int_{x \in X} e^{f_\theta(x)[i]} dx.$$

It follows that

$$\log p_\theta(y=i) = -E_\theta(y=i) - \log Z_\theta,$$

implying $-E_\theta(y=i)$ is linearly aligned with $\log p_\theta(y=i)$. On class-imbalanced data, $f_\theta(x) \in \mathbb{R}^C$0 systematically favors majority classes.

The remedy is to compute per-class shifts $f_\theta(x) \in \mathbb{R}^C$1 such that

$f_\theta(x) \in \mathbb{R}^C$2

using Monte Carlo estimates: $f_\theta(x) \in \mathbb{R}^C$3 where $f_\theta(x) \in \mathbb{R}^C$4 denotes the log-sum-exp over $f_\theta(x) \in \mathbb{R}^C$5 samples.

2.2 TEA for Atomistic Simulation Data

Given total energy $f_\theta(x) \in \mathbb{R}^C$6 computed under method $f_\theta(x) \in \mathbb{R}^C$7, and isolated-atom energies $f_\theta(x) \in \mathbb{R}^C$8, the atomization energy is

$f_\theta(x) \in \mathbb{R}^C$9

TEA involves two main steps:

1. Inner Core Energy Alignment (ICEA):

$x$0

shifted total energy:

$x$1

2. Atomization Energy Correction (AEC): After ICEA, any residual slope offset in atomization energy is corrected by a global scale factor $x$2 from least-squares fit:

$x$3

aligned total energy:

$x$4

Aligned forces are scaled accordingly: $x$5

3. Algorithmic Workflow

The high-level pipelines for TEA are concise.

3.1 Energy Aligning (ML)

1. Offline: Sample a small, balanced set $x$6.
2. For each class, compute
   • $x$7
   • Per-class shift $x$8 as above.
3. At inference:
   • Compute $x$9 for each $p_\theta(y=i|x)$0.
   • Apply softmax for debiased prediction.

3.2 TEA for Simulation Data

Input:
    D_target: { ( {r_i}, E^{[1]}_{\rm tot}, F^{[1]} ) }
    D_aux:    { ( {r_j}, E^{[2]}_{\rm tot}, F^{[2]} ) }
    E_i^{P,[1]}, E_i^{P,[2]}: isolated atom energies

1. Compute Δ_core = sum_i (E_i^{P_i,[1]} - E_i^{P_i,[2]}).

2. For each sample in D_aux:
    a. E_at^{[2]} = sum_i E_i^{P_i,[2]} - E^{[2]}_{\rm tot}
    b. ICEA shift: E' = E^{[2]}_{\rm tot} + Δ_core

3. Fit a by least-squares using D_target & D_aux.

4. For each sample in D_aux:
    E_aligned = sum_i E_i^{P_i,[1]} - a * E_at^{[2]};
    F_aligned = a * F^{[2]}
Output: D_target ∪ D_aligned

4. Empirical Results and Impact

Application of TEA has led to notable advances in both fields.

• In ML, Energy Aligning yields balanced classification performance without retraining, outperforming existing post-hoc calibration schemes and requiring only forward passes for shift estimation (Zhao et al., 2021).
• In atomistic simulations, integration of TEA with the MACE-Osaka24 model permitted seamless fusion of inorganic (MPtrj, VASP/PBE) and organic (OFF23, ωB97M-D3(BJ), Psi4) datasets. This yields state-of-the-art accuracy not only within each chemical domain but also across them, with key benchmarks including:
  • Organic reaction barrier MAE reduction (Transition1x): from 0.937/0.519 eV (MACE-MP-0-large, inorganic-only) to 0.404/0.265 eV (MACE-Osaka24-large, TEA-merged).
  • Inorganic crystal lattice constant MAE: 0.018 Å (Osaka24-large) vs 0.016 Å (MP-0-large).
  • Near-unity agreement in liquid water radial distribution function with PBE-D3(BJ) (Osaka24-large-D3(BJ)) (Shiota et al., 2024).

Alignment RMSE on QM9 calibration (eV):

Comparison	Before TEA	After ICEA	After ICEA + AEC
VASP vs. ADF	0.326	0.326	0.099
VASP vs. Psi4	4.202	4.202	0.839

5. Limitations and Practical Considerations

Both flavors of TEA rest on well-defined assumptions. In energy-based models, balanced test distribution $p_\theta(y=i|x)$1 is presumed, with cluster-shift variants available for high cardinality or data-sparse regimes. The accuracy of the shift estimates $p_\theta(y=i|x)$2 hinges on the representativeness and balance of the calibration set (Zhao et al., 2021). TEA does not correct for covariate shift or distributional mismatch in $p_\theta(y=i|x)$3.

For simulation datasets, TEA’s validity depends on reliable isolated-atom energies under all requisite methods, and assumes residual differences are approximately linear in atomization energy. Nonlinear systematic differences (e.g., multi-reference or charged species) are not captured unless extensions to nonlinear correction are employed. The global scaling factor $p_\theta(y=i|x)$4 may be insufficient for certain classes of methods or molecules (Shiota et al., 2024).

Recommended practices include:

• Using a chemically diverse calibration set for fitting $p_\theta(y=i|x)$5,
• Removing extrinsic dispersion corrections from organic datasets prior to alignment,
• Verification of alignment via parity plots.

6. Extensions and Related Approaches

In high-class-count learning scenarios, a cluster-wise version of Energy Aligning can be deployed, estimating cluster-level shifts rather than per-class $p_\theta(y=i|x)$6 to reduce variance due to limited calibration data. A plausible implication is that, in atomistic data, finer-grained nonlinear corrections—potentially via piecewise regression or learned functions $p_\theta(y=i|x)$7—could further enhance alignment for data spanning wide chemical or method space diversity.

TEA is orthogonal to data reweighting, architectural, and training-based bias mitigation strategies, and can be integrated with representation learning and transfer learning workflows as a light-weight, post-hoc calibration layer.

7. Summary

Total Energy Alignment unifies two core applications: rectifying class bias in energy-based learning and achieving commensurate total energy scales across heterogeneous computational chemistry datasets. TEA operates via simple, physically and statistically motivated post-hoc transformations, enabling scalable deployment, competitive empirical performance, and broad adoption in both machine learning and materials modeling workflows (Zhao et al., 2021, Shiota et al., 2024).