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

AI summary

Overview

• Research area: Scientific machine learning / neural compression of simulation data, spanning implicit neural representations (INRs), hypernetworks, randomized sketching, and continual (non-stationary) learning.
• Technical level: Advanced. The paper combines a practical training protocol with Johnson–Lindenstrauss theory, manifold JL transforms, and randomized numerical linear algebra.
• Scope: The paper proposes an in situ training scheme for INR-based neural compressors that uses a small buffer of full snapshots plus sketched snapshots as a regularizer against catastrophic forgetting, supported by a theoretical JL surrogate result and evaluated on 2D and 3D simulation data. Note that the supplied content is truncated before Section 3, so most of the reported experimental numbers are not available here; where that is the case, this summary says so explicitly.

What This Paper Is About

Large-scale scientific simulations produce datasets far too big to store offline, yet analysts still want access to them, so compression must happen while the simulation is running in situ. Neural compressors based on implicit neural representations can achieve this, but training them in a single pass over streaming snapshots causes catastrophic forgetting of earlier time steps. The paper's goal is a training protocol that keeps a tiny memory footprint while preserving accuracy close to what an equivalent offline method would achieve.

Key Contributions

1. An in situ training protocol for INR-based neural compressors built on two memory buffers: a short queue of full snapshots of size T_f and a much longer queue of sketched snapshots of size T_s (the total stored buffer data is (T_f × n + T_s × k) × c, where n is the number of mesh nodes and k ≪ n is the sketch dimension). The loss is a sum of a full-data term and a weighted sketch term with weight λ > 0.

2. A theoretical justification for sketching as a regularizer. Theorem 3 ("JL Surrogate") shows, for squared ℓ2 loss, that the sketched loss upper-bounds the true loss by a factor 1/(1−ε) and that the minimizer of the sketched loss is within a factor (1+ε)/(1−ε) of the true minimizer, both with probability at least 1−δ. Theorem 4 ("Batch JL Surrogate") extends this to the loss averaged over multiple snapshots, requiring per-sketch failure probability δ/T_s. The authors state that, to their knowledge, this is the first use of sketching for regularizing neural compression, as opposed to using it directly for low-rank compression.

3. A mesh-agnostic formulation that requires only point coordinates and their connectivity-free locations, making the method applicable to uniform/structured, unstructured, and curvilinear meshes, and to non-Cartesian geometries in two and three dimensions.

4. A practical implementation using the fast Johnson–Lindenstrauss transform (FJLT) (Algorithm 3), with random seeds stored per snapshot so the same transform can be regenerated during training rather than stored. Subsampling sketches (Algorithm 2) are also implemented and investigated.

Main Findings

• Sketching lets in situ training approximately match offline. The abstract states that the presented in situ scheme approximately matches the performance of the equivalent offline method, which the authors highlight as the most important result.
• Strong reconstruction at high compression rates. The paper reports strong reconstruction performance at high compression rates on complex 2D and 3D simulation data over long time horizons, across unstructured grids and non-Cartesian geometries. Specific Section 3 tables and figures are not included in the supplied content.
• Motivating compression targets. The introduction sets the goal as compression rates greater than 100× with relative errors close to 1%.
• Offline reference reconstructions (Figure 2). Two offline INR examples are shown: streamwise velocity at 1.1% relative error and 40.5 dB PSNR for a 2,366× compression rate; and species concentration (%) at 0.5% relative error and 62.9 dB PSNR for a 1,392× compression rate. These are offline results, not in situ results.
• Buffer configuration used in experiments. The numerical experiments use T_f = 1, i.e., a full-snapshot queue holding a single snapshot, while the sketched buffer can be long because each sketch is small.
• Sketch theory bounds. The JL lemma gives k ≤ min{n, O(ε^{−2} log(m))}; the manifold JL transform used for the main result gives k = O(M log(NVRτ^{−1}ε^{−1}) log(δ^{−1}) / ε^2) for an M-dimensional manifold in R^N with volume V, condition number τ^{−1}, and geodesic covering regularity R; the FJLT achieves k = O(ε^{−2} log(m(δ − n^{−log^3(n)})^{−1}) log^4(n)) with an O(n log n) time cost.
• Practical sketch size rationale. The authors argue that regularization need not require a small ε, so smaller sketches than theory suggests can still serve as an adequate surrogate for the true loss, with actual learning driven by the full snapshots.
• Comparison limitation. The authors did not compare directly to the KD-INR knowledge distillation approach because its code and data were not publicly available; they leave this to future work. They also note that a prior hypernetwork continual-learning method limits the number of tasks to at most 100, whereas a task here is a single time step and PDE simulations often have far more time steps.
• Design intention behind the wording "in situ." The authors deliberately avoid "online" or "streaming" because their procedure cannot run indefinitely and will not end exactly when data generation ends.

Methodology in Plain English

The compressor is a neural network that takes space-time coordinates and outputs the physical field values, with a hypernetwork generating time-dependent weights. Because the simulation cannot store all snapshots, the training loop maintains a small working memory:

1. When a new snapshot arrives, a sketch matrix S_t is constructed and a random seed is saved. The full snapshot goes into a short FIFO queue of size T_f, and the sketched product S_t U_t goes into a longer queue of size T_s.
2. Each optimization step draws a mini-batch of full snapshots (size b_f) and a mini-batch of sketched snapshots (size b_s). The full batch provides the actual learning signal; the sketched batch provides a regularization term that discourages the network from drifting away from earlier reconstructed states.
3. Because only the sketch is stored for old snapshots, the sketched reconstruction S_t Φ(X, t; θ) is formed by evaluating the network on the subsampled points. Random seeds are stored so the exact same sketch can be reproduced at training time. The same sketch is broadcast across all channels for multi-channel data.

The theoretical justification is that a Johnson–Lindenstrauss sketch approximately preserves pairwise distances, so a small sketched loss implies a small true loss. Standard JL applies to a finite set of points, but the network output varies continuously with its parameters, so the authors invoke a manifold JL transform (Theorem 2) over the image manifold of the network. They choose the FJLT because it mixes all rows via a type II or III discrete cosine transform, is fast, and empirically outperforms alternatives; subsampling is also tested because, while it can fail badly by choosing unhelpful points, it coarsens rather than aggregates and thereby preserves some mesh structure.

Why This Matters

• Research impact: The work connects randomized sketching, formerly used mainly as a direct low-rank compression tool, with the problem of catastrophic forgetting in continual learning, and provides a bound showing when a sketch-based loss is a valid surrogate for the full loss. It also pushes neural compression of scientific data from offline into the in situ regime, while remaining mesh-agnostic.
• Real-world applications:
  • Compressing large time-dependent PDE simulations (e.g., fluid velocity fields) during the run, when the full dataset cannot be written to disk.
  • Compressing multi-channel scalar fields such as species concentration in reacting-flow or combustion simulations.
  • Compressing data on unstructured or curvilinear meshes, where connectivity-based compression schemes do not apply.
  • Long-horizon simulations where hundreds or thousands of time steps would exceed the task limits noted for prior hypernetwork continual-learning methods.
• Industry relevance: Relevant to national laboratories, supercomputing centers, and engineering organizations running large simulation campaigns on HPC systems, where I/O and storage, rather than compute, are often the bottleneck.

Future Directions

• Full experimental detail. The truncated content does not include Section 3's dataset descriptions, full benchmark tables, sketch-size ablations, or comparisons among FJLT, subsampling, and other sketch choices; these are the paper's main remaining evidence to inspect.
• Direct comparison to KD-INR. The authors explicitly leave a head-to-head comparison with the knowledge distillation approach to future work because the code and data were unavailable.
• Broader continual-learning applicability. The authors state that the protocol "may be of wider interest in the field of continual learning" beyond scientific compression, which invites testing on non-simulation streaming tasks.
• Tightening the theory to practice. The paper acknowledges that rigorously verifying whether the network's image is a Riemannian manifold with geometric properties small enough to yield a small sketch size is difficult; the relevant dimension is expected to depend on both the data and the network parameters, and the worst case of entirely uncorrelated data is noted as unresolved in the supplied content.

Target Audience

Researchers and graduate students in scientific machine learning, randomized numerical linear algebra, and continual learning, as well as HPC and computational-physics practitioners who need in situ compression for large simulations. The theoretical portion assumes comfort with Johnson–Lindenstrauss theory, Riemannian manifolds, and low-rank approximation, while the algorithmic portions (Algorithms 1–3) are readable for engineers implementing replay-style buffers.

Authors’ 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.