After the final clusters are formed, parameter optimization is performed using gradient based local optimizers with multi-start as depicted in Fig. . Vectorization based parallelism is employed for both energy minimization and parameter optimization steps shown in this figure . For gradient based optimization to work, JAX traces the error function from Eq. ( ) and computes the gradients of the parameters. However, since typically many geometries require geometry optimization, tracing the gradients through the optimization step is more error prone due to the complex functional form of the ReaxFF. To remedy this problem, we separate the geometry optimization from the error minimization. The error function without the geometry optimization can be thought as a surrogate model since it is a fast way to approximate the true error where the geometry optimization is done as well. The approach accelerates the training and does not require tracing the gradients through the geometry optimization step.