Universal Functional Regression with Neural Operator Flows

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-06 06:49:21

Abstract: Regression on function spaces is typically limited to models with Gaussian process priors. We introduce the notion of universal functional regression, in which we aim to learn a prior distribution over non-Gaussian function spaces that remains mathematically tractable for functional regression. To do this, we develop Neural Operator Flows (OpFlow), an infinite-dimensional extension of normalizing flows. OpFlow is an invertible operator that maps the data space into a Gaussian process, allowing for exact likelihood estimation of point evaluations of functions. OpFlow enables robust and accurate uncertainty quantification via drawing posterior samples. We empirically study the performance of OpFlow on regression tasks with data generated from Gaussian processes with known closed-form posterior distribution as well as highly non-Gaussian real-world earthquake time-series with an unknown closed-form posterior distribution.

Main paper: 
 introduction: 
 Inference on function spaces is essential to the physical sciences and engineering. It is often desirable to infer the function given a sparse number of observations. There are numerous problems in which functional regression plays an important role, such as inverse problems, forecasting, and data imputation/assimilation. However, regression on function spaces can be particularly challenging in real world problems where the underlying stochastic process is often unknown.Much of the work on functional regression and inference has relied on Gaussian processes (GPs) [18] . GP regression (GPR) provides several advantages for function space inference including robustness and mathematical tractability. Despite widespread adoption, the assumption of a GP prior for functional inference problems can be rather limiting, particularly in scenarios where the data is heavy-tailed or distributions are multimodal. This underscores the need for models with greater expressiveness, allowing for regression on data arising from unknown stochastic processes. We refer to such regression problems as universal functional regression (UFR).To solve UFR problems, we believe it requires a model with two primary components. First, the model needs to be capable of learning priors over data that lies in function spaces. There has been much recent progress on learning priors on function spaces [17, 13, 19, 7, 16, 1, 6, 8, 5] . Second, the models need a framework for performing functional regression with learned function space priors and likelihoods-a component that is critically missing from the previous works described above. Addressing these challenges not only expands the models available for functional regression but also enhances our capacity to extract meaningful insights from complex datasets.We aim to address this gap by developing a viable formulation for UFR with a novel learnable bijective operator, OPFLOW, that extends normalizing flows to function spaces. OPFLOW consists of an invertible neural operator trained to map the data point-process distribution to a known and easy-to-sample GP distribution. Given point value samples of a data function, OPFLOW allows for computing the likelihood of the observed values, a principled property a priori held by GP. Using this property, we formally define the problem setting for UFR and develop an approach for posterior estimation with Stochastic Gradient Langevin Dynamics (SGLD) [21] . We demonstrate UFR with OPFLOW on various Gaussian and non-Gaussian processes. 
related work: 
 Neural Operators. Unlike traditional neural networks that primarily work with fixed-dimensional vectors, neural operators are a new paradigm in deep learning designed to operate on functions. Neural operators learn maps between infinite-dimensional function spaces, making them inherently suitable for a wide range of scientific tasks involving partial differential equations (PDEs) [10, 11, 12] .Function space generative modeling with Neural Operators. Several discretization invariant generative models on function spaces have been recently developed using Neural Operators including Generative Adversarial Neural Operator (GANO) [17] , Denoising Diffusion models (DDO) [13] , and Variational Autoencoding Neural Operator (VANO) [19] . All of these methods have exhibited superior performance over their finite-dimensional counterparts by directly learning mappings between function spaces. However, they lack the ability to evaluate exact likelihoods of samples.Normalizing Flows. Normalizing flows are a class of flow-based finite-dimensional generative models, that are usually composed of a sequence of invertible transformations [9, 4, 3] . By gradually transforming a simple probability distribution into a more complex target distribution with the invertible architecture, normalizing flows enable exact likelihood evaluation and direct sampling.Traditional normalizing flows are defined on a finite-dimensional spaces and constrained to evaluating likelihoods on fixed, regular-size grids.Nonparametric Bayesian and Gaussian Process Regression. Nonparametric Bayesian models are defined on an infinite-dimensional parameter space, yet they can be evaluated using only a finite subset of parameter dimensions to explain observed data samples [15] . These models provide a flexible regression framework by offering analytically tractable posteriors. Despite the great advantages of these models, they heavily depend on the chosen prior, which can restrict their adaptability in complex real-world scenarios [2, 20] . Within the family of nonparmateric Bayesian models, GPR is a robust framework characterized by its analytical posteriors [18] . However, GPR assumes both the prior and posterior are Gaussian, which can limit its applicability. 
neural operator flow: 
 We introduce the Neural Operator Flow (OPFLOW), an innovative framework that extends finitedimensional normalizing flows to infinite-dimensional function spaces. The architecture is shown schematically in Fig. 1 . OPFLOW retains the invertible structure of normalizing flows, while directly operating on function spaces. It allows for exact likelihood estimation for point estimated functions.OPFLOW is composed of a sequence of layers, each containing actnorms [9] , domain partitioning, and affine coupling. In particular, domain partitioning mimics the checkerboard pattern in [4, 9] .Training. The goal of OPFLOW is to learn a mapping between samples from the data space U to the latent space A where a ∈ A, a : D A → R d A and u ∈ U, u : D U → R d U . Let P U and P A as probability measures defined on U and A, respectively. Since OPFLOW is a bijective operator, we only need to learn the inverse mapping F θ : U → A, as the forward mapping is immediately available.We train OPFLOW in a similar way to other infinite-dimensional generative models [17, 19, 13] with a special addition of domain alignment. The space of functions for A is taken to be drawn from a GP, which allows for exact likelihood estimation and efficient training of OPFLOW by minimizing the negative log-likelihood. Since OPFLOW is discretization agnostic [17, 10] , it can be trained on various discretizations of D U and can be later applied to a new set of discretizations. In the limit, it can be applied to the whole of D U [10] . We now define the training objective L for OPFLOW training as follows,L = -min θ∈Θ E u∼P U [log p θ (u| D )], log p θ (u| D ) = log p(a| D A ) + s i=1 log | det( ∂(v i | D i V ) ∂(v i-1 | D i-1 V ) )|. (1)The inverse operator is composed of s invertible operator layers withEQUATIONgradually transforms v 0 to v 1 , • • • v s-1 , v s , where v 0 = u, v s = a.The alignment naturally holds within the domains associated with v 0 , v 1 , • • • v s as each layer of F θ is an invertible operator. V i is the collection of positions for function v i with D 0 V = D, D s V = D A , and each data point has its own discretization D.Universal Functional Regression with OPFLOW. We outline an algorithm for UFR with OPFLOW for general Bayesian inference problems. We assume that OPFLOW has been trained previously and sufficiently approximates the true stochastic process of the data. We thus have an invertible mapping between function spaces U and A. We use the trained OPFLOW as a learned prior distribution and fix the parameters for functional regression.Similar to GPR, suppose we are given ũobs which consists of pointwise and potentially noisy observations of a function u observed at points in the set D ⊂ D U . We refer to the noise-free evaluations at D as u obs , with the additive noise ϵ ∼ N (0, σ 2 ). We aim to infer the function's values on a set of new points D ′ where D ⊂ D ′ given the observation ũobs . For the likelihood of the values on points in D ′ , i.e., the likelihood of u ϕ , we have,EQUATIONwhere p θ (u| D ) denotes the probability of u evaluated on D with the learned prior (OPFLOW), and r is the cardinality of set D with r = |D|. Maximizing Eq. 2 results in the maximum a posteriori (MAP) estimate p θ (u ϕ |ũ obs ) where u ϕ denotes the MAP estimate given ũobs . For sampling the posterior in Eq. 2, we utilize the fact that OPFLOW is a bijective operator and the target function u ϕ uniquely defines a ϕ . Therefore, drawing posterior samples of u ϕ is equivalent to drawing posterior samples of a ϕ in latent space where we can use SGLD with Gaussian random field perturbation. We utilize SGLD to sample p θ (u ϕ |ũ obs ). We initialize a ϕ |ũ obs with the MAP estimate of p θ (u ϕ |ũ obs ), i.e., u ϕ , and follow Langevin dynamics in the latent space to sample u ϕ . 
universal functional regression, experiments: 
 In this section, we aim to provide an evaluation of OPFLOW's capabilities for UFR, as well as generation tasks. We consider several datasets, composed of both Gaussian and non-Gaussian processes, as well as a real-world dataset of earthquake recordings with highly non-Gaussian characteristics.Here, we show UFR experiments using OPFLOW. For training, we take A ∼ GP() with a Matern covariance parameterized by length scale l and roughness ν. The mean function is set to zero over D.The noise level (σ 2 ) of observations in Eq. 2 is 0.01 for all regression tasks.Gaussian Processes. This experiment aims to duplicate the results of classical GPR with OPFLOW.For training we use 30,000 GP samples with a Matern covariance parameterized by l U = 0.5, ν U = 1.5. For the latent space GP, we use l A = 0.1, ν A = 0.5. The resolution of D is set to 128. Although GPR requires the GP mean and covariance of the data distribution to be known, OPFLOW learns them instead. For regression, we infer the posterior from just 6 random observation points in D. Gaussian random fields. We build a 2D dataset of 20,000 32x32 Gaussian random fields with l U = 0.5 and ν U = 1.5. To train the OPFLOW prior, we use a GP latent space with l A = 0.1, ν A = 0.5.In this experiment, the posterior can be computed analytically. In Fig. 3 , we perform inference with only 32 random spatial observations. OPFLOW reliably recovers the true posterior mean and variance without assuming that the true data are from a Gaussian process a priori. However, OPFLOW does have a large amount of error in the upper right corner where there are few observations. Seismic waveforms regression. In this experiment, OPFLOW regression is applied to a real-world non-Gaussian dataset consisting of earthquake ground motion records from Japan [14] . We downsampled the time-series to 10Hz and used l A = 0.05, ν U = 0.5 for the latent space A. After training, we perform inference on an unseen time series and show results from UFR (Fig. 4 ). OPFLOW is able to generate samples that match the data and can accurately capture key characteristics in the data including various types of seismic waves arriving at different times. Tthe standard GP model struggles with the complexity of the seismic data. 
conclusion: 
 We introduced Neural Operator Flows (OPFLOW), an invertible infinite-dimensional generative model that offers a learning-based solution for end-to-end universal functional regression, which has the potential to outperform functional regression with Gaussian processes. Looking ahead, OPFLOW may enable new avenues for analyzing functional data by learning directly from the data.  
 
 Appendix:  



Reviews: review: title: title: Official review for paper 102
main_review: main_review: Overall, the paper lacks clarity and structure, making it difficult to follow the main contributions. The method is barely explained, and the authors provide only a simple diagram, which is poorly described. Without a clear description of the model, it is difficult to see where Equation 2 comes from; it would be helpful to have a proper explanation of it. Moreover, as far as I understand, the model seems to be a transformed Gaussian Process (GP) with normalizing flows, a concept already explored by Maroña et al. (2021). The authors should have discussed how their model relates to this previous work and highlight any differences.

All the figures were barely readable; I had to zoom in significantly to understand the images.
title: title: A potential contribution to functional regression in machine learning
main_review: main_review: This paper proposes an infinite-dimensional extension of normalizing flows, Neural Operator Flows, that map to Gaussian process base processes. The authors provide a framework for Bayesian functional regression with priors learned via Neural Operator Flows (akin to posterior inference in a Gaussian process in which the prior hyperparameters have been tuned via marginal likelihood). The proposed method relies on MAP estimation and SGLD to sample from the posterior. 

The work is important in that flexible functional prior specification and functional inference could improve on current methods for approximate Bayesian inference in machine learning. The paper is mostly clear, but more detail on OpFlow (particularly how domain partitioning works on $\mathbb{R}^d$) would be beneficial. Furthermore, the originality of the paper is not entirely clear; the related work section should probably touch on some related papers using normalizing flows and GPs (e.g.,  https://arxiv.org/abs/2011.01596) and on other functional regression approaches (e.g., functional variational inference, such as https://arxiv.org/abs/1903.05779). The examples presented are somewhat limited, and the seismic example does not in my view contribute much to the paper, as it’s clear that a stationary GP will underperform and it's not clear that OpFlow performs exceedingly well in terms of usual metrics like MSE of the mean prediction or coverage of uncertainty intervals (even though the result sort of “looks like” seismic data).

Some questions for the authors to consider: 

1. The presented examples all use Matern kernels for the base GP; how much does the choice of GP in the latent space affect the results? 
2. Can you elaborate on connections to other literature involving GPs and normalizing flows, including https://arxiv.org/abs/2011.01596 ? 
3. None of the results really show “extrapolation uncertainty”; how does the model perform compared to GPs when extrapolating outside the observed data points? 
4. How well does the proposed method achieve coverage via its uncertainty intervals?