In the original MBAR theory, Eqn.2 needs to be solved iteratively as it is a set of coupled nonlinear equations for all ĉi . However, in the spirit of reweighting, the computation of the estimator can be largely simplified. Noticing that in a gradient descent training process, the parameters are only slightly perturbed in each training cycle, thus the samplings from the previous cycles can be reused. A resample is not necessary until the current ensemble (i.e., the ensemble being optimized) deviates from the sampling ensembles significantly. Therefore, in our implementation of reweighting MBAR, we define two types of ensembles: the sampling ensembles, from which all samples are drawn (assuming there are M of them, labeled by index m = 1, 2, . . . , M ), and the target ensembles (labeled by indices p, q), which are the ensembles subject to optimization. The sampling ensembles are only updated when necessary and they do not need to be differentiable, so their data can be generated using external packages such as OpenMM. Since all samples are drawn from the sampling ensembles, the N p for the target ensembles are always zero. Consequently, in Eqn.2, the ĉp is well decoupled from ĉm and each other, and can be computed from ĉm in one shot: