Variable-Context Conditional Diffusion for Nonlinear Dynamics Inversion

Reconstructing the initial state of a dynamical system governed by partial differential equations (PDEs) from its final state and external forcing is a fundamental challenge in inverse problems. This work reframes the task as a conditional generative problem, enabling the use of diffusion models to learn physically consistent initial conditions. A modified U-Net architecture with cross-attention layers conditions on variable-length trajectory context, allowing the model to exploit intermediate observations from the forward dynamics. Model training employs a progressively harder schedule in which fewer time slices are provided as context, encouraging the model to accurately infer initial states even from minimal information. A temporally weighted training objective further enforces consistency across the evolution trajectory. These design choices promote accurate reconstruction of the initial state while preserving fidelity to the governing dynamics. Experiments on 1D and 2D synthetic Burgers' equation datasets demonstrate that this method achieves accurate reconstructions of the initial state while maintaining consistency with the governing dynamics.