[Paper Review] Variational Autoencoding Neural Operators

Jacob H. Seidman, Georgios Kissas+2 more|arXiv (Cornell University)|Feb 20, 2023

Machine Learning in Materials Science4 citations

TL;DR

This paper introduces Variational Autoencoding Neural Operators (VANO), a novel framework that integrates variational autoencoders with operator learning architectures to enable unsupervised representation learning and generative modeling of functional data in infinite-dimensional spaces. By formulating a rigorous variational objective in function space, VANO learns disentangled, latent representations of input functions and reconstructs them via differentiable operators, achieving state-of-the-art performance on Gaussian random fields, Cahn-Hilliard phase separation, and real satellite interferogram data with improved generalization and discretization invariance.

ABSTRACT

Unsupervised learning with functional data is an emerging paradigm of machine learning research with applications to computer vision, climate modeling and physical systems. A natural way of modeling functional data is by learning operators between infinite dimensional spaces, leading to discretization invariant representations that scale independently of the sample grid resolution. Here we present Variational Autoencoding Neural Operators (VANO), a general strategy for making a large class of operator learning architectures act as variational autoencoders. For this purpose, we provide a novel rigorous mathematical formulation of the variational objective in function spaces for training. VANO first maps an input function to a distribution over a latent space using a parametric encoder and then decodes a sample from the latent distribution to reconstruct the input, as in classic variational autoencoders. We test VANO with different model set-ups and architecture choices for a variety of benchmarks. We start from a simple Gaussian random field where we can analytically track what the model learns and progressively transition to more challenging benchmarks including modeling phase separation in Cahn-Hilliard systems and real world satellite data for measuring Earth surface deformation.

Motivation & Objective

• To develop a general method for adapting operator learning architectures to perform unsupervised representation learning and generative modeling on functional data.
• To address the limitations of point-wise training objectives in existing VAEs for functional data, which are sensitive to sampling density and lack function-level generalization.
• To establish a rigorous mathematical formulation of the variational objective in infinite-dimensional function spaces, enabling training of VAEs over continuous fields.
• To enable latent space disentanglement and reconstruction of functional data at arbitrary spatial locations, independent of input discretization.
• To demonstrate the framework’s effectiveness across diverse benchmarks, including synthetic Gaussian fields, nonlinear PDEs, and real Earth observation data.

Proposed method

• Proposes a variational objective formulated in function space using kernel mean embeddings and generalized MMD (GMMD) to compare probability measures over functions.
• Employs a parametric encoder $\mathcal{E}_{\phi}$ to map input functions to a latent distribution $\mathbb{Q}^{\phi}_{z|u}$, and a decoder $\mathcal{D}_{\theta}$ to reconstruct the input from sampled latent variables.
• Uses a neural operator architecture (e.g., FNO or U-Net variants) as the encoder and decoder, enabling discretization-invariant inference across arbitrary spatial grids.
• Applies GMMD with a family of Gaussian kernels over a range of length-scales to robustly measure distributional similarity between real and generated functional data.
• Trains the model end-to-end via backpropagation, minimizing the variational lower bound (ELBO) with a reparameterization trick for stochastic gradient optimization.
• Employs circular variance and skewness metrics to evaluate angular function generation quality in periodic data, such as interferometric phases.

Experimental results

Research questions

• RQ1 Can operator learning architectures be systematically adapted to perform unsupervised representation learning and generative modeling of functional data in infinite-dimensional spaces?
• RQ2 How can a rigorous variational objective be formulated in function space to enable training of VAEs on continuous fields rather than pointwise samples?
• RQ3 Does VANO achieve better generalization and robustness compared to existing VAEs that optimize pointwise likelihoods, especially under varying sampling densities?
• RQ4 Can VANO learn disentangled, meaningful latent representations of complex functional data such as phase separation patterns and satellite interferograms?
• RQ5 To what extent does the use of GMMD with a kernel family improve distributional matching compared to single-kernel MMD in functional generative modeling?

Key findings

• VANO achieves state-of-the-art performance on 2D Gaussian density generation, with GMMD values indicating high distributional similarity between real and generated data across all tested length-scales.
• On the Cahn-Hilliard system at $256 \times 256$ resolution, VANO (nonlinear) required 1,667,000 trainable parameters and 1,660 seconds of training time, outperforming baseline VAEs in reconstruction fidelity and generalization.
• For InSAR interferogram data, VANO achieved a test loss of 0.0048 with 11.1 million parameters, demonstrating strong performance on real-world geophysical data with complex phase patterns.
• The model successfully captured the circular nature of interferometric phases, with generated data showing low circular variance (0.08) and skewness (0.12), indicating accurate modeling of angular distributions.
• VANO’s latent space disentanglement was validated via interpolation and conditional sampling, producing smooth, physically plausible transitions in functional outputs across latent dimensions.
• The use of GMMD with a kernel family ($\sigma \in [0.1, 20]$) led to more robust and stable training compared to single-kernel MMD, which often collapsed under high-resolution or sparse sampling.