Title: Sparse Diffusion Autoencoder for Test-time Adapting Prediction of Complex Systems
Venue: NeurIPS 2025
Accepted: yes
Keywords: Complex system

Diffusion
Abstract: Predicting the behavior of complex systems is critical in many scientific and engineering domains, and hinges on the model’s ability to capture their underlying dynamics.
Existing methods encode the intrinsic dynamics of high-dimensional observations through latent representations and predict autoregressively. 
However, these latent representations lose the inherent spatial structure of spatiotemporal dynamics, leading to the predictor's inability to effectively model spatial interactions and neglect emerging dynamics during long-term prediction.
In this work, we propose SparseDiff, introducing a test-time adaptation strategy to dynamically update the encoding scheme to accommodate emergent spatiotemporal structures during the long-term evolution of the system. 
Specifically, we first design a codebook-based sparse encoder, which coarsens the continuous spatial domain into a sparse graph topology. Then, we employ a graph neural ordinary differential equation to model the dynamics and guide a diffusion decoder for reconstruction.
SparseDiff autoregressively predicts the spatiotemporal evolution and adjust the sparse topological structure to adapt to emergent spatiotemporal patterns by adaptive re-encoding. 
Extensive evaluations on representative systems demonstrate that SparseDiff achieves an average prediction error reduction of 49.99\% compared to baselines, requiring only 1\% of the spatial resolution.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper “Sparse Diffusion Autoencoder for Test-time Adapting Prediction of Complex Systems” introduces SparseDiff, a framework for long-term forecasting of high-dimensional spatiotemporal dynamics (e.g., fluid flow, climate fields). SparseDiff compresses each high-dimensional frame into a handful of learned “probes” with a codebook-based sparse encoder, advances those probes through an edge-aware graph neural ODE, and then uses a guided diffusion decoder to regenerate the full field, thus preserving spatial structure while working in an extremely low-dimensional space.  Evaluated on PDE benchmarks and the SEVIR weather dataset, this test-time adaptation cuts long-range RMSE by roughly 50 % versus state-of-the-art baselines while requiring fewer than 1 % of grid points and remaining robust to noise and out-of-distribution Reynolds numbers.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. Robustness and generalization: maintains performance under noisy inputs and extrapolates to unseen Reynolds numbers without retraining. \n\n2. Test-time adaptive encoding automatically refreshes the probe graph during inference, preventing drift in very long forecasts and preserving accuracy. \n\n3. State-of-the-art long-horizon accuracy, reducing RMSE by about 50 % versus FNO, UNet and other strong baselines on four PDEs and a weather dataset.\n\nWeaknesses:\n\n1.  The introduction points out “a lack of theoretical understanding for aggregating rich spatiotemporal dynamics from a continuous domain into a small set of discrete probes”; yet the paper itself provides no such theoretical analysis—most notably, it does not quantify how prediction accuracy scales with the forecast horizon t, leaving readers without guidance on the temporal range over which the method’s outputs remain reliable.\n\n2. The Methods section introduces a “Sparse Encoder,” yet it never defines what “sparse” specifically refers to—is it merely the step of replacing the latent vector z with its nearest codeword c before decoding? In addition, the paper offers no theoretical justification for why sparsity is needed or how it benefits the model.\n\n3. In the PDE context $\nh_t = h_0 + \\int_{0}^{t} A\\,h_t \\, dt,\n$, we also care about estimating and inferring the parameter A, as it is crucial for stability analysis and control design. Could your method be extended to provide reliable inference of A, rather than focusing solely on long-term state prediction?\n\nQUESTIONS:\nPlease see the “Weaknesses” section. If the authors can supply adequate theoretical justification, I would be willing to consider raising my score.\n\nLIMITATIONS:\nPlease see the “Weaknesses” section.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents SparseDiff, a novel autoencoder framework for predicting complex spatiotemporal dynamics. The model first encodes the system's state into a sparse probe topology using a learned codebook of dynamic patterns. It then employs a graph neural ODE to model the dynamics on this graph and guides a diffusion decoder for reconstruction. A test-time adaptation strategy is adopted to dynamically update the probe topology, ensuring the model remains accurate during long-term prediction. The authors validate the accuracy and efficiency of SparseDiff on various PDE systems and real-world datasets.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- novel idea of test-time adapting prediction\n- Strong performance over baselines across multiple PDE systems and a real-world dataset.\n- Ablation studies on probe selection and edge weighting\n\n\n**Weaknesses**\n\nsee questions\n\nQUESTIONS:\n1. The current framework is limited to 2D regular spatial domain. Many real-world systems, such as those in climate science or sensor networks, involve data sampled from irregular or sparse spatial domains. Can you comment on the applicability of SparseDiff to such irregular or sparse data? \n\n2. l. 143 $v_i \\in R^{t\\times h\\times w}$. l.38  $v_i \\in R^{h\\times w}$. There seems to be a notational inconsistency. Can you further clarify the precise algorithm for computing the probe features? \n\n3. How were the two hyperparameters (number of codewords $K$, and the threshold $\\tau$ ) selected in practice for the experiments? Can you provide some insight into the influence of the hyperparameters on the performance? \n\n4. It is unclear to me why the test-time adapting prediction improves the performance. Can you provide a quantitative explanation for why this re-encoding step leads to a better-aligned probe topology and improved performance? It will be good to do an ablation study that isolates the effect of adaptation. To do this, can you compare the following two models on the Navier-Stokes experiment and report the result as in Table 1? \n    - SparseDiff without test-time adaptation\n    - G-LED with test-time adaptation\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes SparseDiff, a Sparse Diffusion Autoencoder for efficient long-term prediction in spatiotemporal systems. It encodes global dynamics via a sparse set of probes using a codebook-based encoder, models them with a graph neural ODE, and reconstructs the full state with a diffusion decoder. The probe topology adapts at test time, enabling dynamic tracking. Experiments on PDE benchmarks and a climate dataset show superior accuracy and efficiency compared to state-of-the-art baselines.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n1. The manuscript introduces a practical strategy for updating probe selections and topologies dynamically at test time via codebook re-encoding. This design is argued to provide an advantage for capturing emergent spatiotemporal dynamics compared to conventional static encoders and predictors.\n2. The experiments are thorough and comprehensive. On PDE benchmarks and the SEVIR climate dataset, SparseDiff reduces prediction error by an average of 49.99% compared to leading baselines (Table 1), while only using about 1% of the input spatial grid. Such results suggest the framework can maintain accuracy with substantial efficiency gains. Moreover, they conduct ablation studies  to emphasize the necessity of the codebook-driven probe encoder and edge-aware attention.\n\n### Weaknesses\n1. Notation and method description are occasionally ambiguous. For example, Section 3.1.1 is overly difficult to follow, with dense explanations and somewhat confusing notation (though not incorrect). I recommend revising this section for clarity—streamlining the writing and organizing the symbols more clearly would greatly improve readability.\n2. The edge weighting scheme for probe topology (Section 3.1.1) is domain-agnostic and not learned end-to-end, as acknowledged under limitations. While the ablation (Figure 6c) shows this is better than unweighted edges, it leaves open whether optimal graph connectivity could be learned or data-driven; alternatively, comparisons to other heuristic graph construction methods (e.g., k-nearest neighbors, spectral approaches) are not included.\n\nQUESTIONS:\n1. Is the topological structure present in every frame? And is it just the probe states that incorporate information from all previous frames up to the current one? If so, it would be helpful to clarify this distinction in the paper, as it’s not immediately clear whether the topology itself evolves over time or remains static while only the probe representations are temporally informed.\n2. Why is the dimension of $v_i$ given as $h \\times w$ ? Typically, one would expect $v_i$ to represent a vector or feature at a probe, which would usually have a fixed channel or embedding dimension. If $h\\times w$ refers to spatial dimensions, this seems inconsistent with the idea of sparse probe representations. Could the authors clarify the meaning and rationale behind this shape?\n3. In the first equation of Section 3.2, it’s recommended not to use $t$ as the integration variable, as it may cause confusion with the timestep index. Additionally, it's unclear whether $h_t$ and $h_i$ share the same meaning or dimensionality—does $h_t$ refer to the representation at time $t$, while $h_i$ denotes the representation of the $i$-th probe? Clarifying this distinction would help avoid ambiguity in notation.\n4. In line 177, should the superscript be $N$ instead? \n5. Are the Codebook Encoder and Diffusion Decoder trained separately, or is the entire model trained end-to-end?\n\nLIMITATIONS:\nSee Weakness", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes SparseDiff, a novel framework for long-term prediction of complex spatiotemporal systems. SparseDiff uses a codebook-based sparse encoder to construct a “probe topology” capturing low-dimensional spatial structures, and a guided diffusion model to reconstruct full-domain states. The key innovation lies in the test-time adaptation mechanism that dynamically updates the encoding based on the current system state. The approach is validated on both synthetic PDE systems and real-world climate data, with strong quantitative and qualitative performance.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n- The Test-Time Adaptation  mechanism based on latent similarity offers a practical solution for capturing emerging spatiotemporal structures, which help improve long-term prediction accuracy and inference efficiency.\n\n- Extensive experiments on diverse PDE benchmarks and real-world data demonstrate clear advantages over baselines, including FNO and G-LED.\n\n- The method that uses sparse probes to construct a reduced-order model is very interesting. However, I'm not very familiar with this field, so I'm not sure about its novelty.\n\n## Weeknesses :\n\n- The number of trainable parameters in SparseDiff is larger than that of all other baselines, and the training efficiency compared to the other methods is not clearly discussed.\n\n- The experimental configurations for the 'Robustness,' 'Generalization,' and 'Trade-off between Accuracy & Efficiency' sections are not consistent. For example, the Robustness experiment compares against G-LED on the Navier-Stokes system, while the Generalization experiment compares against FNO on the Cylinder Flow system. Could you clarify the rationale behind these choices, or consider adding experiments on additional PDE systems with more consistent baseline comparisons?\n\nQUESTIONS:\n-\tIn Section 3.1.2, you highlight the differences between your method and previous works [14, 13, 12], but these methods are not included in the experimental comparisons. Could you clarify why they were omitted?\n\n-\tHow is the value of k, described in Section 3.1.1, chosen, and how does it affect the final performance and effiency?\n\n-\tThe trainable parameter of SparseDiff are larger than all of other baseline, could you also compare  training GPU hours for SparseDiff and other baselines?\n\n\n## Typos:\n\n-\tIn line 101, \\epolsion should be \\epolsion_n? Keep consistent with Equation (2), where the subscript denotes the specific noise term at step n.\n\n-\tFigure 7(b) lacks explanation; it is unclear what the term 'Rollout step' refers to.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}