4.14. Fit energy Hessian 
#
Note
Supported backends: PyTorch-TorchScript , JAX
To train a model that takes Hessian matrices, i.e., the second order derivatives of energies w.r.t coordinates as input, you only need to prepare full Hessian matrices and modify the loss section to define the Hessian-specific settings, keeping other sections the same as the normal energy model’s input script.
4.14.1. Energy Hessian Loss#
If you want to train with Hessians, add the start and limit prefactors of Hessians, i.e., start_pref_h and limit_pref_h to the standard energy loss section in the input.json:
"loss": {
"type": "ener",
"start_pref_e": 0.02,
"limit_pref_e": 1,
"start_pref_f": 1000,
"limit_pref_f": 1,
"start_pref_v": 0,
"limit_pref_v": 0,
"start_pref_h": 10,
"limit_pref_h": 1
},
The legacy loss type "ener_hess" remains accepted as an alias of "ener", but new input files should use the canonical "ener" type.
Setting either Hessian prefactor to a nonzero value enables Hessian supervision, so schedules may ramp the term up from zero or down to zero. Both prefactors must be zero to disable the term. Earlier releases could silently disable the term on some backends when exactly one endpoint was zero.
The options start_pref_e, limit_pref_e, start_pref_f, limit_pref_f, start_pref_v and limit_pref_v determine the start and limit prefactors of energy, force, and virial, respectively. The calculation and definition of Hessian loss are the same as for the other terms.
If one does not want to train with virial, set the virial prefactors start_pref_v and limit_pref_v to 0.
4.14.2. Hessian Data Format#
In the PyTorch-TorchScript and JAX backends, Hessian matrices are listed in hessian.npy files, and the data format may contain the following files:
type.raw
set.*/box.npy
set.*/coord.npy
set.*/energy.npy
set.*/force.npy
set.*/hessian.npy
This system contains Nframes frames with the same atom number Natoms, the total number of elements contained in all frames is Ntypes. Most files are the same as those in standard formats, here we only list the distinct ones:
ID | Property | Raw file | Unit | Shape | Description |
|---|---|---|---|---|---|
hessian | Hessian matrices | hessian.npy | eV/Å^2 | Nframes * (Natoms * 3 * Natoms * 3) | Second-order derivatives of energies w.r.t coordinates. |
Note that the hessian.npy should contain the full Hessian matrices with shape of (3Natoms * 3Natoms) for each frame, rather than the upper or lower triangular matrices with shape of (3Natoms * (3Natoms + 1) / 2) for each frame.
4.14.3. Train the Model#
There are two approaches to training a Hessian model. The first method involves training the model from scratch using the same command as in the ener mode:
dp --pt train input.json
dp --jax train input.json
The second approach is to train a Hessian model from a pretrained energy model, following the same command as the finetune strategy:
dp --pt train input.json --finetune pretrained_energy.pt
dp --jax train input.json --finetune pretrained_energy.jax
The detailed loss can be found in lcurve.out:
# step rmse_val rmse_trn rmse_e_val rmse_e_trn rmse_f_val rmse_f_trn rmse_h_val rmse_h_trn lr
0 1.05e+02 2.28e+01 2.11e-01 1.59e+00 3.25e+00 3.37e-01 6.00e+00 6.37e+00 1.0e-03
200 1.86e+01 3.23e+01 9.24e-03 1.54e-01 2.51e-01 4.70e-01 5.31e+00 9.05e+00 1.0e-03
400 2.69e+01 2.98e+01 1.03e-01 1.07e-01 5.67e-01 4.17e-01 6.35e+00 8.47e+00 1.0e-03
600 2.00e+01 1.90e+01 7.23e-02 6.90e-03 3.35e-01 2.58e-01 5.37e+00 5.41e+00 1.0e-03
800 1.68e+01 1.48e+01 4.06e-02 2.27e-01 2.35e-01 1.98e-01 4.76e+00 4.24e+00 1.0e-03
1000 1.70e+01 1.81e+01 3.90e-01 1.66e-01 2.02e-01 1.99e-01 4.98e+00 5.37e+00 1.0e-03
4.14.4. Test the Model#
Warning
The PyTorch-TorchScript freeze route does not preserve Hessian output. A PyTorch-TorchScript model frozen with dp --pt freeze is treated as a standard energy model. PyTorch-Exportable cannot construct a Hessian model for freezing or inference. The JAX backend can preserve Hessian output in a frozen model with dp --jax freeze --hessian.
The PyTorch-TorchScript tab below tests the frozen model as a standard energy model, while the JAX tab preserves and tests the Hessian output:
dp --pt freeze -o frozen_model.pth
dp --pt test -m frozen_model.pth -s test_system -d ${output_prefix} -a -n 1
dp --jax freeze -c . -o frozen_model.hlo --hessian
dp --jax test -m frozen_model.hlo -s test_system -d ${output_prefix} -a -n 1
For the PyTorch-TorchScript frozen-model command, the output files are the same as those in the ener mode, i.e.,
${output_prefix}.e.out ${output_prefix}.e_peratom.out ${output_prefix}.f.out
${output_prefix}.v.out ${output_prefix}.v_peratom.out
The PyTorch-TorchScript checkpoint can also be tested directly without freezing:
dp --pt test -m model.pt -s test_system -d ${output_prefix} -a -n 1
The JAX training checkpoint (.jax) cannot be tested directly: the JAX dp test route accepts only frozen .hlo or .savedmodel artifacts, so a JAX Hessian model must be frozen with dp --jax freeze --hessian (as shown above) before testing.
When either backend tests a Hessian-capable checkpoint with -d ${output_prefix} -a, the predicted Hessian for each frame is written to an additional file in the working directory:
${output_prefix}.h.out
For *.h.out.*, it contains matrix with shape of (2, n_hess):
# frame - 0: data_h pred_h (3Na*3Na matrix in row-major order)
5.897392891323943331e+01 2.909700516268236825e+01
-7.682282297964052376e+00 2.535680817045881774e+00
-1.266442953072092514e+01 -2.127310638041492652e+01
5.442541716174009031e-02 7.202825779190234756e-02
5.198263170894957939e-05 -8.110080221576332349e-02
7.443552765043950914e-02 -2.248597801730128215e-02
1.029910175689553675e+00 1.938646932394622047e-03
1.213862217511276764e+00 5.344132558814301825e-02
-1.221943904909605250e+00 1.602557574981743893e-01
The full Hessian matrices are stored in a flattened form in the row-major order. Here, n_hess is the total number of Hessian matrix elements across all frames, calculated as:
where \(N_{\text{atom}, i}\) represents the number of atoms in the \(i^{\text{th}}\) frame.