Predicting Dynamical Systems across Environments via Diffusive Model Weight Generation

We consider time-varying dynamical systems, the general form of which is described by Equation 1. Given environmental conditions $e\in\mathcal{E}$, the system dynamics function is instantiated as $\frac{dx}{dt}=f(x,t,e)=f_{e}(x,t)\in\mathcal{F}$. The environment space $\mathcal{E}$ and the function space $\mathcal{F}$ are linked by the governing equations $f$, forming a joint set $\{e,f_{e}\}$. We employ a data-driven model $f_{\theta,e}$, parameterized by $\theta$, to learn $f_{e}$, thereby formalizing the function space $\mathcal{F}$ as the model’s weight space $\Theta$. The environment space is divided into an observed environment set $\mathcal{E}_{tr}$ and an unseen environment set $\mathcal{E}_{te}$, and consequently, the weight space is also partitioned into corresponding subspaces $\Theta_{tr}$ and $\Theta_{tr}$. Treating model weights as the modeling object, we learn the inherent joint distribution of environments and weights from the joint observation space $\{\mathcal{E}_{tr},\Theta_{tr}\}$. Once learning is complete, we generate a corresponding predictive function $f_{\theta,e}$ for a new environment $e\in\mathcal{E}_{te}$.

Notably, we posit that even when sharing the same governing equations, each environment possesses a unique dynamical function (though it may not be complex). We directly sample the complete model weights $\theta\sim P(\theta|e)$ for a given environment without data for finetuning, which significantly differs from existing practices in dynamics prediction.

Diffusion models [18, 19] learn a probabilistic transformation from a prior Gaussian $p_{prior}\in\mathcal{N}(\mathbf{0},\mathbf{I})$ distribution to a target distribution $p_{target}$. It perturbs data distributions by adding noise and learn to reverse this process through denoising, demonstrating strong fitting capabilities for data across modalities like images, language, and speech [20, 21]. We denote the original sample as $x_{0}$. The forward noising process in standard diffusion models is computed as $x_{n}=\sqrt{\overline{a}_{n}}x_{0}+\sqrt{1-\overline{a}_{n}}\epsilon$, where $\epsilon$ and $\{\overline{a}_{n}\}$ represent the Gaussian noise and noise schedule [22], respectively. The reverse process gradually denoises from Gaussian noise to sample data as

$$p_{\theta}(x_{n-1}|x_{n}):=\mathcal{N}(x_{n-1};\mu_{\theta}(x_{n},n),\sigma^{2}_{n}\mathbf{I}),$$

(2)

where $\mu_{\theta}=\frac{1}{\sqrt{\alpha_{n}}}(x_{n}-\frac{1-\alpha_{n}}{\sqrt{1-\overline{\alpha}_{n}}}\epsilon_{\theta}(x_{n},n))$ and $\{\sigma_{n}\}$ are step dependent constants. The noise $\epsilon_{\theta}$ is computed by a parameterized neural network, typically implemented as a UNet or Transformer architecture. The network’s parameters are optimized through an objective function [22]

$$L_{n}=\mathbb{E}_{n,\epsilon_{n},x_{0}}||\epsilon_{n}-\epsilon_{\theta}(\sqrt{\overline{\alpha}_{n}}x_{0}+\sqrt{1-\overline{\alpha}_{n}}\epsilon_{n},n)||^{2}$$

(3)

to minimize the negative log-likelihood $\mathbb{E}_{x_{0}\sim q(x_{0})}[-p_{\theta}(x_{0})]$. To model conditional distributions $p(x|c)$, state-of-the-art methods inject conditional information during noise prediction using techniques like adaptive layer normalization [23], as $\epsilon_{\theta}(x_{n},n,c)$.

EnvAd-Diff first organizes the expert small model weights into a weight graph. It then pretrains a weight VAE, yielding a high-quality latent space. Finally, an environment-conditioned diffusion model is trained within this latent space.

In the previous section, we assume the environmental condition $e$ is given. However, real-world applications often lack prior knowledge of the physical environment generating the observation trajectories. Here, we propose using a physics-informed surrogate label $c$ to differentiate observation trajectories originating from unknown environments. Within the constructed model zoo, the parameterized prediction models $f_{\theta,e}$ capture the dynamical behavior of environment $e$, thus serving as reliable physical proxies. Inspired by the concept of function distance in functional analysis, we quantify the differences between environments by computing the L2 distance between the responses of their corresponding physical proxies $f_{\theta,e}$ to the same observation samples, calculated as

$$D_{ij}=\mathbb{E}_{x_{k}\in X}||f_{\theta,i}(x_{k})-f_{\theta,j}(x_{k})||_{2}^{2}.$$

(6)

By orderly stacking the distances between each environment’s physical proxy and those of other environments, we obtain the feature vector for each environment, $D_{i}=[D_{i0},D_{i1},D_{i2},...]$. We use these feature vectors to distinguish different physical environments.

To identify the environmental condition for new environments in the test set, we apply principal component analysis, extracting the 1-dimensional principal component along the direction of the feature vectors’ largest eigenvalue. This 1D principal component serves as the surrogate label $c$, which is used to train a regression model, termed "Prompter". Prompter takes the initial frame of a trajectory as input and predicts the physical environment $c$ to which that trajectory belongs. Here we employ classic support vector regression for prompter to prevent model complexity.

We construct a model zoo to serve as the training corpus for EnvAd-Diff, adhering to three principles: 1) creating powerful expert models tailored to specific environments; 2) ensuring efficient and rapid construction; and 3) guaranteeing the stability of EnvAd-Diff training.

Neural operators [24] are employ as expert models, with minimum parameter count required for effective prediction. Prior to large-scale training of expert models for each environment, a global model is pretrained on trajectories from all available environments. This pretraining provides domain-adaptive initialization [14] for environment-specific expert training, significantly reducing the required number of training epochs. To encourage exploration of the weight landscape for different expert models within the same environment, we randomly select a layer of weights from the domain-adaptive initialization and introduce noise. We demonstrate that domain-adaptive initialization not only accelerates model zoo construction but also reduces the degrees of freedom of the weights while maintaining expert model accuracy. This significantly stabilizes the subsequent generative learning.

We assume unknown environmental conditions for all dynamical systems, training models solely on observed trajectories across diverse visible environments. At test time, models autoregressively predict future states given a single initial frame. Test environments are categorized as in-domain (seen during training, novel initial conditions) and out-domain (unseen environments). We evaluate prediction quality using root mean square error (RMSE) and structural similarity index (SSIM).

We validate the model’s effectiveness on four time-dependent PDE systems and one real-world dataset: 1) Cylinder Flow [27]; 2) Lambda-Omega [28]; 3) Kolmogorov Flow [29]; 4) Navier-Stokes Equations [7]; and 5) ERA5 Dataset [30]. For the PDE systems, we use equation coefficients or external forcing as environmental variables and simulate multiple trajectories under different environments for training and testing. We train 100 FNO weight sets for each seen environment across all systems to serve as EnvAd-Diff’s model zoo (size 100). Detailed descriptions and data generation procedures for each system are provided in Appendix A and B.

We visualize the joint distribution of weights and environments using the Cylinder Flow system as an example to aid qualitative analysis, where over 80% environments are unseen by EnvAd-Diff. In Figure 5, the x-axis represents the surrogate environment labels predicted by the prompter, and the y-axis represents the first principal component of the weights of a specific layer. The weight-environment landscape learned by EnvAd-Diff closely resembles that learned by One-per-Env through optimizer training. This indicates that EnvAd-Diff successfully models the joint distribution of weights and environments, thereby explaining its superior performance in Table 1.

To quantitatively analyze the environmental-weight joint distribution fitted by EnvAd-Diff, we introduce a simple ODE system, the Lotka-Volterra (LV) equations [7], as a toy example. We use a symbolic regression algorithm [33] to distill the predictive model generated by EnvAd-Diff for specific environmental conditions ($\beta$ and $\delta$) into an equation expression, as shown in Figure 6. The equivalent equation for the weights generated by EnvAd-Diff is consistent in form with the LV equations, and the environmental coefficients are close. This quantitatively demonstrates that EnvAd-Diff can fit the generalizable function $f$ in equation 1 rather than an environment-specific function. We detail the experimental setup in Appendix E.2.

We assess EnvAd-Diff’s robustness to data variations. Specifically, we investigate the impact of the number of seen environments and the model zoo size. We first examine the effect of model zoo size on Cylinder Flow and Lambda-Omega systems, as depicted in Figures 7a and b. The results indicate that EnvAd-Diff exhibits relatively stable performance with a zoo size of 50. As the zoo size decreases further, performance begins to deteriorate, even within the distribution. Subsequently, we test the influence of the number of seen environments on the Kolmogorov Flow and Navier-Stokes systems. The number of environments ranged from approximately 5% to 20% of the total. The findings reveal that increasing the number of seen environments reduces prediction error, but the gains become marginal after reaching around 20%. This suggests that EnvAd-Diff learns the underlying joint distribution of weights and environments from a small number of environments, rather than overfitting to trajectory samples within those environments.

The weight graph structure proposed in Section 3.1.1 is capable of organizing neural networks of arbitrary architectures. Here, we extent to more neural operators as expert models within EnvAd-Diff, including Wavelet Neural Operator [34] (WNO), and U-shape Neural Operator [35] (UNO). Our experimental results on Cylinder Flow are presented in Figure 8. EnvAd-Diff, when using different neural operators, consistently achieves excellent generalization performance, with actual performance showing only minor variations depending on the specific operator architecture. This demonstrates that EnvAd-Diff is a model-agnostic framework capable of benefiting from the sophisticated architectural designs of its expert models. Detailed architectures of these neural operators are provided in Appendix F.

We conduct an ablation study on the key components of EnvAd-Diff. Specifically, we verify the necessity of domain initialization when building the model zoo and the function loss used during VAE training. Experimental results on the Kolmogorov Flow and Navier-Stokes systems are presented in Table 4. When function loss is omitted, the VAE relies solely on MSE for reconstruction similarity, leading to suboptimal generated weights, particularly out-domain. Function loss relaxes VAE encoding constraints and helps prevent overfitting by prioritizing functional consistency over exact reconstruction. Removing domain initialization results in a significant deterioration in generated weight performance. This is attributed to the high complexity and distribution of a randomly initialized model zoo, which significantly increases the diffusion model’s modeling difficulty. We conclude that for weight generation aimed at generalization, sample quality is far more critical than diversity.

In addition, we also compare real environmental conditions with surrogate labels in Appendix E.1.

Developing dynamic prediction models with cross-environment generalization is a crucial problem in scientific machine learning and has garnered significant research interest. We review the main approaches and related work in this area. The first category trains large-parameter neural solvers as foundation models using extensive simulated data [36, 37, 38]. Subramanian et al. [5] explore the generalization performance of classical FNO architectures across different parameter sizes. Subsequently, models such as MPP [12], DPOT [11], and Poseidon [10] employed more advanced architectures to improve computational efficiency and approximation capabilities. The second approach is meta-learning [39]. These methods capture cross-environment invariants through environment-shared weights and fine-tune environment-specific weights on limited data from new environments for adaptation, including DyAd [8], LEADS [32], CoDA [7], GEPS [29], CAMEL [9], and NCF [40]. Additionally, other methods exist, like in-context learning [14]. Yang et al. [13] frame differential equation forward and inverse problems as natural language statements, pre-train transformers, and provide solution examples for new environments as context to enhance model performance. Compared to these works, we innovatively treat the complete model weights as generated objects and explicitly model their joint distribution with the environment.

Generating neural network weights is a relatively nascent research area [41]. An initial line of work involved training MLPs to overfit implicit neural fields, distilling them into model weights, and subsequently generating these MLP weights as an alternative to directly generating the fields [42]. Another category proposes using generated weights to replace hand-crafted initialization, thereby accelerating and improving the neural network training process [43, 44, 45]. These efforts primarily focus on image modalities. More recent studies leverage diffusion models to address generalization in various domains. Yuan et al. [46] employ urban knowledge graph as prompts to guide diffusion for generating spatio-temporal prediction model weights for new cities. Zhang et al. [47] replace the inner loop gradient updates of the meta learning with diffusion-generated weights. A recent work [48] explores extracting features from unseen datasets and controlling diffusion to generate adapted model weights for them. However, these methods exhibit limited zero-shot performance. This may stem from them disrupting the neural network’s inherent topological connections by directly flattening the weights, which constrains the representational capacity of the generative model. In contrast, we organize weights in the form of a neural graph and introduce a function loss to guide their representation. The zero-shot prediction performance of our generated model surpasses pre-trained models with hundreds of times more weights.