Overcoming Diffusion Underestimation in Generative Molecular Dynamics of TiO2-SiO2 Membranes via Displacement-Weighted Fine-Tuning

Flow Matching-based generative models can be used to generate molecular dynamics (MD) trajectories over long time scales. However, in systems where rare hopping events dominate transport, standard models underestimate the long-time mean squared displacement (MSD) and diffusion coefficient. For $\mathrm{H_2}$ permeation in an amorphous $\mathrm{TiO_2\text{-}SiO_2}$ membrane, local small displacements are frequent, whereas large displacements associated with hopping are rare. Consequently, a model trained with uniform weighting mainly reproduces frequent small displacements and underrepresents rare hopping events. We propose displacement-weighted fine-tuning, which changes the training loss according to the ratio between the one-step displacement distributions of target H atoms obtained from the MD trajectory and from the baseline model. For 10 ns generated trajectories, fine-tuning increases the diffusion coefficient from $1.58 \times 10^{-4}\,\mathrm{cm^2/s}$ to $3.36 \times 10^{-4}\,\mathrm{cm^2/s}$, a 2.1-fold increase, bringing it closer to the MD value of $5.06 \times 10^{-4}\,\mathrm{cm^2/s}$. Fine-tuning also reduces the underestimation of large-displacement events. These results show that emphasizing rare large displacements during training improves the long-time transport behavior of $\mathrm{H_2}$ in the membrane.