Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

In Situ Training of Implicit Neural Compressors for Scientific Simulations via Sketch-Based Regularization

Cooper Simpson, Stephen Becker, Alireza Doostan

TL;DR

This work tackles the challenge of compressing and storing data from large-scale simulations in situ by leveraging implicit neural representations (INRs) paired with a time-aware hypernetwork. A sketch-based regularization, justified via a Johnson-Lindenstrauss-type argument, mitigates catastrophic forgetting while enabling high compression rates on unstructured meshes. The authors provide theoretical bounds showing that the sketch loss serves as a surrogate to the full loss and demonstrate empirically that in situ training with FJLT sketching closely matches offline performance across multiple datasets, with a memory-efficient buffer design (full snapshots $T_f=1$ and sketch snapshots $T_s=T-1$). They further discuss extensions to more sketching strategies, adaptive buffering, and physics-informed objectives, highlighting practical pathways to deploy neural compression in distributed, real-time simulations.

Abstract

Focusing on implicit neural representations, we present a novel in situ training protocol that employs limited memory buffers of full and sketched data samples, where the sketched data are leveraged to prevent catastrophic forgetting. The theoretical motivation for our use of sketching as a regularizer is presented via a simple Johnson-Lindenstrauss-informed result. While our methods may be of wider interest in the field of continual learning, we specifically target in situ neural compression using implicit neural representation-based hypernetworks. We evaluate our method on a variety of complex simulation data in two and three dimensions, over long time horizons, and across unstructured grids and non-Cartesian geometries. On these tasks, we show strong reconstruction performance at high compression rates. Most importantly, we demonstrate that sketching enables the presented in situ scheme to approximately match the performance of the equivalent offline method.

In Situ Training of Implicit Neural Compressors for Scientific Simulations via Sketch-Based Regularization

TL;DR

This work tackles the challenge of compressing and storing data from large-scale simulations in situ by leveraging implicit neural representations (INRs) paired with a time-aware hypernetwork. A sketch-based regularization, justified via a Johnson-Lindenstrauss-type argument, mitigates catastrophic forgetting while enabling high compression rates on unstructured meshes. The authors provide theoretical bounds showing that the sketch loss serves as a surrogate to the full loss and demonstrate empirically that in situ training with FJLT sketching closely matches offline performance across multiple datasets, with a memory-efficient buffer design (full snapshots and sketch snapshots ). They further discuss extensions to more sketching strategies, adaptive buffering, and physics-informed objectives, highlighting practical pathways to deploy neural compression in distributed, real-time simulations.

Abstract

Focusing on implicit neural representations, we present a novel in situ training protocol that employs limited memory buffers of full and sketched data samples, where the sketched data are leveraged to prevent catastrophic forgetting. The theoretical motivation for our use of sketching as a regularizer is presented via a simple Johnson-Lindenstrauss-informed result. While our methods may be of wider interest in the field of continual learning, we specifically target in situ neural compression using implicit neural representation-based hypernetworks. We evaluate our method on a variety of complex simulation data in two and three dimensions, over long time horizons, and across unstructured grids and non-Cartesian geometries. On these tasks, we show strong reconstruction performance at high compression rates. Most importantly, we demonstrate that sketching enables the presented in situ scheme to approximately match the performance of the equivalent offline method.

Paper Structure

This paper contains 14 sections, 4 theorems, 13 equations, 8 figures, 4 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries and Related Work
  3. Implicit neural representation (INR)
  4. Hypernetworks
  5. Non-stationary learning
  6. Compression
  7. In Situ Training
  8. Sketching
  9. JL Regularization for In Situ Learning
  10. Experiments
  11. Results
  12. Discussion and Future Work
  13. Sketching for In Situ Learning
  14. Neural Compression

Key Result

Theorem 1

For any $\epsilon>0$ and $n\in \mathbb{N}$, there exists a probability distribution $\mathcal{D}$ with support in $\mathbb{R}^{k\times n}$, where $k\leq \min\{n,\mathcal{O}(\epsilon^{-2}\log(m))\}$, such that for any set of $m$ points $\{\bm{p}\}_{i=1}^m\subset\mathbb{R}^n$, there exists $1>\delta>0 In particular, for sufficiently large dimensions, using the probabilistic method, we see that there

Figures (8)

  • Figure 1: The compression approach of this study applies to all mesh/geometry types, including (uniform) structured (left), unstructured (middle), and curvilinear (right), as it uses the node locations without connectivity information.
  • Figure 2: Compression results for INRs trained offline. See \ref{['sec:experiments']} for details on the datasets.
  • Figure 3: Our in situ INC training approach relies on a sketched data buffer to avoid catastrophic forgetting.
  • Figure 4: INC network structure.
  • Figure 5: Reconstruction comparison of the Ignition data at snapshot $t=13$. The absolute error in the cutout is computed with respect to the original image, which is presented for reference. The colorbar only references the error cutouts.
  • ...and 3 more figures

Theorems & Definitions (6)

  • Theorem 1: Johnson-Lindenstrauss Transformations johnson1984extensions
  • Definition 1: JL Transform
  • Theorem 2: Manifold JL Transform baraniuk2009random
  • Theorem 3: JL Surrogate
  • proof
  • Theorem 4: Batch JL Surrogate