Training Tasks We identified three training tasks that form a well-rounded test bench with their varying degrees of complexity. These tasks include different system types (Cobalt, a metal; Silica, an amorphous material; Disulfide, a molecular system), different types and numbers of items in the training datasets, and different number of parameters with their respective ranges to be optimized. While structures in the Cobalt and Silica datasets mostly require energy minimization, those in the Disulfide case mostly require single-step energy evaluations. Table summarizes the specifications of the selected training tasks. JAX-ReaxFF currently does not support training items with simulation cell optimization, as such these are ignored. This only affects the Silica dataset which has 5 of them (out of 296 cases).