Volumetric and Grid VAEs

• Volumetric and grid VAEs are generative models employing spatially-structured latent representations to capture geometric and anatomical features in high-dimensional data.
• They integrate innovations such as low-rank tensor normal approaches, hierarchical vector quantization, and slice-wise covariance modeling to enhance reconstruction and compression.
• These architectures apply to diverse domains including medical imaging, 3D scene synthesis, and unsupervised representation learning, offering improved fidelity and interpretability.

A variational autoencoder (VAE) is a generative model in which observed data is reconstructed via a latent representation sampled from an inferred posterior. When extending the VAE paradigm to data with strong spatial or geometric structure—such as 2D image grids, volumetric medical images, or 3D scene representations—several methodological innovations have emerged: grid VAEs structured around spatial latent matrices or tensors, and volumetric VAEs designed for high-dimensional cube-like or radiance field data. This article surveys the concrete methodologies, architectures, and empirical properties of volumetric and grid VAEs, referencing matrix- and tensor-variate normal latent variables, vector-quantized frameworks, slice-covariance models, and explicit geometry-aware scene likelihoods.

1. Mathematical Foundations of Volumetric and Grid Latent Spaces

Canonical VAEs use vectorial latent variables $z\sim\mathcal N(\mu,\Sigma)$, but this representation does not explicitly model spatial structure. Grid VAEs ("spatial VAEs") replace vector latents with matrix- or tensor-variate normal distributions. For the matrix case, a latent $Z\in\mathbb R^{d\times d}$ follows

$Z \sim \mathcal N_{d,d}(M, U\otimes V)$

with mean matrix $M$, row-covariance $U$, column-covariance $V$, and density

$$p(Z) = \frac{\exp(-\frac{1}{2} \operatorname{tr}[V^{-1}(Z - M)^\top U^{-1}(Z - M)])}{(2\pi)^{d^2/2} |U|^{d/2} |V|^{d/2}}$$

(Wang et al., 2017). The low-rank MVN approach further reduces overparameterization of $M$ via an outer product $M=AB^\top$, $A,B\in\mathbb R^{d\times r}$ ($r\ll d$). These latent feature maps are stacked per channel to form a 3D tensor $Z\in\mathbb R^{d\times d\times N}$.

For volumetric data, tensor-variate normal distributions generalize this framework:

$$T \sim \mathcal N_{d_1,d_2,d_3}(M, \Sigma_1\otimes \Sigma_2\otimes \Sigma_3)$$

where $T$ encodes volumetric spatial structure (Wang et al., 2017). Parameter count is reduced by a low-rank Tucker decomposition of the mean.

2. Architectural Implementations and Inference Mechanics

Grid VAEs and Spatial VAEs

Spatial VAEs employ convolutional encoders generating feature maps, with latent matrices/tensors parameterized by means $M$ (via low-rank factors) and diagonal or Kronecker-structured covariance matrices. The corresponding decoder is a deconv/upsampling network that reconstructs outputs from the latent grid(s). For each map, samples are generated as

$$Z_{(i,j,k)} = (A_kB_k^\top)_{ij} + \sqrt{u_{k,i} v_{k,j}} \epsilon_{(i,j,k)}$$

with diagonal $u_{k}, v_{k}$ (Wang et al., 2017).

Hierarchical Volumetric VQ-VAE

Volumetric vector-quantized VAEs (VQ-VAE) use 3D encoders/decoders paired with multi-scale latent quantization. The network encodes a 192×256×192 input volume into three levels of quantized grids (top: $48^3$, middle: $12^3$, bottom: $3^3$), each with its own codebook. Decoding proceeds hierarchically, with higher-resolution grids conditioned on lower-resolution quantized codes. Training employs both a reconstruction loss (with pixel- and gradient-matching variants) and a VQ codebook loss with exponential moving average (EMA) code update (Tudosiu et al., 2020).

Volumetric VAEs for Medical Images via Slice Covariances

Modelling volumetric MRIs with computational efficiency is enabled by VAEs trained on 2D slices, followed by fitting a Gaussian model in the $L$-dimensional slice-wise latent space across the slice axis. The joint prior over concatenated latent coordinates across all slices is a block-diagonal Gaussian with empirically estimated means and covariances. New 3D volumes are sampled by drawing slice-stack latents with the appropriate covariance, and decoding per-slice with the 2D VAE decoder (Volokitin et al., 2020).

NeRF-VAE: Probabilistic Volumetric Scene Models

NeRF-VAE augments the neural radiance field (NeRF) paradigm with a VAE framework: the latent $z$ encodes the scene, with a conditional NeRF function $G_\theta(x, d; z)\to(c, \sigma)$ parameterizing color and volume density. Given $z$, images are rendered by differentiable volumetric integration along rays:

$$C(r; \theta, z) = \int_{t_n}^{t_f} T(t) \sigma(x(t)) c(x(t), d) dt, \qquad T(t) = \exp(-\int_{t_n}^t \sigma(x(s)) ds)$$

The likelihood is Gaussian per pixel, and the overall evidence lower bound (ELBO) is optimized jointly over encoder and decoder parameters. Encoder amortizes inference over context images and camera poses, optionally refined by iterative steps, while decoder conditioning is performed by concatenation, adaptive normalization, or a spatial attention mechanism over the latent map (Kosiorek et al., 2021).

3. Training Objectives and Optimization Strategies

For both grid and volumetric VAEs based on (tensor-)normal distributions, the ELBO includes a closed-form KL between the matrix- or tensor-normal approximate posterior and the prior. For the hierarchical VQ-VAE, the loss combines reconstruction and VQ codebook terms:

$L = L_{\rm recon}(x, \hat x) + L_{\rm codebook}$

where $L_{\rm recon}$ may operate in voxel space or a 3D DCT domain, and $L_{\rm codebook}$ includes commitment and embedding updates.

In NeRF-VAE, the ELBO with optional $\beta$-annealing is sampled over rays for efficiency and supports both coarse and fine sampling per the original NeRF protocol, using Adam for optimization. Iterative inference with gradient-based refinement of the posterior parameters further reduces the amortization gap (Kosiorek et al., 2021).

4. Empirical Properties and Quantitative Results

Compression and Reconstruction

The 3D VQ-VAE compresses full-resolution brain MRIs linearly to 0.825% of the original size, with all latents stored as 8-bit indices. Reconstruction metrics improve substantially over GAN baselines: Multi-Scale SSIM increases from $\sim$0.5 (GAN) to $\sim$0.99, log MMD reduces from $\sim$15.7 to $\sim$6.7, and Dice overlaps increase for white matter (WM), gray matter (GM), and CSF segmentations, with all improvements statistically significant ($p<0.01$) (Tudosiu et al., 2020).

Generative Modeling of Volumes

The slice-VAE + Gaussian model for MR volumes achieves superior fidelity and sample diversity compared to 3D GAN/VAEs, as measured by MMD, MS-SSIM, and the Realistic Atlas Score (RAS), which quantifies anatomical plausibility based on segmentation overlap after volume registration. For example, at $128^3$ resolution, the proposed method reports MMD=19890, RAS$\approx$0.845, outperforming a 3D $\alpha$-WGAN baseline (MMD=64446, RAS$\approx$0.842) (Volokitin et al., 2020).

Geometric Consistency in 3D Scene Generation

NeRF-VAE reconstructs novel views from as few as 1–4 context images, compared to per-scene NeRF requiring dozens or hundreds for equivalent MSE. It maintains geometric and color consistency even under out-of-distribution camera poses, whereas convolutional baselines such as GQN break down. Attention-based conditioning yields the best MSE vs. latent KL tradeoff, particularly in complex environments. Posterior sampling reflects calibrated uncertainty in ambiguous contexts (Kosiorek et al., 2021).

Effect of Spatially Structured Latent Spaces

Low-rank-matrix or tensor-grid VAEs provide sharper, more globally coherent samples than naïve vector-latent VAEs, both visually (CelebA, CIFAR-10) and quantitatively (MNIST log-likelihood), with only moderate training cost increase. Kronecker-structured covariances and low-rank means are essential; naïve grid restructurings or full covariances are less effective (Wang et al., 2017).

5. Applications, Limitations, and Extensions

Grid and volumetric VAEs are central in unsupervised representation learning and generative modeling for high-dimensional spatial data domains such as medical imaging (MRI, CT, PET), 3D scenes, and voxel-based shape synthesis. The morphological integrity and anatomical faithfulness attained by VQ-VAE models can enable downstream analysis tasks and transfer learning across populations. Volumetric VAE designs facilitate efficient sampling, memory scaling, and separation of global and local information.

Documented limitations include high GPU memory requirements for 3D models, evaluation metrics focused on pixel/voxel similarity rather than clinical outcomes, and the absence of explicit non-uniform quantization. Extensions comprise autoregressive priors (e.g., PixelCNN3D) over discrete codes, graph-based or spatially adaptive quantization, and modality-fused volumetric backbones (Tudosiu et al., 2020).

NeRF-VAE opens new directions for amortized inference over 3D environments, probabilistic uncertainty modeling in scene synthesis, and generalization to out-of-distribution observations, paving the way for broader application of VAEs in geometry-aware and physics-based domains (Kosiorek et al., 2021).

6. Comparative Summary of Principal Architectures

Model / Reference	Latent Structure	Main Application	Key Feature
Matrix-variate/tensor-variate VAE (Wang et al., 2017)	$d\times d$ (2D grid) / $d_1\times d_2\times d_3$ (3D)	Images, Volumes	Explicit spatial structure in prior, low-rank means
3D VQ-VAE (Tudosiu et al., 2020)	Hierarchical multi-scale discrete grids	Volumetric medical images	Compression $<1\%$, morphological fidelity
Slice-VAE + Gaussian (Volokitin et al., 2020)	Per-slice latent vectors with across-slice covariance	Volumetric generative modeling	Block-diagonal covariance, Realistic Atlas Score
NeRF-VAE (Kosiorek et al., 2021)	Latent code + MLP, optionally spatial/attention map	Generative 3D scenes	Differentiable volume rendering, attention, amortized inference

Each approach leverages spatial or geometric structure at the latent or likelihood level, yielding improved sample quality, interpretability, and downstream task utility over vector-based VAEs or convolution-only decoders.

---

In sum, volumetric and grid VAEs—spanning low-rank tensor-normal frameworks, hierarchical vector quantization, spatial-structured slice models, and geometry-aware NeRF hybrids—constitute a principled set of methods for generative modeling of high-dimensional data with spatial, anatomical, or geometric regularity. Their architectures and training losses directly encode the inductive biases and modeling requirements of their target domains, and empirical results demonstrate measurable gains in fidelity, sample coherence, and compression.