SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jin Xu $^{1*}$  Emilien Dupont $^{1\dagger}$  Kaspar Märtens $^{2}$  Tom Rainforth $^{1}$  Yee Whye Teh $^{1\dagger}$
$^{1}$  Department of Statistics, University of Oxford, UK.  $^{2}$  Big Data Institute, University of Oxford, UK.
We introduce Markov Neural Processes (MNPs), a new class of Stochastic Processes (SPs) which are constructed by stacking sequences of neural parameterised Markov transition operators in function space. We prove that these Markov transition operators can preserve the exchangeability and consistency of SPs. Therefore, the proposed iterative construction adds substantial flexibility and expressivity to the original framework of Neural Processes (NPs) without compromising consistency or adding restrictions. Our experiments demonstrate clear advantages of MNPs over baseline models on a variety of tasks.
Stochastic Processes (SPs) are widely used in many scientific disciplines, including biology [3], chemistry [57], neuroscience [34], physics [45] and control theory [1]. They are formed by a (typically infinite) collection of random variables and can be used to model data by considering the conditional distribution of target variables given observed context variables. In machine learning, SPs in the form of Bayesian nonparametric models—such as Gaussian Processes (GPs) [47] and Dirichlet processes [56]—are used in tasks such as regression, classification, and clustering. SPs parameterised by neural networks have also been used for meta-learning [21, 60] and generative modelling [15, 39].
With the increasing amount of data available, and the complex patterns arising in many applications, more flexible and scalable SP models with greater learning capacity are required. The Neural Process (NP) family [18, 21, 22, 25, 30] meets this demand by parameterising SPs with neural networks, and enjoys greater flexibility and computational efficiency compared to traditional nonparametric models.
Unfortunately, the original version of NPs [22] lacks expressivity and often underfits the data in practice [30]. Various extensions have therefore been proposed—such as Attentive Neural Processes (ANPs) [30], Convolutional Neural Processes (CONVNPs) [18, 25] and Gaussian Neural Processes (GNPs) [4]—to improve expressivity. However, these models—which we refer to as predictive SPs—are based around directly constructing mappings from contexts to predictive distributions, and therefore forgo the fully generative nature of the original NP formulation. This can be problematic as it means that they are no longer consistent under conditioning: their predictive distributions no longer correspond to the conditional distribution of an underlying SP prior. In turn, this can cause a variety of issues, such as conflicts or inconsistencies between predictions and miscalibrated uncertainties.
To address these issues, we propose an alternative mechanism to extend the NP family and provide increased expressivity while maintaining their original fully generative nature and consistency. We
*Corresponding author: <jin.xu@stats.ox.ac.uk>  
†ED is now at Google DeepMind; this work was done while ED was at Oxford.  
‡YWT is at both Google DeepMind and Oxford; this work was done at Oxford.
Step 0 (trivial SP)
Step 1
Figure 1: MNPs construct expressive SPs via iterative transformations. Here we use MNPs to model distributions over monotonic functions. We start with trivial initial SPs and gradually transform them into more complex SPs conditioned on observed context points marked as black crosses. The blue lines represent sampled mean functions and the shaded region indicates one standard deviation.
Step 3 (final SP)
begin by generalising the marginal density functions of NPs. The generalised form can be viewed as Markov transition operators with functions as states. This lays the foundation for stacking these operators to construct more powerful generative SP models that we call Markov Neural Processes (MNPs). MNPs can be seen as Markov chains in function space, parameterised by neural networks: they make iterative transformations from simple initial SPs into more flexible, yet properly defined, SPs (as illustrated in Figure 1), without compromising consistency or introducing additional assumptions.
To empirically demonstrate the value of our proposed approach, we first benchmark MNPs against baselines on 1D function regression tasks, and conduct ablation studies in this controlled setting. We then show that they can be used as a high-performing surrogate for contextual bandit problems. Finally, we apply them to geological data, for which they demonstrate encouraging performance.
A SP is a (typically infinite) collection of random variables defined on a common probability space. We can consider a SP as a random function  $F: \mathcal{X} \to \mathcal{Y}$  where inputs can be regarded as indexing the output random variables. With a relaxed use of notation, we employ  $p(f)$  in denoting a SP, where  $f$  maps inputs  $x \in \mathcal{X}$  to outputs  $y \in \mathcal{Y}$ . Kolmogorov's Consistency Theorem shows that a SP can be indirectly defined via a collection of marginal distributions,  $\{p_{x_{1:n}}(y_{1:n})\}_{x_{1:n} \in \mathcal{X}^n}$  (we drop  $\mathcal{X}^n$  from here on for conciseness) if they, for any permutation  $\pi$  and all possible sets of inputs  $x_{1:n} \in \mathcal{X}^n$ , satisfy the exchangeability condition:
$$
p _ {x _ {1: n}} \left(y _ {1: n}\right) = p _ {\pi \left(x _ {1: n}\right)} \left(\pi \left(y _ {1: n}\right)\right) := p _ {x _ {\pi (1)}, \dots , x _ {\pi (n)}} \left(y _ {\pi (1)}, \dots , y _ {\pi (n)}\right), \tag {1}
$$
and the (marginal) consistency condition:
$$
p _ {x _ {1: m}} \left(y _ {1: m}\right) = \int p _ {x _ {1: n}} \left(y _ {1: n}\right) \mathrm {d} y _ {m + 1: n} \quad \forall 1 \leq m <   n. \tag {2}
$$
If none of the random variables in a SP are observed, we call it a prior SP. For any two distinct subsets of datapoints  $\mathcal{C} = \{(x_i,y_i)\}_{i = 1}^m$  (the context) and  $\mathcal{T} = \{(x_i,y_i)\}_{i = m + 1}^n$  (the target), one can use this prior SP to compute the conditional density  $p_{x_{m + 1:n}|x_{1:m}}(y_{m + 1:n}|y_{1:m})$  via Bayesian inference. If inference is exact, it can be proved that the collection of conditional distributions  $\{p_{x_{m + 1:n}|x_{1:m}}(y_{m + 1:n}|y_{1:m})\}_{x_{m + 1:n}}$  also satisfy exchangeability and consistency, and hence define a valid posterior SP,  $p(f|\mathcal{C})$ . We discuss the impact of approximate inference in Section 4.3.
Based on a conditional version of de Finetti's Theorem, a NP [22] defines a SP by providing a collection of exchangeable and consistent marginal densities parameterized by neural networks. Specifically, they set up their marginal densities to have the form:
$$
p _ {x _ {1: n}} \left(y _ {1: n}; \theta\right) = \int p _ {\theta} (z) \prod_ {i = 1} ^ {n} \mathcal {N} \left(y _ {i} \mid \mu_ {\theta} \left(x _ {i}, z\right), \sigma_ {\theta} ^ {2} \left(x _ {i}, z\right)\right) d z, \tag {3}
$$
where  $z$  is a latent variable which captures dependencies across different input locations,  $y_{i} \coloneqq f(x_{i})$ ,  $\mu_{\theta}$  and  $\sigma_{\theta}$  are deep neural networks, and  $p_{\theta}(z)$  is a, typically Gaussian, prior distribution on the latents. Note that this form is more general than the one in [22] where both  $p_{\theta}$  and  $\sigma_{\theta}$  are not learnt.
3 Generative versus predictive stochastic process models
Before introducing our MNP approach, we first delve into the distinctions between conventional, consistent, SP models like NPs, and the predictive SP models corresponding to popular NP extensions, such as ANPs, CONVNPs, and GNPs, highlighting some of the potential drawbacks with the latter.
We introduce the term generative SP model to refer to the standard case where one first specifies a prior SP,  $p(f)$ , and then relies on Bayesian inference to compute posterior a SP,  $p(f|\mathcal{C})$ . The context and the target are only distinguished for inference, and they are treated equally in model specification. Traditional nonparametric models such as GPs, Student-t processes and the original NPs belong to this category. Since posterior SPs are inferred under the same prior SP, under the condition of exact inference, for two different contexts  $\mathcal{C} = \{(x_i,y_i)\}_{i\in \mathcal{C}}$  and  $\mathcal{C}' = \{(x_i,y_i)\}_{i\in \mathcal{C}'}$ , we have
$$
\int p (f | \mathcal {C}) p _ {x _ {\mathcal {C}}} \left(y _ {\mathcal {C}}\right) \mathrm {d} y _ {\mathcal {C}} = \int p (f | \mathcal {C} ^ {\prime}) p _ {x _ {\mathcal {C} ^ {\prime}}} \left(y _ {\mathcal {C} ^ {\prime}}\right) \mathrm {d} y _ {\mathcal {C} ^ {\prime}}, \tag {4}
$$
where  $x_{\mathcal{C}} \coloneqq \{x_i\}_{i \in \mathcal{C}}$  and  $x_{\mathcal{C}'} \coloneqq \{x_i\}_{i \in \mathcal{C}'}$  are taken as fixed, and  $p_{x_{\mathcal{C}}}(\mathcal{Y}_{\mathcal{C}})$  and  $p_{x_{\mathcal{C}'}}(y_{\mathcal{C}'})$  are defined through the prior  $\mathrm{SP} p(f)$ . We call this property conditional consistency to set apart from marginal consistency defined in Equation (2). Note that conditional consistency is a prerequisite for marginal consistency to hold for the prior  $p(f)$ : if the two integrals in Equation (4) are not equal, at least one must be different to the prior distribution  $p(f)$ .
By contrast, predictive SP models directly construct mappings from a context  $\mathcal{C}$  to a predictive SP  $p(f;\mathcal{C})$  and make a clear distinction between the context and the target in model specification. Consequently, for these models  $p(f;\mathcal{C})$  no longer corresponds to a posterior derived from a certain prior  $p(f)$ , so they no longer need to satisfy conditional consistency. As such, they no longer form a valid conditioning of a SP: though they are typically constructed to ensure  $p(f;\mathcal{C})$  is itself a valid SP for new evaluations of  $f$  with  $\mathcal{C}$  held fixed, this SP is no longer itself derived as the conditional of a prior SP. As such, predictive SP models are not consistent under updating and no longer update their uncertainties in a Bayesian manner as new data is observed. In short, they are no longer treating the data itself as being drawn from an SP.
Though predictive SP models have proven to be effective tools for meta-learning tasks, such as 1-D regression, image regression, and few-shot image classification [4, 18, 30, 60], this lack of consistent updating can be problematic in scenarios where the context is not fixed, such as when performing sequential updating. In particular, their uncertainties will be mismatched with how the model is updated in practice as new data is observed. This has the potential to be problematic for a wide variety of possible applications involving sequential decision-making, such as Bayesian experimental design [7, 46], active learning [2, 27, 51], Bayesian optimization [19, 23], and contextual bandit problems [35, 50]. There is also the potential for such models to fall foul of so-called Dutch Book scenarios [26], wherein the inconsistencies of the model can lead to conflicting predictions.
4 Markov Neural Processes
To correct the shortfalls of NPs while maintaining their conditional consistency, we now introduce a more expressive family of generative SP models termed Markov Neural Processes (MNPs). Our starting point is to extend NP density functions into a generalised form that can be viewed as a transition operator which transforms a trivial SP to a more flexible one. MNPs are then formed by stacking sequences of these transition operators to form a highly expressive and flexible model class. The training and inference procedure for MNP mirrors that of NP, but we also introduce a novel inference model that allows for efficient learning in our scenario.
4.1 A more general form of Neural Process density functions
Recall the form of the NP marginal density functions from Equation (3). One can draw joint samples from  $p_{x_{1:n}}(y_{1:n};\theta)$  via reparameterization using:
$$
y _ {1: n} ^ {(0)} \sim \mathcal {N} (\mathbf {0}, \mathbf {1}), \quad z \sim p _ {\theta} (z), \quad y _ {i} = \sigma_ {\theta} (z, x _ {i}) \cdot y _ {i} ^ {(0)} + \mu_ {\theta} (z, x _ {i}). \tag {5}
$$
The key starting point for MNPs is to show that this can be generalised to the case where  $y_{1:n}^{(0)}$  is drawn from any given SP of its own,  $p(f^{(0)})$ , and each  $y_i$  is any invertible transformation,  $F_\theta$ , of
$y_{i}^{(0)}$ , parameterized by  $x_{i}$  and  $z$ . Specifically, we introduce the following result (see Appendix A for proof).
Proposition 4.1. Let  $F_{\theta}(\cdot ;x,z):\mathcal{V}\mapsto \mathcal{V}$  denote an invertible transformation between outputs, parameterized by the input and latent. If  $\{p_{x_{1:n}}(y_{1:n}^{(0)})\}_{x_{1:n}}$  is a valid SP (i.e. it satisfies Equation (1) and Equation (2)) and
$$
y _ {1: n} ^ {(0)} \sim \left\{p _ {x _ {1: n}} \left(y _ {1: n} ^ {(0)}\right) \right\} _ {x _ {1: n}}, \quad z \sim p _ {\theta} (z), \quad y _ {i} = F _ {\theta} \left(y _ {i} ^ {(0)}; x _ {i}, z\right). \tag {6}
$$
then  $\{p_{x_{1:n}}(y_{1:n})\}_{x_{1:n}}$  is also a valid SP.
Our next step is to realise that Equation (6) can be viewed as a Markov transition in function space, denoted as  $p(f|f^{(0)})$ , which transforms a simpler SP  $p(f^{(0)})$  to a more expressive one  $p(f)$ . We can thus generalise things further by stacking sequences of Markov transition operators in function spaces (FMTOs), denoted by  $p(f^{(1)}|f^{(0)})$ , ...,  $p(f^{(T)}|f^{(T - 1)})$ , together to form a Markov chain  $f^{(0)}\to \dots \to f^{(T)}$  in function space. This will be the basis of MNPs.
4.2 Markov chains in function space
Analogously to defining a SPs through its marginals, we indirectly specify FMTOs using a collection of marginal Markov transition operators (MMTOs), denoted by  $\{p_{x_{1:n}}(y_{1:n}|y_{1:n}^{(0)})\}_{x_{1:n}}$ , where each MMTO is just Markov transition operator over a finite sequence of function outputs.
To adapt the definitions of consistency and exchangeability for SPs to the transition operator setting, we say that the MMTOs are consistent if and only if, for any  $1 < m < n$  and sequence  $x_{1:n} \in \mathcal{X}^n$ ,
$$
\int p _ {x _ {1: n}} \left(y _ {1: n} \mid y _ {1: n} ^ {(0)}\right) d y _ {m + 1: n} = p _ {x _ {1: m}} \left(y _ {1: m} \mid y _ {1: n} ^ {(0)}\right) = p _ {x _ {1: m}} \left(y _ {1: m} \mid y _ {1: m} ^ {(0)}\right). \tag {7}
$$
Similarly, MMTOs are exchangeable if, and only if, for all possible permutations  $\pi$ ,
$$
p _ {x _ {1: n}} \left(y _ {1: n} \mid y _ {1: n} ^ {(0)}\right) = p _ {\pi \left(x _ {1: n}\right)} \left(\pi \left(y _ {1: n}\right) \mid \pi \left(y _ {1: n} ^ {(0)}\right)\right). \tag {8}
$$
Note that, if we consider  $(x_{i},y_{i}^{(0)})$  as inputs, and  $y_{i}$  as function outputs, the transition operator  $p_{x_{1:n}}(y_{1:n}|y_{1:n}^{(0)})$  can be seen as the finite marginals of a random function  $F^{\prime}:\mathcal{X}\times \mathcal{Y}\to \mathcal{Y}$  whose distribution is a SP, and the conditions (7) and (8) correspond to (2) and (1).
Provided that these MMTOs are consistent and exchangeable, the FMTOs in function space will also be well-defined indirectly—i.e. the transition produces a well-defined SP given input random functions from a SP—as per the following result (see Appendix A for proof):
Proposition 4.2. If the collection of marginals before transition  $\{p_{x_{1:n}}(y_{1:n}^{(0)})\}_{x_{1:n}}$  is consistent and exchangeable (as per Equations (1) and (2)) and the collection of MMTOs  $\{p_{x_{1:n}}(y_{1:n}|y_{1:n}^{(0)})\}_{x_{1:n}}$  is also consistent and exchangeable (as per Equations (7) and (8)), then the collection of marginals after transition  $\{p_{x_{1:n}}(y_{1:n})\}$  is also consistent and exchangeable, hence defining a valid SP.
Furthermore, a Markov chain in function space can be constructed by a sequence of FMTOs  $p(f^{(1)}|f^{(0)}), \ldots, p(f^{(T)}|f^{(T - 1)})$  where  $p(f^{(t)}|f^{(t - 1)}) \coloneqq \{p_{x_{1:n}}(y_{1:n}^{(t)}|y_{1:n}^{(t - 1)})\}_{x_{1:n}}$ . With repeated applications of Proposition 4.2, if the initial state  $\{p_{x_{1:n}}(y_{1:n}^{(0)})\}_{x_{1:n}}$  is exchangeable and consistent,  $\{p_{x_{1:n}}(y_{1:n}^{(T)})\}_{x_{1:n}}$  at time  $T$  is also exchangeable and consistent, hence defining a SP  $p(f^{(T)})$ .
Equation (6) then provides a valid construction of consistent and exchangeable MMTOs by introducing an auxiliary latent variable  $z$ . The transition operator is written as
$$
p _ {x _ {1: n}} \left(y _ {1: n} \mid y _ {1: n} ^ {(0)}; \theta\right) = \int p _ {\theta} (z) \prod_ {i = 1} ^ {n} \delta \left(y _ {i} - F _ {\theta} \left(y _ {i} ^ {(0)}; x _ {i}, z\right)\right) d z, \tag {9}
$$
where both  $p_{\theta}$  and  $F_{\theta}$  in Equation (9) can be parameterised by neural networks. Critically, we can extend Proposition 4.2 to cover these auxiliary settings in Equation (9) as well, as per the following result (see Appendix A for proof):
Proposition 4.3. MMTOs in the form of Equation (9) are consistent and exchangeable.
4.3 Parameterisation, inference and training
We can now define a MNP as a sequence of neural FMTOs, with each FMTO indirectly specified via a collection of MMTOs in the form of Equation (9). If we specify a distribution over the sequence of auxiliary latent variables  $p_{\theta}(z^{(1:T)})$  along with an initial SP with marginals  $p_{x_{1:n}}(y_{1:n}^{(0)})$ , we can write down the marginal distribution over function outputs  $y_{1:n} \coloneqq y_{1:n}^{(T)}$  for a sequence of inputs  $x_{1:n}$  under the MNP model as (see Figure 2a for illustration and Appendix A for derivation):
$$
\begin{array}{l} p _ {x _ {1: n}} \left(y _ {1: n}; \theta\right) = \int p _ {\theta} \left(z ^ {(1: T)}\right) p _ {x _ {1: n}} \left(y _ {1: n} ^ {(0)}\right) \prod_ {t = 1} ^ {T} \prod_ {i = 1} ^ {n} p _ {\theta} ^ {(t)} \left(y _ {i} ^ {(t)} \mid y _ {i} ^ {(t - 1)}, x _ {i}, z ^ {(t)}\right) d y _ {1: n} ^ {(0: T - 1)} d z ^ {(1: T)} \tag {10} \\ = \int p _ {\theta} \left(z ^ {(1: T)}\right) p _ {x _ {1: n}} \left(y _ {1: n} ^ {(0)}\right) \prod_ {t = 1} ^ {T} \prod_ {i = 1} ^ {n} \left| \det  \frac {\partial F _ {\theta} ^ {(t)} \left(y _ {i} ^ {(t - 1)} ; x _ {i} , z ^ {(t)}\right)}{\partial y _ {i} ^ {(t - 1)}} \right| d z ^ {(1: T)} \\ \end{array}
$$
According to Propositions 4.2 and 4.3,  $\{p_{x_{1:n}}(y_{1:n};\theta)\}_{x_{1:n}}$  defines a valid SP parameterised by  $\theta$ . The initial SP can be arbitrarily chosen, as long as we can evaluate its marginals  $p_{x_{1:n}}(y_{1:n}^{(0)})$ . In our experiments, we use a trivial SP where all the outputs are i.i.d. standard normal distributed. We adopt normalising flows [16, 43, 49] to parameterise the invertible transformations  $F_{\theta}^{(t)}$ .
To integrate over latent variables  $z^{(1:T)}$  in Equation (10), we introduce a posterior approximation  $q_{x_{1:n}}(z^{(1:T)}|y_{1:n};\phi)$ . For many applications, we need to learn and query the conditional distributions of a target  $\mathcal{T} = \{(x_i,y_i)\}_{i=m+1}^n$  given context  $\mathcal{C} = \{(x_i,y_i)\}_{i=1}^m$ . To better reflect the desired model behaviour at test time, similar to NPs [22], we train the model by maximising the following approximation of the conditional log-likelihood w.r.t. both model  $\theta$  and variational  $\phi$  parameters:
$$
\begin{array}{l} \mathcal {L} _ {\theta , \phi} \left(y _ {1: n}; x _ {1: n}\right) := \\ \mathbb {E} _ {q _ {x 1: n} \left(z ^ {(1: T)} \mid y _ {1: n}; \phi\right)} \left[ \log p _ {x _ {m + 1: n}} \left(y _ {m + 1: n} \mid z ^ {(1: T)}, \theta\right) + \log \frac {q _ {x 1 : m} \left(z ^ {(1 : T)} \mid y _ {1 : m} ; \phi\right)}{q _ {x 1 : n} \left(z ^ {(1 : T)} \mid y _ {1 : n} ; \phi\right)} \right] \tag {11} \\ \end{array}
$$
where  $\log p_{x_{m+1:n}|x_{1:m}}(y_{m+1:n} | y_{1:m}; \theta) \gtrsim \mathcal{L}_{\theta, \phi}(y_{1:n}; x_{1:n})$ . Here  $q_{x_{1:m}}(z^{(1:T)} | y_{1:m}; \phi)$  can be seen as an approximate prior conditioned on the context  $\{(x_i, y_i)\}_{i=1}^m$ , while  $q_{x_{1:n}}(z^{(1:T)} | y_{1:n}; \phi)$  is the approximate posterior after the target  $\{(x_i, y_i)\}_{i=m+1}^n$  is observed. Thus the prior and the posterior share the same inference model with parameters  $\phi$ . The training objective is optimised by stochastic gradient descent, and the gradients  $\nabla_{\theta, \phi}\mathcal{L}_{\theta, \phi}(y_{1:n}; x_{1:n})$  can be efficiently estimated with low variance using the reparameterisation trick [40].
However, different from NPs, the latent variables  $z^{(1:T)}$  for MNPs are a sequence. Therefore, we propose a different inference model which shares parameters with the generative model above. Firstly, we write  $q_{x_{1:m}}(z^{(1:T)} \mid y_{1:m}; \phi)$  in a factorised form:
$$
q _ {x _ {1: n}} \left(z ^ {(1: T)} \mid y _ {1: n}; \phi\right) = \prod_ {t = 1} ^ {T} q _ {x _ {1: n}} ^ {(t)} \left(z ^ {(t)} \mid z ^ {(t + 1)}, y _ {1: n} ^ {(t)}; \phi\right) \tag {12}
$$
where  $\phi$  are variational parameters and we set  $z^{(T + 1)} = \mathbf{0}$ ,  $y_{1:n} = y_{1:n}^{(T)}$ . In our implementation, we use a Gaussian distribution for each factor  $q_{x_{1:n}}^{(t)}$ , with mean  $\mu_{\phi}^{(t)}(z^{(t + 1)},y_{1:n}^{(t)},x_{1:n})$  and variance  $\Sigma_{\phi}^{(t)}(z^{(t + 1)},y_{1:n}^{(t)},x_{1:n}))$ . Here  $\mu_{\phi}^{(t)}$  and  $\Sigma_{\phi}^{(t)}$  are parameterised functions invariant to the permutation of data points  $\{(x_i,y_i)\}_{i = 1}^n$ , i.e.
$$
\mu_ {\phi} ^ {(t)} (z ^ {(t + 1)}, y _ {1: n} ^ {(t)}, x _ {1: n}) = \mu_ {\phi} ^ {(t)} (z ^ {(t + 1)}, y _ {\pi (1: n)} ^ {(t)}, x _ {\pi (1: n)})
$$
$$
\Sigma_ {\phi} ^ {(t)} (z ^ {(t + 1)}, y _ {1: n} ^ {(t)}, x _ {1: n}) = \Sigma_ {\phi} ^ {(t)} (z ^ {(t + 1)}, y _ {\pi (1: n)} ^ {(t)}, x _ {\pi (1: n)}).
$$
We can parameterise them by first having
$$
r ^ {(t)} = \operatorname {S E T E N C O D E R} \left(y _ {1: n} ^ {(t)}, x _ {1: n}\right) \tag {13}
$$
where SETENCODER could be a Set Transformer [36], Deep Sets [61], then taking
$$
\mu_ {\phi} ^ {(t)} \left(z ^ {(t + 1)}, y _ {1: n} ^ {(t)}, x _ {1: n}\right) = \operatorname {M L P} _ {\mu} \left(z ^ {(t + 1)}, r ^ {(t)}\right) \tag {14}
$$
$$
\sigma_ {\phi} ^ {(t)} \left(z ^ {(t + 1)}, y _ {1: n} ^ {(t)}, x _ {1: n}\right) = \mathrm {M L P} _ {\sigma} \left(z ^ {(t + 1)}, r ^ {(t)}\right). \tag {15}
$$
(a) Generative models
(b) Inference models
Figure 2: (a) Consistent generative models on finite sets. Observed variables are shown in white, and hidden variables are shaded. Stochastic paths are indicated with dashed lines while deterministic paths are solid lines. Conditioned on the function inputs  $x_{1:n}$  and the auxiliary latent variables  $z^{(1:T)}$ , the function outputs are transformed independently using instance-wise conditional normalising flows (CNFs). If the initial states  $\{(x_i, y_i)\}_{i=1:n}$  come from a SP, the construction will ensure that the collection of marginals  $p(y_{1:n}^{(t)}|x_{1:n})$  are consistent under marginalisation for any  $t$ . (b) Permutation-invariant inference models. Auxiliary latent variables  $z^{(1:T)}$  are inferred in reverse order in our inference model. We reuse the parameters of CNFs in the generative model to compute function values  $y_{1:n}^{t}$  at all time steps, where each step has a different colour. We parameterise the conditional distribution  $q_{x_{1:n}}^{(t)}(z^{(t)}|z^{(t+1)}, y_{1:n}^{(t)}; \phi)$  with permutation-invariant neural networks.
In Equation (12), the intermediate function outputs  $y_{1:n}^{(1:T-1)}$  are not observed. However, we can sample them autoregressively by iterating between sampling  $z^{t+1}$  and calculating  $y_i^{(t)} = (F_\theta^{(t+1)})^{-1}(y_i^{(t+1)}; x_i, z^{(t+1)})$  for each  $i$  (see Figure 2b). Therefore, we share the parameters of the normalising flows  $F_\theta^{(1:T)}$  between the generative and inference models, leading to better performance. This idea was originally proposed by Cornish et al. [11], except that our normalising flows are applied to a sequence of function outputs  $y_{1:n}$  given the inputs  $x_{1:n}$ .