Diffusion models for robust full-field prediction of soliton propagation

The propagation of ultrashort optical pulses in fibres is governed by the Nonlinear Schrödinger Equation (NLSE), whose numerical integration remains computationally expensive for high-dimensional applications. Traditional solvers such as the Split Step Fourier Method are accurate but limited by step-size constraints and scaling complexity. In response to these drawbacks, machine learning methods have emerged as promising alternatives to act as surrogate propagators. Specifically, deep neural networks offer robust capabilities for learning patterns and predicting solutions of nonlinear differential equations. Despite their performance, these methods are inherently limited due to sensitivity to noise and error accumulation over long sequence prediction tasks. In this work, we present a generative surrogate model for the NLSE based on denoising diffusion probabilistic models (DDPMs). Diffusion models have recently shown remarkable capability in learning complex physical distributions and reconstructing missing information from partial or noisy data. Our framework learns the spatiotemporal evolution of the optical field from simulated NLSE datasets, predicting the full-field intensity and phase from partially masked intensity inputs. Using a U-Net backbone within the diffusion process, the model generates physically consistent propagation maps across varying soliton numbers, capturing the nonlinear interplay between dispersion and self-phase modulation. We benchmark the DDPM against a deterministic U-Net trained under equivalent conditions and show superior generalisation to out-of-distribution nonlinear regimes and masked inputs. Quantitative evaluation using standard metrics confirms that the diffusion-based surrogate accurately reproduces the nonlinear dynamics of ultrashort pulse propagation. This work highlights the potential of diffusion models as fast, datadriven solvers for complex photonic partial-differential equations and paves the way to the application of generative models to the prediction of nonlinear optical phenomena.