Predicting partially observable dynamical systems via diffusion models with a multiscale inference scheme

We thank the reviewer for their feedback, especially regarding the motivation behind our approach and the types of physical systems we address, which has greatly helped improve the paper. Below, we respond to the main points concerning the motivation and applicability of our method.

Question 1. Motivation of our multiscale templates

Thanks to the reviewer’s point, we would like to clarify the following. The paper tackles probabilistic forecasting of a dynamical system, which has two main properties.

• Partial observability. The state of the system at any given time cannot be fully determined by the observations. Consequently, the distribution of future scenarios conditioned on past observed data may not be restricted to a Dirac measure. It is this point that justifies the task of probabilistic forecasting, rather than deterministic forecasting.

• Long memory. The system has slowly varying components that are important to model (e.g., the motion of plasma across several millions of kilometers on the Sun).

Long memory alone is not sufficient to justify our method. For example, consider the synthetic case of Fig. 2 in our paper. It exhibits long memory, since it has a slowly varying sinusoidal component. If the process had zero stochasticity (making it fully observable), then all trajectories would exactly follow this smooth component, and a single state would be sufficient to determine the position on the sinusoid and what comes next. In that case, there would be no need for our multiscale templates. Now, if the system exhibits stochasticity, imagine we are at the negative peak (around $t=30$). Depending on the noise realization, the local trend may be positive or negative, making the state difficult to recover locally. In other words, the system’s partial observability, due to noise in the data, prevents accurate estimation of the slowly varying component. In this case, it is thus necessary to look further into the past, which is precisely what our multiscale inference scheme achieves, at no additional computational cost.

This is supported by the experimental results reported below: as the noise level increases, the performance gain (measured as the ratio of Wasserstein distances; see Fig. 2) of our multiscale template becomes more pronounced. A ratio greater than 1 indicates improved performance with our model.

Following your comments, we would like to clarify that the “stochasticity” mentioned in the paper (line 164), such as measurement noise, can be considered part of the partial observability of the system. This has now been reflected in the abstract (line 9), which reads: “In this paper, we tackle the probabilistic prediction of partially observable dynamical systems with long-memory.”

Question 2. Design principle

The method is not specific to the Sun and accommodates a wide range of timescales, as shown in Fig. 4 (Appendix). The general algorithm is described in lines 227–242 in the main text. To improve clarity and highlight the generality of our approach, we will include an algorithm box in the paper.

The multiscale templates are derived from first principles and are not hand-designed. The only dataset-specific adaptation lies in two hyperparameters: the time horizon that is generated at once (see Fig. 4, Appendix) and the number of future steps predicted per diffusion sampling process ($K$ in the paper). The later is driven by the computational budget allocated to the model (the more steps to generate at once, the larger the denoiser). The former depends on the time range of the dependencies. For example, if the system exhibits a maximum timescale $T$, after which any two observations $x_{t},x_{t+T}$ are independent, then this is a natural candidate for the horizon parameter. Otherwise, this can be tuned, and also provides a way to reveal the horizon of long-range dependencies in a dynamical system.

To showcase the genereticity of our multiscale inference scheme, we applied it to the Navier-Stokes data from PDEArena [Gupta and Brandstetter 2022]. To make the dynamical system partially observable, we downsampled it to 32x32 (from 128x128), considered only the density field (discarding velocity) and added a small Gaussian noise on the observations. The results on this new example of partially observable process are presented in the Table below with the same metric than in the paper (MAE on the power-spectrum, see line 296).

We observe that our multiscale inference scheme improves predictions over the autoregressive baseline, particularly in the long-term interval (16–36), where it achieves a significantly lower MAE on the power spectrum. This shows that our approach is not specific to solar prediction. Rather, we considered solar forecasting as a challenging scientific application where our model also performs well.

We thank the reviewer for their insightful comments and suggestions, and for highlighting both the clarity of the manuscript and the value of releasing the solar dataset to the community.

Below are our answers to your questions.

Question 1

Yes, we can extend to the case where the observations are available as functions of the state without a lot of effort, so long as the function in question is differentiable. Indeed, in that case, one can learn a diffusion model on the states themselves through an expectation-maximization algorithm for example [Rozet et al., 2024]. Then, one can apply our multiscale inference to the diffusion model in the same way as in our paper. The expectation-maximization algorithm and the inference scheme are essentially orthogonal.

Question 2

While interesting, the tasks introduced in the Long Range Arena benchmark [Tay, Dehghani et al., 2020] focus on compositionality, document similarity, and classification, and thus fall outside the scope of our work, which addresses probabilistic forecasting of partially observable dynamical systems. Since we agree that validating our method beyond the solar problem is important, we applied it to the PDEArena dataset [Gupta et al., 2022], specifically to the Navier-Stokes trajectories (see their Fig. 1). To make this system partially observable, we downsampled it from 128×128 to 32×32, retained only the density field (discarding velocity), and added small Gaussian noise to the observations. Results for this new partially observable system are reported in the table below, using the same metric as in our paper (MAE on the power spectrum, see line 296).

We observe that our multiscale inference scheme improves predictions over the autoregressive baseline, particularly in the long-term interval (16–36), where it achieves a significantly lower MAE on the power spectrum. This shows that our approach is not specific to solar prediction. Rather, we considered solar forecasting as a challenging scientific application where our model also performs well.

Question 3

Thank you for the insightful comment. We now include a discussion of state-space models in the "Related Works" section. We would like to emphasize that the focus of our paper is not on the choice of architecture, but rather on how to perform prediction within a conditional diffusion model. Diffusion models are increasingly being applied to dynamical systems [63, 67, 73, 34], and have demonstrated strong performance in generative modeling, which aligns well with our task: probabilistic forecasting.

Our paper does focus on the solar prediction problem (although we also include a synthetic example and will add the results on PDEArena), because it is currently a very challenging task in probabilistic prediction, with the potential to significantly advance solar flare prediction in the future [41, 10, 43].

We thank the reviewer for their positive feedback on our work, particularly for recognizing that our method is "breaking the limitations of traditional autoregressive rolling prediction" and for highlighting its "possibility of easy deployment for real applications".

Question 1. Part 1.

We are not directly choosing the template {-5, -4, -2, 0, 2, 4, 5} because it is not, strictly speaking, a multiscale template in the sense of Definition Eq. 7 in our paper. This template is refined near $t = -5$ and $t = 5$, but not around $t = 0$, whereas the motivation behind our multiscale templates is to provide finer resolution around $t = 0$ and become progressively coarser farther from it (see line 199). Moreover, if we were to first use the template {-9, -3, -1, 0, 1, 3, 9} and then switch to {-5, -4, -2, 0, 2, 4, 5} as suggested, this would mean that future steps {1, 3, 9} are generated before steps {2, 4, 5}, potentially introducing temporal discontinuities. For instance, the model could generate a scenario where {1, 3, 9} become more active, while {2, 4, 5} become less active — even though both sets are conditioned on the same past steps {-5, -4, -2, 0}. This is precisely why our multiscale scheme (see Fig. 1, right) is designed to always condition on already generated future steps, thereby preserving temporal coherence.

Question 1. Part 2.

Indeed, the second multiscale template {-6,-2,-1,0,1,2,6} should be shifted by 3 (not by 2) to obtain {-3,1,2,3,4,5,9}, which overlaps with the steps {-3,1,3}. These correspond to either past steps {-3} or already generated future steps {1,3}. This has been updated in the manuscript (lines 234–235).

Question 2

Handling unevenly sampled trajectories is primarily achieved by concatenating time information to the input of our denoiser, as described on line 289 of the paper. More precisely, we add an additional channel to the input tensor that indicates, for each pixel in the input trajectory, its relative time index (e.g., -9, -3, -1, 0, 1, 3, 9). Further details are provided in Appendix (lines 99–104)

Question 3

Table 1 in our paper reports the results of a non-diffusion-based architecture, specifically a recent transformer architecture designed for next token prediction [52], that has been consistently used for the prediction of physical dynamical systems. Below are the results for such baseline, as reported in our paper. Note that published models on solar prediction [1,66] operate in substantially worse spatiotemporal resolutions (e.g., a time resolution of 12h while we predict hourly, or 2x spatially downsampled compared to our data) than this baseline.

Note that a transformer, as any deterministic baseline, will suffer from the well-known smoothing effect, as shown below. The transformer model produces a prediction with significantly less power in the high-frequencies than the observed spectral content (measured as the sum of the power-spectrum (PS) over the top 75% high frequencies).

This smoothing behavior would occur with any deterministic baseline, and we preferred to focus our paper on comparing different schemes within the same diffusion model.

Thank you for your response. Your explanation has addressed most of my concerns, and I now have a better understanding of your motivation and the rationale behind your template choice. After careful consideration, I think this paper is above the acceptance bar of NeurIPS, and I'm willing to increase my rating from 4 to 5. Good luck.

We thank the reviewer for their positive feedback, in particular for acknowledging the strength of the application, the readability of our paper, and for the valuable insights on data assimilation (DA).

Connections with related prediction approaches in dynamical systems

Although we follow developments in data assimilation (DA) with interest, our paper does not directly address this topic. Instead, it focuses on probabilistic prediction of the observations themselves, without attempting to reconstruct the system's underlying states. We updated the “Related Works” section based on your comment about data assimilation (DA).

• Line 68 now reads: “our work falls within the scope of modeling the dynamics of the observations”.
• Line 76 now reads: “Thus, while [40, 63] see no benefit from using more than two past observations, incorporating additional past steps substantially improves results in our setting”.
• Line 79 now reads: “Diffusion models can perform data assimilation and prediction from incomplete observations simultaneously [67, 73, 34], but this requires a dataset of underlying system states to train the model -- an assumption we do not make in this paper”, where the last part of the sentence refers to [67, 73, 34].

In our paper, partial observability is defined by the fact that, at any time, the conditional distribution of the state given its observation is not restricted to a Dirac measure, characterizing the impossibility of reconstructing the system’s state from the observation. This will be stated more clearly in the introduction and in section 4.1. In that sense, we align with the standard mathematical framework used in Bayesian filtering and DA, without explicitly writing out the equations, since, as you mention, we aimed to avoid any “mathiness” not directly useful for developing our approach. In the solar prediction application tackled in our paper, there are substantial unknown latent physical variables (e.g., all state variables of the interior, the magnetic field in the Sun's atmosphere) as well as open questions about dynamics (e.g., how the atmosphere of the Sun is actually heated) and entirely missing subgrid physics that cannot be observed in our data.

Clarifications needed regarding assumption (b) in Sec 4.1

We thank the reviewer for pointing out this sentence. It has been replaced by:
" Diffusion models have been applied to prediction on dynamical systems by using additional methods to remove assumption (b). For example, [40] apply a diffusion model to fully resolved fluids which are effectively Markovian. In weather prediction, although the observed data are sparse, data assimilation -- also known as reanalysis -- enables the reconstruction of missing information, resulting in large datasets of highly informative states [28], on which diffusion models have been successfully trained [63]. Other models handle missing states, but require clean sates for training [73,34], which is not always available.

In this paper, we tackle the challenging problem of predicting a dynamical system that presents the three difficulties (a), (b), and (c) simultaneously, without assuming that (b) can be circumvented, as is often the case in many disciplines. "

Other clarifications

As you point out, the system has to be non-first-order Markov for our method to show an advantage over autoregressive systems. We did not mention this assumption as we allow ourselves to test it empirically. In practice, if a standard autoregressive scheme outperforms our multiscale approach, this may suggest that the process is Markovian of some order.

Indeed, it is important to clarify that we learn the score of each of the distribution $p(x_{I_n}| x_{C_n})$ in Eq.6 in the same way as [Harvey et al., 2022] did, by training a denoiser which takes as input the relative positions of the time steps in the context $C_n \cup I_n$ (see line 289-290 and Appendix, line 99-104).

Section 3 will be renamed "Background: Conditional Diffusion Models".

Questions

As stated above, we learn the score of the distributions $p(x_{I_n}| x_{C_n})$, which can be used to approximate the distribution $p(x_{1:T} | x_{t\leq0})$ (see Eq. 6). Let us recall that we assume the set of indices $I_n$ and $C_n$ to have $K$ and $K+1$ elements respectively.

In certain restricted cases, one can prove that Eq. 6 is actually an equality, for example, under a Markovian assumption of order $K+1$, and employing a standard autoregressive inference scheme (see line 155 in our paper). In that case, it makes sense to model the target distribution $p(x_{1:T} | x_{t\leq0})$ by learning the score of $p(x_{I_n}| x_{C_n})$ for all $n$.

Studying this approximation in the case of generic $I_n,C_n$ is hard, and it is not clear that a Markovianity assumption of order $K+1$ is well suited for justifying the choice of the maximum $\alpha$ in our inference scheme. At this stage, a proper explanation would require a separate LaTeX note, which is not permitted this year.

Intuitively, if we assume that the process has a maximum timescale of dependence, that is, there exists a time $T$ such that for any time $\tau\geq T$, the variables $x_t$ and $x_{t+\tau}$ are independent (often referred to as the integral scale in signal processing), then it is natural to restrict the extent of the largest multiscale template $\mathbb{T}^{\alpha_{max}}_K$ (defined in Eq. 7) to be smaller than $T$, since there is no need to model dependencies beyond $T$. However, this intuition would need to be made more rigorous.