Zhong Yi Wan
Zhong Yi Wan is a research scientist at Google Research. He obtained his PhD degree in computational engineering from MIT at 2020. His research specialization is dynamical system and generative modeling for scientific applications.
Research Areas
Authored Publications
Sort By
DySLIM: Dynamics Stable Learning by Invariant Measure for Chaotic Systems
Yair Schiff
Jeff Parker
Volodymyr Kuleshov
International Conference on Machine Learning (ICML) (2024)
Preview abstract
Learning dynamics from dissipative chaotic systems is notoriously difficult due to their inherent instability, as formalized by their positive Lyapunov exponents, which exponentially amplify errors in the learned dynamics. However, many of these systems exhibit ergodicity and an attractor: a compact and highly complex manifold, to which trajectories converge in finite-time, that supports an invariant measure, i.e., a probability distribution that is invariant under the action of the dynamics, which dictates the long-term statistical behavior of the system. In this work, we leverage this structure to propose a new framework that targets learning the invariant measure as well as the dynamics, in contrast with typical methods that only target the misfit between trajectories, which often leads to divergence as the trajectories’ length increases. We use our framework to propose a tractable and sample efficient objective that can be used with any existing learning objectives. Our Dynamics Stable Learning by Invariant Measure (DySLIM) objective enables model training that achieves better point-wise tracking and long-term statistical accuracy relative to other learning objectives. By targeting the distribution with a scalable regularization term, we hope that this approach can be extended to more complex systems exhibiting slowly-variant distributions, such as weather and climate models. Code to reproduce our experiments is available here: https://github.com/google-research/swirl-dynamics/tree/main/swirl_dynamics/projects/ergodic.
View details
Evolve Smoothly, Fit Consistently: Learning Smooth Latent Dynamics For Advection-Dominated Systems
Anudhyan Boral
International Conference on Learning Representations (ICLR) (2023)
Preview abstract
We present a data-driven, space-time continuous framework to learn surrogate models for complex physical systems described by advection-dominated partial differential equations. Those systems have slow-decaying Kolmogorov n-width that hinders standard methods, including reduced order modeling, from producing high-fidelity simulations at low cost. In this work, we construct hypernetwork-based latent dynamical models directly on the parameter space of a compact representation network. We leverage the expressive power of the network and a specially designed consistency-inducing regularization to obtain latent trajectories that are both low-dimensional and smooth. These properties render our surrogate models highly efficient at inference time. We show the efficacy of our framework by learning models that generate accurate multi-step rollout predictions at much faster inference speed compared to competitors, for several challenging examples.
View details
Debias Coarsely, Sample Conditionally: Statistical Downscaling through Optimal Transport and Probabilistic Diffusion Models
Ricardo Baptista
Yi-fan Chen
John Anderson
Anudhyan Boral
Advances in Neural Information Processing Systems (NeurIPS) 36 (2023)
Preview abstract
We introduce a two-stage probabilistic framework for statistical downscaling between unpaired data. Statistical downscaling seeks a probabilistic map to transform low-resolution data from a (possibly biased) coarse-grained numerical scheme to high-resolution data that is consistent with a high-fidelity scheme. Our framework tackles the problem by tandeming two transformations: a de-biasing step that is performed by an optimal transport map, and a super-resolution step that is achieved via a probabilistic diffusion model with a posteriori conditional sampling. This approach characterizes a conditional distribution without the need for paired data, and faithfully recovers relevant physical statistics from biased samples.
We demonstrate the utility of the proposed approach on one- and two-dimensional fluid flow problems; they are representative of the core difficulties present in numerical simulations of weather and climate. Our method produces realistic high-resolution outputs from low-resolution inputs, by upsampling resolutions of 8x and 16x. Moreover, the procedure is faithful to the correct statistics of physical quantities, even when the low-frequency energy profiles of the inputs and the desired outputs do not match, a crucial but difficult-to-satisfy assumption by current state-of-the-art alternatives.
View details
Neural Ideal Large Eddy Simulation: Modeling Turbulence with Neural Stochastic Differential Equations
Anudhyan Boral
James Lottes
Yi-fan Chen
John Anderson
Advances in Neural Information Processing Systems (NeurIPS) 36 (2023)
Preview abstract
We introduce a data-driven learning framework that assimilates two powerful ideas: ideal large-eddy-simulation (LES) from turbulence closure modeling and neural stochastic differential equations (SDE) for modeling stochastic dynamical systems. ideal LES identifies the optimal reduced-order flow fields of the large-scale features by marginalizing out the effect of small-scales in stochastic turbulent trajectories. However, ideal LES is analytically intractable. In our work, we use a latent neural SDE to model the evolution of the stochastic process and a pair of encoder and decoder for transforming between the latent space and the desired optimal flow field. This stands in sharp contrast to other types of neural parameterization of the closure models where each trajectory is treated as a deterministic realization of the dynamics. We show the effectiveness of our approach (niLES – neural ideal LES) on a challenging chaotic dynamical systems: Kolmogorov flow at a Reynolds number of 20,000. Compared to prior works, our method is also able to handle non-uniform geometries and unstructured meshes. In particular, niLES leads to more accurate long term statistics, and is stable even when rolling out to long horizons.
View details