Flow Annealed Kalman Inversion for Gradient-Free Inference in Bayesian Inverse Problems ^†

Abstract

For many scientific inverse problems, we are required to evaluate an expensive forward model. Moreover, the model is often given in such a form that it is unrealistic to access its gradients. In such a scenario, standard Markov Chain Monte Carlo algorithms quickly become impractical, requiring a large number of serial model evaluations to converge on the target distribution. In this paper, we introduce Flow Annealed Kalman Inversion (FAKI). This is a generalization of Ensemble Kalman Inversion (EKI) where we embed the Kalman filter updates in a temperature annealing scheme and use normalizing flows (NFs) to map the intermediate measures corresponding to each temperature level to the standard Gaussian. Thus, we relax the Gaussian ansatz for the intermediate measures used in standard EKI, allowing us to achieve higher-fidelity approximations to non-Gaussian targets. We demonstrate the performance of FAKI on two numerical benchmarks, showing dramatic improvements over standard EKI in terms of accuracy whilst accelerating its already rapid convergence properties (typically in $\mathcal{O}\left(10\right)$ steps).

1. Introduction

Many scientific inference tasks are concerned with inverse problems of the form

$$y=\mathcal{G}\left(x\right)+\eta ,$$

(1)

where $y\in {\mathbb{R}}^{d_y}$ are the data, $x\in {\mathbb{R}}^d$ are the model parameters, $\mathcal{G}$ is the forward map, and $\eta$ is the observation noise. Throughout this work, we will assume that we do not have access to gradients of $\mathcal{G}$ with respect to the parameters, and that $\eta \sim \mathcal{N}(0,\Gamma )$, where $\Gamma$ is a fixed noise covariance. The assumption of additive Gaussian noise is the standard setting for Ensemble Kalman Inversion (EKI) [1,2,3,4,5,6,7,8], and whilst we are restricted to problems with Gaussian likelihoods, this covers a large family of scientific inverse problems. The goal of the Bayesian inverse problem is then to recover the posterior distribution over the model parameters given our observations, $p\left(x\right|y)$.

Typical gradient-free inference methods often involve some variant on Markov Chain Monte Carlo (MCMC) algorithms, e.g., random walk Metropolis [9,10,11] or Sequential Monte Carlo (SMC) [12]. However, these methods typically require $\gtrsim {10}^3$ serial model evaluations to achieve convergence, making them intractable for problems with expensive forward models. In contrast, EKI utilizes embarrassingly parallel model evaluations to update parameter estimates, typically converging to an approximate solution in $\mathcal{O}\left(10\right)$ iterations [2,6,7,8].

EKI leverages ideas originally developed in the context of Ensemble Kalman Filtering (EKF) for data assimilation [13]. Since its development, EKI has seen applications across a range of disciplines, including studies of fluid flow [14], climate models [15], and machine learning tasks [16]. EKI can be understood in the context of annealing, where we seek to move from the prior to the posterior through a sequence of intermediate measures. In standard EKI, this involves constructing a sequence of Gaussian approximations to the intermediate measures. In the regime where we have a Gaussian prior ${\pi }_0\left(x\right)=\mathcal{N}(m_0,C_0)$ and a linear forward model $\mathcal{G}\left(x\right)=Gx$, the particle distribution obtained via EKI converges to the true posterior in the limit where the ensemble size $J\to \infty$. However, outside this linear Gaussian regime, EKI is an uncontrolled approximation to the posterior that is constructed on the basis of matching first and second moments of the target distribution. Nonetheless, EKI has been shown to perform well on problems with non-linear forward models and slightly non-Gaussian targets [1,2,6].

In this work, we propose the application of normalizing flows (NFs) [17,18,19,20] to relax the Gaussian ansatz made by standard EKI for the intermediate measures. Instead of assuming a Gaussian particle distribution at each iteration, the NF is used to fit the empirical particle distribution and map it to a Gaussian latent space, where the EKI updates are performed. Thus, we are better able to capture non-Gaussian target geometries. The structure of this paper is as follows: in Section 2 we describe the Flow Annealed Kalman Inversion (FAKI) algorithm; in Section 3, we demonstrate the performance of the method on two Bayesian inference tasks with non-Gaussian target geometries; and we summarize our work in Section 4.

2. Methods

2.1. Regularized Ensemble Kalman Inversion

A number of versions of EKI have been proposed in the literature. Of interest here is the regularized, perturbed observation form of EKI [6]. Starting with an ensemble of particles drawn from the prior, ${\left\{x_0^j\right\}}_{j=1}^J$, the particles are updated at each iteration according to

$$x_{n+1}^j=x_n^j+C_n^{x\mathcal{G}}{(C_n^{\mathcal{G}\mathcal{G}}+{\alpha }_n\Gamma )}^{-1}(y-\mathcal{G}\left(x_n^j\right)+\sqrt{{\alpha }_n}{\xi }_n^j).$$

(2)

The empirical covariances $C_n^{x\mathcal{G}}$ and $C_n^{\mathcal{G}\mathcal{G}}$ are given by

$$\begin{array}{cc}\hfill C_n^{x\mathcal{G}}& =\frac{1}{J-1}\sum _{j=1}^J(x_n^j-\langle x_n\rangle )\otimes (\mathcal{G}\left(x_n^j\right)-\langle {\mathcal{G}}_n\rangle ),\hfill \end{array}$$

(3)

$$\begin{array}{cc}\hfill C_n^{\mathcal{G}\mathcal{G}}& =\frac{1}{J-1}\sum _{j=1}^J(\mathcal{G}\left(x_n^j\right)-\langle {\mathcal{G}}_n\rangle )\otimes (\mathcal{G}\left(x_n^j\right)-\langle {\mathcal{G}}_n\rangle ).\hfill \end{array}$$

(4)

At each iteration, we perturb the forward model evaluations with the Gaussian observation noise ${\xi }_n^j\sim \mathcal{N}(0,\Gamma )$. The parameter ${\alpha }_n$ is a Tikhonov regularization parameter, which can be viewed as an inverse step size in the Bayesian annealing context. In particular, given a set of annealing parameters ${\beta }_0\equiv 0<{\beta }_1<\dots <{\beta }_N<{\beta }_{N+1}\equiv 1$, we have the corresponding set of target distributions

$${\pi }_n\left(x\right)\propto {\pi }_0\left(x\right)exp\left(-\frac{{\beta }_n}{2}{{\Gamma }^{-1/2}(y-\mathcal{G}\left(x\right))}^2\right),$$

(5)

with

$${\alpha }_n={\beta }_{n+1}-{\beta }_n.$$

(6)

EKI proceeds by constructing a sequence of ensemble approximations to Gaussian distributions that approximate the intermediate targets.

The choice of the regularization parameter, ${\alpha }_n$, controls the transition from the prior to the posterior. Previous proposals for an adaptive choice have taken inspiration from SMC by using a threshold on the effective sample size (ESS) of the particles at each temperature level [21,22]. In this work, we adopt the same approach, calculating pseudo-importance weights at each temperature given by

$$w_n^j=exp\left(-\frac{1}{2}({\beta }_{n+1}-{\beta }_n)\Vert {{\Gamma }^{-1/2}(y-\mathcal{G}\left(x_n^j\right))\Vert }^2\right).$$

(7)

The next temperature level can then be selected by solving

$${\left(\sum _{j=1}^J{w}_n^j{\left({\beta }_{n+1}\right)}^2\right)}^{-1}{\left(\sum _{j=1}^J{w}_n^j\left({\beta }_{n+1}\right)\right)}^2=\tau J,$$

(8)

using the bisection method, where $0<\tau <1$ is the target fractional ESS threshold. Throughout our work we set $\tau =0.5$. The full pseudo-code for EKI is given in Algorithm 1.

Algorithm 1 Ensemble Kalman Inversion

1: Input: J prior samples ${\{x_0^j\sim {\pi }_0\left(x\right)\}}_{j=1}^J$, data y, observation error covariance $\Gamma$ and fractional ESS target threshold $\tau$ 2: Initialize inverse temperature ${\beta }_0=0$, iteration counter $n=0$ 3: while  $\beta <1$ do 4:    Evaluate ${\mathcal{G}}_j=\mathcal{G}\left(x_n^j\right)$, $j\in \{1,\dots ,J\}$ 5:    Solve for ${\beta }_{n+1}$ using the bisection method with Equation (8) 6:    ${\alpha }_n\leftarrow {\beta }_{n+1}-{\beta }_n$ 7:    Update particles using Equation (2) to obtain ${\left\{x_{n+1}^j\right\}}_{j=1}^J$ 8:    $n\leftarrow n+1$ 9: end while 10: Output: Converged particle ensemble ${\left\{x_N^j\right\}}_{j=1}^J$

2.2. Normalizing Flows

As discussed above, standard EKI proceeds by constructing a sequence of ensemble approximations to Gaussian distributions. The procedure works well in the situation where the target and all the intermediate measures are close to Gaussian. However, when any of these measures are far from Gaussian, EKI can dramatically fail to capture the final target geometry.

To address this shortcoming, we propose the use of NFs to approximate each intermediate target, instead of using the Gaussian ansatz of standard EKI. NFs are powerful generative models that can be used for flexible density estimation and sampling [17,18,19,20]. An NF model maps from the original space $x\in {\mathbb{R}}^d$ to a latent space $z\in {\mathbb{R}}^d$ through a sequence of invertible transformations $f=f_1\circ f_2\circ \dots \circ f_L$, such that we have a bijective mapping $z=f\left(x\right)$. The mapping is such that the latent variables are mapped to some simple base distribution, typically chosen to be the standard Normal distribution, giving $z\sim p_z\left(z\right)=\mathcal{N}(0,I)$.

The NF density can be evaluated through the change of variables formula,

$$q\left(x\right)=p_z\left(f\left(x\right)\right)\left|\det Df\left(x\right)\right|=p_z\left(f\left(x\right)\right)\prod _{l=1}^L\left|\det D{f}_l\left(x\right)\right|,$$

(9)

where $Df\left(x\right)=\partial f\left(x\right)/\partial x$ denotes the Jacobian of f. The efficient evaluation of this density requires the Jacobian of the transformation to be easy to evaluate, and efficient sampling requires the inverse of the mapping f to be easy to calculate. In this work, we use Masked Autoregressive Flows (MAFs) [18], which have previously been found to perform well in the context of preconditioned MCMC sampling within SMC without the need for expensive hyper-parameter searches during sampling [23].

2.3. Flow Annealed Kalman Inversion

Given particles distributed as ${\pi }_n\left(x\right)$, the subsequent target can be written as

$${\pi }_{n+1}\left(x\right)\propto {\pi }_n\left(x\right)exp\left(-\frac{1}{2{\alpha }_n}{\Vert {\Gamma }^{-1/2}(y-\mathcal{G}\left(x\right))\Vert }^2\right).$$

(10)

We may therefore view ${\pi }_n\left(x\right)$, i.e., the posterior at the temperature level ${\beta }_n$, as an effective prior for ${\pi }_{n+1}\left(x\right)$, with a data likelihood annealed by ${\alpha }_n^{-1}$. By fitting an NF to the particles ${\left\{x_n^j\right\}}_{j=1}^J$, we obtain an approximate map from the intermediate target ${\pi }_n\left(x\right)$ to $\mathcal{N}\left(z\right|0,I)$. The latent space target is then given by the change of variables formula as

$${\pi }_{n+1}\left(z\right)={\pi }_{n+1}(x=f_n^{-1}\left(z\right))\left|\det D{f}_n^{-1}\left(z\right)\right|.$$

(11)

By controlling the choice of ${\alpha }_n$, we control the distance between the Gaussianized effective prior and this latent space target density. For FAKI, we therefore perform the EKI updates in the NF latent space at each temperature level, allowing us to relax the Gaussian ansatz of standard EKI by constructing an approximate map from each ${\pi }_n\left(x\right)$ to a Gaussian latent space. It is worth noting that, whilst this method relaxes the Gaussianity assumptions of standard EKI, it does not address the linearity assumptions used in deriving EKI.

The FAKI update for the latent space particle locations is given by

$$z_{n+1}^j=z_n^j+{\mathcal{C}}_n^{z\mathcal{G}}{({\mathcal{C}}_n^{\mathcal{G}\mathcal{G}}+{\alpha }_n\Gamma )}^{-1}(y-\mathcal{G}\left(f_n^{-1}\left(z_n^j\right)\right)+\sqrt{{\alpha }_n}{\xi }_n^j),$$

(12)

where the latent space empirical covariances are given by

$$\begin{array}{cc}\hfill {\mathcal{C}}_n^{z\mathcal{G}}& =\frac{1}{J-1}\sum _{j=1}^J(z_n^j-\langle z_n\rangle )\otimes (\mathcal{G}\left(f_n^{-1}\left(z_n^j\right)\right)-\langle {\mathcal{G}}_n\rangle ),\hfill \end{array}$$

(13)

$$\begin{array}{cc}\hfill {\mathcal{C}}_n^{\mathcal{G}\mathcal{G}}& =\frac{1}{J-1}\sum _{j=1}^J(\mathcal{G}\left(f_n^{-1}\left(z_n^j\right)\right)-\langle {\mathcal{G}}_n\rangle )\otimes (\mathcal{G}\left(f_n^{-1}\left(z_n^j\right)\right)-\langle {\mathcal{G}}_n\rangle ).\hfill \end{array}$$

(14)

The full pseudocode for FAKI is given in Algorithm 2.

Algorithm 2 Flow Annealed Kalman Inversion

1: Input: J prior samples  ${\{x_0^j\sim {\pi }_0\left(x\right)\}}_{j=1}^J$, data y, observation error covariance $\Gamma$ and fractional ESS target threshold $\tau$ 2: Initialize inverse temperature ${\beta }_0=0$, iteration counter $n=0$ 3: while  $\beta <1$ do 4:    Evaluate ${\mathcal{G}}_j=\mathcal{G}\left(x_n^j\right)$, $j\in \{1,\dots ,J\}$ 5:    Solve for ${\beta }_{n+1}$ using the bisection method with Equation (8) 6:    ${\alpha }_n\leftarrow {\beta }_{n+1}-{\beta }_n$ 7:    Fit NF map $f_n$ to current samples ${\left\{x_n^j\right\}}_{j=1}^J$ 8:    Map particles to latent space $z_n^j=f_n\left(x_n^j\right)$, $j\in \{1,\dots ,J\}$ 9:    Update particles using Equation (12) to obtain ${\left\{z_{n+1}^j\right\}}_{j=1}^J$ 10:    Map back to the data space $x_{n+1}^j=f_n^{-1}\left(z_{n+1}^j\right)$, $j\in \{1,\dots ,J\}$ 11:    $n\leftarrow n+1$ 12: end while 13: Output: Converged particle ensemble ${\left\{x_N^j\right\}}_{j=1}^J$

3. Results

In this section, we demonstrate the performance of FAKI compared to standard EKI on two numerical benchmarks, a two-dimensional Rosenbrock distribution and a stochastic Lorenz system [24,25]. Both models display significant non-Gaussianity at some point during the transition from prior to posterior, severely frustrating the performance of EKI. This manifests in both the reduced fidelity of the final ensemble approximations to the posterior and a larger number of iterations being required for convergence following the ESS-based annealing scheme described in Section 2.1.

In

Table 1. Median and MAD values of $N_{\mathrm{iter}}$ and 1-Wasserstein distances for each model and algorithm combination, calculated over 10 independent runs using different random seeds. For both numerical benchmarks, we see that FAKI resulted in a reduced number of iterations for convergence and a lower value of the 1-Wasserstein distance between the converged samples and the ground truth.

Model	Algorithm	$\mathbf{Median}\left({\mathit{N}}_{\mathbf{iter}}\right)$	$\mathbf{MAD}\left({\mathit{N}}_{\mathbf{iter}}\right)$	$\mathbf{Median}\left({\mathit{W}}_1\right)$	$\mathbf{MAD}\left({\mathit{W}}_1\right)$
Rosenbrock	EKI	100	7.0	0.72	0.05
Rosenbrock	FAKI	34.0	7.0	0.43	0.14
Lorenz	EKI	10.0	0.0	69.8	1.08
Lorenz	FAKI	8.0	0.0	5.65	0.86

, we provide statistics summarizing the performance of EKI and FAKI on our numerical benchmarks. We measured the quality of the posterior approximations by computing the 1-Wasserstein distance, $W_1$ [26,27], between the samples obtained through FAKI and EKI against reference posterior samples obtained via long runs of the Hamiltonian Monte Carlo (HMC) algorithm [28,29]. These reference samples were thinned to be approximately independent when computing the 1-Wasserstein distances (we use the Python Wasserstein library: https://github.com/pkomiske/Wasserstein/ (accessed on 30 June 2023)). The 1-Wasserstein distance may be interpreted as the cost involved in rearranging one probability measure to look like another, with lower values indicating the two probability measures are closer to one another. In addition to this assessment of the approximation quality, we report the number of iterations, $N_{\mathrm{iter}}$, required by FAKI and EKI for convergence. For both quantities, we report the median and median absolute deviation (MAD) estimated over 10 independent runs using different random seeds.

3.1. $d=2$ Rosenbrock

In our first numerical experiment, we considered the two-dimensional Rosenbrock distribution. This toy model allowed us to clearly see the impact of non-Gaussianity on the performance of EKI and how FAKI is able to alleviate these issues. For the Rosenbrock model, we assumed a Gaussian prior over the parameters $x\in {\mathbb{R}}^2$,

$$x\sim \mathcal{N}(0,{10}^{2}I).$$

(15)

The data, $y\in {\mathbb{R}}^2$, were distributed according to the likelihood

$$y\sim \mathcal{N}(\mathcal{G}\left(x\right)={(x_1-x_0^2,x_0)}^{⊺},\Gamma =\mathrm{diag}({0.01}^2,1^2)).$$

(16)

To generate simulated data, we evaluated $y=G\left({(1,1)}^{⊺}\right)+\eta$, where $\eta \sim \mathcal{N}(0,\mathrm{diag}({0.01}^2,1^2))$. The large difference in noise scales resulted in a highly non-Gaussian posterior geometry that posed a significant challenge for EKI. For each run of EKI and FAKI, we used 100 particles.

In Figure 1, we show pair plots comparing the final particle distributions obtained with EKI and FAKI against samples obtained through a long run of HMC. The NF mapping meant that the ensemble approximation obtained by FAKI was able to capture the highly non-linear target geometry. In comparison, EKI struggled to fill the tails of the Rosenbrock target. Moreover, whilst FAKI converged within ∼34 iterations, EKI required a median number of ∼100 iterations to converge using the ESS-based annealing scheme.

3.2. Stochastic Lorenz System

The Lorenz equations are a set of coupled differential equations used as a simple model of atmospheric convection. Notably, for certain parameter values, the Lorenz equations are known to exhibit chaotic behavior [24]. In this work, we followed [25] and considered the stochastic Lorenz system

$$\begin{array}{cc}\hfill \mathrm{d}{X}_t& =10(Y_t-X_t)\mathrm{d}t+\mathrm{d}{W}_t^x,\hfill \end{array}$$

(17)

$$\begin{array}{cc}\hfill \mathrm{d}{Y}_t& =X_t(28-Z_t)\mathrm{d}t-Y_t\mathrm{d}t+\mathrm{d}{W}_t^y,\hfill \end{array}$$

(18)

$$\begin{array}{cc}\hfill \mathrm{d}{Z}_t& =X_t{Y}_t\mathrm{d}t-\frac{8}{3}{Z}_t\mathrm{d}t+\mathrm{d}{W}_t^z,\hfill \end{array}$$

(19)

where $W_t^x$, $W_t^y$, and $W_t^z$ are Gaussian white noise processes with standard deviation ${\sigma }_0=0.1$. To generate simulated, data we integrated these equations using a Euler–Maruyama scheme with $\mathrm{d}t=0.02$ for 30 steps, with initial conditions $X_0,Y_0,Z_0\sim \mathcal{N}(0,1^2)$. The observations were then taken to be the $X_t$ values over these 30 time steps, with Gaussian observational noise ${\eta }_t\sim \mathcal{N}(0,{\sigma }^2=1^2)$.

The goal of our inference here was to recover the initial conditions, the trajectories $(X_t,Y_t,Z_t)$ and the innovation noise scale ${\sigma }_0$, giving a parameter space of $d=94$ dimensions. We assigned priors over these parameters as

$$\begin{array}{cc}\hfill log{\sigma }_0& \sim \mathcal{N}(-1,1^2),\hfill \end{array}$$

(20)

$$\begin{array}{cc}\hfill X_0,Y_0,Z_0& \sim \mathcal{N}(0,1^2),\hfill \end{array}$$

(21)

$$\begin{array}{cc}\hfill X_t& \sim \mathcal{N}(X_{t-1}+f_X(X_{t-1},Y_{t-1},Z_{t-1},t-1)\mathrm{d}t,{\sigma }_0^2\mathrm{d}t),t\in \{1,\dots ,30\},\hfill \end{array}$$

(22)

$$\begin{array}{cc}\hfill Y_t& \sim \mathcal{N}(Y_{t-1}+f_Y(X_{t-1},Y_{t-1},Z_{t-1},t-1)\mathrm{d}t,{\sigma }_0^2\mathrm{d}t),t\in \{1,\dots ,30\},\hfill \end{array}$$

(23)

$$\begin{array}{cc}\hfill Z_t& \sim \mathcal{N}(Z_{t-1}+f_Z(X_{t-1},Y_{t-1},Z_{t-1},t-1)\mathrm{d}t,{\sigma }_0^2\mathrm{d}t),t\in \{1,\dots ,30\},\hfill \end{array}$$

(24)

where $f_X$, $f_Y$, and $f_Z$ are the transition functions corresponding to Equations (17)–(19), respectively. The Gaussian likelihood has the form

$${\widehat{X}}_t\sim \mathcal{N}(X_t,{\sigma }^2),t\in \{1,\dots ,30\},$$

(25)

where ${\widehat{X}}_t$ are the observations of the $X_t$ trajectory. The chaotic dynamics of the Lorenz system resulted in a highly non-Gaussian prior distribution, with the inversion having to proceed through a sequence of highly non-Gaussian intermediate measures towards the posterior. This severely frustrated the performance of EKI, with the Gaussian ansatz failing to describe the geometry of the intermediate measures. For each run of EKI and FAKI, we used 940 particles.

In Figure 2, we show the ensemble estimates for the mean and standard deviation along each dimension obtained by EKI and FAKI, compared to reference estimates obtained through long runs of HMC. FAKI was able to obtain accurate mean estimates along each dimension, whereas EKI was unable to obtain the correct means for much of the $Z_t$ trajectory. EKI severely overestimated the marginal standard deviations along many dimensions. This situation was alleviated by the NF mappings learned by FAKI. The greater fidelity of the FAKI posterior approximations is reflected in the median estimates for the 1-Wasserstein distances, with a value of $5.65$ for FAKI and $69.8$ for EKI.

4. Conclusions

In this work, we introduced Flow Annealed Kalman Inversion (FAKI), a gradient-free inference algorithm for Bayesian inverse problems with expensive forward models. This is a generalization of Ensemble Kalman Inversion (EKI) where we utilize normalizing flows (NFs) to replace the Gaussian ansatz made in EKI. Instead of constructing a sequence of ensemble approximations to Gaussian measures that approximate a sequence of intermediate measures, as we move from the prior to the posterior, we learn an NF mapping at each iteration to a Gaussian latent space. Provided the transition between temperature levels is controlled, we can perform Kalman inversion updates in the NF latent space. In the NF latent space, the Gaussianity assumptions of EKI are more closely satisfied, resulting in a more stable inversion at each temperature level.

We demonstrated the performance of FAKI on two numerical benchmarks, a $d=2$ Rosenbrock distribution and a $d=94$ stochastic Lorenz system. Both examples exhibited significant non-Gaussianity in the transition from prior to posterior that frustrated standard EKI. In the presence of strong non-Gaussianity, we found that FAKI produced higher-fidelity posterior approximations compared to EKI, as measured by the 1-Wasserstein distance between FAKI/EKI samples and reference HMC samples. In addition to the improved fidelity of the posterior approximations, we found that FAKI tended to reduce the number of iterations required for convergence.

Whilst the application of NFs is able to relax the Gaussian ansatz of EKI, it does not address the linearity assumptions used in deriving EKI. As such, FAKI is still not exact for general forward models. In future work, it will be interesting to explore methods to address this: for example, the combination of FAKI with unbiased MCMC or importance sampling methods. It would also be interesting to consider generalizations of FAKI that are able to accommodate non-Gaussian likelihoods and/or parameter-dependent noise covariances. The use of NFs means that we typically require ensemble sizes $J\gtrsim 10d$ to learn accurate NF maps with the MAF architecture employed in this work. It would be useful to explore alternative NF architectures and regularization schemes that are able to learn accurate NF maps with smaller ensemble sizes, in order to enable FAKI to scale to higher dimensions. In this work, we found that the MAF architecture is able to capture a wide range of target geometries without the need for expensive NF hyper-parameter searches. However, it may be possible to exploit NF architectures with inductive biases that are particularly suited to common target geometries, e.g., the non-linear correlations that often appear in hierarchical models.