Sparse Diffusion Autoencoder for Test-time Adapting Prediction of Complex Systems

TL;DR

This work designs a codebook-based sparse encoder, which coarsens the continuous spatial domain into a sparse graph topology and employs a graph neural ordinary differential equation to model the dynamics and guide a diffusion decoder for reconstruction.

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.

Figures (11)

• Figure 1: Adapting probe topologies.
• Figure 2: Overall framework of proposed Sparse Diffusion Autoencoder.
• Figure 3: Real-world dataset. (a) Prediction visualization of different models. (b) RMSE comparison of different models.
• Figure 4: Ablation studies. Prediction performance with (a) uniform and (b) random probe selection. (c) Impact of probe topology edge weights on prediction.
• Figure 5: Robustness experiments. (a) Impact of codebook size on SparseDiff's performance on the Navier-Stokes system. (b) Impact of the noise on SparseDiff's performance on the real climate dataset.