Compressive Neural Representations of Volumetric Scalar Fields

Lossy compression of large scale volumes has been an important topic in recent years, particularly as our ability to simulate and acquire data grows. Many of the early works in this area focused on using discrete cosine transforms [YL95] or wavelet-based compression [GWGS02, IP99, Mur93, WMB^∗11], often specifically focusing on interactive rendering applications [RGG^∗13, SW03]. By using such transforms, these techniques enable compression by allowing the end user to separate features within a transformed space (i.e. removing high frequencies while preserving low frequencies) and ultimately sparsify the data representation. ZFP performs transformations customized at the block level to achieve compression with fast I/O access [Lin14]. More recent techniques, such as TAMRESH [SMP13] and TTHRESH [BRLP19] employ tensor decomposition [BRP16]. TTHRESH is particularly notable as a recent state-of-the-art compression technique that we directly compare our results against in this work.

Other techniques focus on data fitting as the workhorse for compression. For example, ISABELA [LSE^∗11] and SZ [DC16, TDCC17] use curve fitting to approximate the input data. Liang et al. outperform SZ and ZFP using an improved Lorenzo prediction [LDT^∗18]. Data fitting is related to our neural network based approach, although the method to achieve the fit is different. Using data models can be a powerful mechanism for providing a compression technique. Peterka et al. [PNG^∗18] study the use of nonuniform rational B-spline functions (NURBS) for compressive volume representations via adaptive refinement of control points. This work shows the benefit of smooth representations for analytically-defined gradients, however an important distinction in our approach is that we can directly optimize for gradients. Other examples focus on models that preserve specific aspects of the data, e.g. topological features [SPCT18], graph-based models like SQ [IKK12], and dictionaries [DMG20]. COVRA uses an octree decomposition to model data blocks which are compressed at multiple resolutions [GIGM12]. Data fitting broadly considers the tradeoff between how many samples are used and how much precision is devoted to each sample. Recently, Hoang et al. provide a precision-resolution tree that is used for efficient selection of either fewer data samples or data samples encoded with fewer bits [HSB^∗20].

Recently, numerous researchers have deployed deep learning and used neural networks to enable a variety of visualization tasks. Berger et al. utilize generative adversarial networks (GANs) to analyzing the space of visual parameters in a volume renderer, providing a new interface for transfer function design [BLL18]. He et al. generalize the idea of learning conditioned on parameter spaces to include ensembles of volumes with InSituNet [HWG^∗19]. Han et al. encode multivariate volume data for variable-to-variable translation [HZX^∗20]. Jakob et al. consider the problem of interpolation using neural networks defined on a large ensemble of vector fields [JGG20]. While the applications are different, all of these works share a common thread in that they construct representations that blend visualization tasks (e.g. rendering, interpolation) with data representations.

Closely related to compression is the task of super-resolution, for which many techniques have been developed in the computer vision communities. More recently, researchers have shown that deep learning can be used to achieve super-resolution on volumetric data. Xie et al. developed tempoGAN for super-resolution of volumetric fluids in computer graphics [XFCT18]. Weiss et al. developed new techniques for isosurface rendering that similarly used deep learning for super-resolution [WCTW19]. Han and Wang developed neural networks for super-resolution in temporal [HW19] and spatial [HW20] super-resolution, and Guo et al. achieved spatial super-resolution in vector fields [GYH^∗20]. Notably, many of these techniques are tested against possible compression techniques, but none of the above were specifically designed for compression. Thus, while they are compressive, their main strengths lie elsewhere.

Recently, the machine learning community has explored using deep learning to construct implicit representations, which can be used to model a variety of data. We utilize a similar approach in our network design, and as seen in recent works this can achieve superior performance. Park et al. developed DeepSDF for representing shape using signed distance fields [PFS^∗19]. Recent extensions to this work include overfit neural networks [DNJ20] and DualSDF [HAESB20] that allow for tradeoffs in the representation precision and flexbility. Similarly, Chen and Zhang use implicit fields for generative shape modeling [CZ19]. Alternatively, instead of training to produce a signed distance field, Mescheder et al. propose to construct occupancy fields [MON^∗19]. Besides just representing shape, other recent applications include texture synthesis [OMN^∗19] and view synthesis [MST^∗20]. As this is an emerging area, we are still discovering new ways to improve these networks. Notable recent works that informed our network design include using periodic activations [SMB^∗20] and Fourier feature mappings [TSM^∗20]. Our approach is distinct from prior works in that we investigate how to make implicit neural representations compressive for arbitrary scalar fields [DNJ20, SMB^∗20, TSM^∗20].

The purpose of the encode step is to train the neural network to accurately model the field at its provided samples. More specifically, we aim to minimize the following mean-squared error loss:

In practice, the loss function is optimized by performing stochastic gradient descent, assembling minibatches at each iteration of training by randomly selecting samples from the provided field.

What remains is the design of the network. We opt for simplicity in network architectures, to avoid the brittleness of hyperparameter tuning and support a diverse set of volumetric fields. To this end, we build off of the approach of SIREN [SMB^∗20], where a neural network is defined by a sequence of fully-connected layers that use sinusoidal activation functions – no normalization, e.g. batch normalization [IS15], is necessary for successful training. SIREN has several benefits. First, as demonstrated by Sitzmann et al., training is ensured stable through a principled initialization scheme that controls for the distribution of activations at layers, namely that all layers are arcsine-distributed. Secondly, the use of sinusoidal activations admits functions that are $C^{\infty}$ differentiable. Our approach takes advantage of both of these properties.

Given the network design of SIREN, to setup the network we need only define the fully-connected layers. To this end, our method accepts two inputs: the number of layers $l$, and the total number of network weights $m$. We derive an integer $k$ such that the first weight matrix has size $k\times d$, the last weight matrix has size $d\times 1$, and all intermediary weight matrices are of size $k\times k$, such that the total number of weights of the network is approximately $m$. Bias vectors are similarly derived from inputs $m$ and $l$. All weights are of 32-bit floating-point precision. The resulting set of weight matrices and bias vectors, collectively, represents our vector of parameters $\Theta$, e.g. the compressed representation. Note that it is also possible to vary the sizes of the weight matrices, e.g. to gradually increase the number of hidden units as a function of layer depth. Experimentally, we found this made little difference, thus for simplicity we chose a fixed number of hidden units in intermediary layers.

The number of layers, $l$, should be large enough to handle the modeling of arbitrarily complex fields. In practice, however, setting $l$ too large can be detrimental, posing challenges for training due to vanishing gradients. As an example, we show results for two different volumes in Fig. 2 (red solid lines), varying the number of layers, where the quality of the volume is measured through peak signal-to-noise ratio (PSNR), the mean-squared error normalized by the data range. Note that, at a certain point, as the model depth increases performance decreases. Further, the network depth that leads to highest PSNR varies between the two volumes.

In order to ensure our models are robust with respect to the number of layers across different volumes, we enrich SIREN with residual connections [HZRS16]. Specifically, we modify the identity mappings of He et al. [HZRS16] to ensure that the distribution of activations in each layer are within the range $[-1,1]$, in accordance with the assumptions of SIREN [SMB^∗20]. Each residual block in our network takes the following form, omitting bias vectors for clarity (c.f. Fig. 1 dashed black arrows):

where $\mathbf{a}_{i}$ represents a vector of activations from a previous block $i$, $\mathbf{a}_{i+1}$ is a vector of activations at the next block $i+1$, and $\mathbf{M}^{1}_{i+1}$ and $\mathbf{M}^{2}_{i+1}$ are a pair of learnable matrices associated with block $i+1$. In Fig. 2 (green solid lines) we show the effect of residual connections – note we treat each block as containing 2 layers. Note that (a) training is more stable as the network depth increases, and (b) we obtain a boost in performance over non-residual connections. Consequently, unless otherwise stated, all of our networks employ 8 residual blocks.

Once the network has been trained, and the volume encoded, decoding is rather straightforward, as it amounts to evaluating the neural network. To reconstruct the scalar field we simply evaluate our network at all grid vertices within the volume. This evaluation is relatively efficient given the data parallelism of the network, e.g. a sequence of matrix multiplications and $\sin$ function evaluations.

Beyond reconstruction, a key advantage of neural scalar field representations is the support for random access function evaluation. We may evaluate our function at any position, not just grid vertices, hence the network serves as an interpolant. Furthermore, we can compute higher-order information from the network directly, eliminating the need to numerically estimate derivatives. These features allow us to use the network directly for visualization purposes, rather than reconstructing the full volume a priori. As an illustration, we have developed a neural volume renderer with our network, where at each step of ray marching, we evaluate our network at all current ray positions to obtain function values, see Fig. 5a for an illustration. Moreover, it is trivial to support direct illumination via the computation of function gradients as we ray march, subsequently normalized for shading purposes. We highlight this feature in Fig. 5b. The quality of the volume-rendered images suggests that our network serves as a good interpolant, not merely overfitting to the sampled field’s values.

We first compare our method with tthresh for 3D scalar fields. We run our method over a sequence of compression levels, and plot the resulting compression ratio against PSNR. For these experiments we do not use gradient regularization, but rather, delegate this to Section 4.4. As tthresh accepts an error/accuracy tolerance (e.g. PSNR) as input parameter, rather than compression ratio, we run tthresh for a sequence of PSNR values such that they give compression ratios that are roughly in the range of what we consider.

Fig. 7 shows the quantitative results. We find that our method is, largely, an improvement over tthresh, with the only exception being the jet volume for low compression ratios. In particular, we find that our method obtains larger gains in performance the higher the compression ratio. This suggests that when our networks are underparameterized (e.g. much fewer parameters than the resolution of the data), they remain robust, and can still obtain good approximations to the given scalar field.

To further qualitatively assess this robustness in high compression regimes, we visually compare volume renderings of our method with tthresh for the magneto-hydrodynamic simulation, namely its magnetic field (mhd_bx) and pressure field (mhd_p) in Fig. 8. We find that our method is able to reproduce smoother surfaces in the rendering, as shown in the magnetic field plot (a), while in the pressure field (b), we find that components of the volume, corresponding to the chosen transfer function, are better retained through our method. Indeed, we find that our method gracefully degrades for extreme compression ratios, unlike block-based methods [Lin14, DC16] that produce visible artifacts at block boundaries. Global decomposition-based methods [BRP16, BRLP19] also do not suffer from locality issues, but we find that they can produce high-frequency noise in high compression regimes. We demonstrate this in Fig. 9 for the rt dataset, comparing our method with tthresh under approximately the same PSNR values. We first note that our method produces compressed representations approximately half the size of tthresh. Further, we find that our method produces smoother scalar fields, more reflective of the ground truth as shown on top, unlike the high frequencies reproduced by tthresh.

We evaluate our approach’s capability to compress time-varying data by comparing our method with tthresh on isabel and asteroid, assembling the first twelve time steps from isabel and ten consecutive time steps towards the beginning of the asteroid simulation. While we show results per time-step, we note that both methods treat the data as 4D volumes.

Fig. 10 show comparisons between our method and tthresh over multiple timesteps with regard to 3 different compression ratios. Note that the compression ratios for our method and tthresh are slightly different, since it is difficult to precisely set the compression ratio for tthresh. As the top of Fig. 10 shows, our approach achieves superior results on asteroid dataset when compared with tthresh, with more stability on both ends. In part, this is due to the fine resolution of sampling in time, where the simulation evolves slowly. For the isabel dataset, our method performs comparable to tthresh at low compression ratio, except for the start and end timesteps where we outperform, as shown in the bottom of Fig.  10. For higher compression levels, however, our method gains larger boosts in performance over tthresh at all timesteps, showing the preference of our method for high-compression ratio scenarios.

Here we examine the impact of spatial resolution on the quality of our compression method. In particular, we wish to gain insight on the following question: do coordinate-based MLPs benefit from high spatial resolution? To test this, we took the isotropic_p $512^{3}$ volume, and grabbed centered crops of size $256^{3}$ and $128^{3}$. We expect this volume to contain similar statistics for sufficiently large spatial crops, since it corresponds to isotropic turbulence at relatively small scales. We then ran our method on the volumes, setting the network size-to-volume ratio ($\frac{C}{m})$ constant, namely 50, 100, 150, 200, and quantize the network weights to 9 bits. Note that for larger volumes, the storage of the cluster centers leads to overall higher compression ratios, thus for equivalent $\frac{C}{m}$ values lower resolution volumes are given the benefit of the doubt.

Fig. 11 (top) shows the results, plotting performance as PSNR with respect to ground truth. We can see that the higher the spatial resolution, the better the performance, despite (a) the network sizes setup in the same manner, and (b) the higher compression ratios for the larger volumes. We view this as an encouraging property, namely that network capacity is not strictly tied with the spatial resolution. On the bottom of the figure we show volume renderings from low and high-end compression ratios, illustrating that our approach captures the predominant features of the volume even for high levels of compression.

Last, we show how gradient regularization can be used to produce better-behaved scalar fields in comparison to prior works. Specifically, for the rt dataset we obtain target gradients through central finite differencing, and train our method using the loss of Eq. 3. We provide our resulting PSNR as a target accuracy for tthresh, yielding comparable PSNR values as shown in Table 3. However, a gap exists between the methods in terms of the Gradient-based PSNR. As discussed in Sec. 3.1.2, the effect of gradient regularization will be prominent when isosurfacing the volume, and we can observe in Fig. 12 that tthresh indeed produces isosurfaces that contain considerable noise – this corroborates findings in Fig. 8(a) and Fig. 9. Indeed, the preservation of higher-order information is nontrivial through existing compression methods, whereas our method can incorporate gradient preservation as part of optimization.