Question: Represent the scientific article to find a short TLDR

Question: We introduce the Convolutional Conditional Neural Process (ConvCNP), a new member of the Neural Process family that models translation equivariance in the data. Translation equivariance is an important inductive bias for many learning problems including time series modelling, spatial data, and images. The model embeds data sets into an infinite-dimensional function space, as opposed to finite-dimensional vector spaces. To formalize this notion, we extend the theory of neural representations of sets to include functional representations, and demonstrate that any translation-equivariant embedding can be represented using a convolutional deep-set. We evaluate ConvCNPs in several settings, demonstrating that they achieve state-of-the-art performance compared to existing NPs. We demonstrate that building in translation equivariance enables zero-shot generalization to challenging, out-of-domain tasks. Neural Processes (NPs; b; a) are a rich class of models that define a conditional distribution p(y|x, Z, θ) over output variables y given input variables x, parameters θ, and a set of observed data points in a context set Z = {x m, y m} M m=1. A key component of NPs is the embedding of context sets Z into a representation space through an encoder Z → E(Z), which is achieved using a DEEPSETS function approximator . This simple model specification allows NPs to be used for (i) meta-learning , since predictions can be generated on the fly from new context sets at test time; and (ii) multi-task or transfer learning , since they provide a natural way of sharing information between data sets. Moreover, conditional NPs (CNPs; a), a deterministic variant of NPs, can be trained in a particularly simple way with maximum likelihood learning of the parameters θ, which mimics how the system is used at test time, leading to strong performance . Natural application areas of NPs include time series, spatial data, and images with missing values. Consequently, such domains have been used extensively to benchmark current NPs (a; b;). Often, ideal solutions to prediction problems in such domains should be translation equivariant: if the data are translated in time or space, then the predictions should be translated correspondingly . This relates to the notion of stationarity. As such, NPs would ideally have translation equivariance built directly into the modelling assumptions as an inductive bias. Unfortunately, current NP models must learn this structure from the data set instead, which is sample and parameter inefficient as well as impacting the ability of the models to generalize. The goal of this paper is to build translation equivariance into NPs. Famously, convolutional neural networks (CNNs) added translation equivariance to standard multilayer perceptrons . However, it is not straightforward to generalize NPs in an analogous way: (i) CNNs require data to live "on the grid" (e.g. image pixels form a regularly spaced grid), while many of the above domains have data that live "off the grid" (e.g. time series data may be observed irregularly at any time t ∈ R). (ii) NPs operate on partially observed context sets whereas CNNs typically do not. (iii) NPs rely on embedding sets into a finite-dimensional vector space for which the notion of equivariance with respect to input translations is not natural, as we detail in Section 3. In this work, we introduce the CONVCNP, a new member of the NP family that accounts for translation equivariance. 1 This is achieved by extending the theory of learning on sets to include functional representations, which in turn can be used to express any translation-equivariant NP model. Our key contributions can be summarized as follows. (i) We provide a representation theorem for translation-equivariant functions on sets, extending a key of to functional embeddings, including sets of varying size. (ii) We extend the NP family of models to include translation equivariance. (iii) We evaluate the CONVCNP and demonstrate that it exhibits excellent performance on several synthetic and real-world benchmarks. In this section we introduce the notation and precisely define the problem this paper addresses. Notation. In the following, let X = R d and Y ⊆ R d (with Y compact) be the spaces of inputs and outputs respectively. To ease notation, we often assume scalar outputs Y ⊆ R. Define Z M = (X × Y) M as the collection of M input-output pairs, Z ≤M = M m=1 Z m as the collection of at most M pairs, and Z = ∞ m=1 Z M as the collection of finitely many pairs. Since we will consider permutation-invariant functions on Z (defined later in Property 1), we may refer to elements of Z as sets or data sets. Furthermore, we will use the notation [n] = {1, . . ., n}. CNPs model predictive distributions as p(y|x, Z) = p(y|Φ(x, Z), θ), where Φ is defined as a composition ρ • E of an encoder E: Z → R e mapping into the embedding space R e and a decoder ρ: R e → C b (X, Y). Here E(Z) ∈ R e is a vector representation of the set Z, and C b (X, Y) is the space of continuous, bounded functions X → Y endowed with the supremum norm. While NPs (b) employ latent variables to indirectly specify predictive distributions, in this work we focus on CNP models which do not. Then a mapping Φ: Z → H is called translation equivariant if Φ(T τ Z) = T τ Φ(Z) for all τ ∈ X and Z ∈ Z. Having formalized the problem, we now describe how to construct CNPs that translation equivariant. We are interested in translation equivariance (Property 2) with respect to translations on X. The NP family encoder maps sets Z to an embedding in a vector space R d, for which the notion of equivariance with respect to input translations in X is not well defined. For example, a function f on X can be translated by τ ∈ X: f (· − τ). However, for a vector x ∈ R d, which can be seen as a function [d] → R, x(i) = x i, the translation x(· − τ) is not well-defined. To overcome this issue, we enrich the encoder E: Z → H to map into a function space H containing functions on X. Since functions in H map from X, our notion of translation equivariance (Property 2) is now also well defined for E(Z). As we demonstrate below, every translation-equivariant function on sets has a representation in terms of a specific functional embedding. Definition 1 (Functional mappings on sets and functional representations of sets). Call a map E: Z → H a functional mapping on sets if it maps from sets Z to an appropriate space of functions H. Furthermore, call E(Z) the functional representation of the set Z. Considering functional representations of sets leads to the key of this work, which can be summarized as follows. For Z ⊆ Z appropriate, a continuous function Φ: Z → C b (X, Y) satisfies Properties 1 and 2 if and only if it has a representation of the form Φ(Z) = ρ (E(Z)), E(Z) = (x,y)∈Z φ(y)ψ(· − x) ∈ H, for some continuous and translation-equivariant ρ: H → C b (X, Y), and appropriate φ and ψ. Note that ρ is a map between function spaces. We also remark that continuity of Φ is not in the usual sense; we return to this below. Equation defines the encoder used by our proposed model, the CONVCNP. In Section 3.1, we present our theoretical in more detail. In particular, Theorem 1 establishes equivalence between any function satisfying Properties 1 and 2 and the representational form in Equation. In doing so, we provide an extension of the key of to functional representations on sets, and show that it can naturally be extended to handle varying-size sets. The practical implementation of CONVCNPs -the design of ρ, φ, and ψ -is informed by our in Section 3.1 (as well as the proofs, provided in Appendix A), and is discussed for domains of interest in Section 4. In this section we establish the theoretical foundation of the CONVCNP. We begin by stating a definition that is used in our main . We denote [m] = {1, . . ., m}. Definition 2 (Multiplicity). A collection Z ⊆ Z is said to have multiplicity K if, for every set Z ∈ Z, every x occurs at most K times: For example, in the case of real-world data like time series and images, we often observe only one (possibly multi-dimensional) observation per input location, which corresponds to multiplicity one. We are now ready to state our key theorem. 3, permutation invariant (Property 1), and translation Apply CNN and predict require: ρ = (CNN, ψρ), ψ, and density γ require: context (xn, yn) 10 end (b) require: ρ = CNN and E = CONV θ require: image I, context Mc, and target mask Mt 1 begin 2 // We discretize at the pixel locations. Figure 1: (a) Illustration of the CONVCNP forward pass in the off-the-grid case and pseudo-code for (b) off-the-grid and (c) on-the-grid data. The function pos: R → (0, ∞) is used to enforce positivity. for some continuous and translation-equivariant ρ: H → C b (X, Y) and some continuous φ: Y → R K+1 and ψ: X → R, where H is an appropriate space of functions that includes the image of E. We call a function Φ of the above form CONVDEEPSET. The proof of Theorem 1 is provided in Appendix A. We here discuss several key points from the proof that have practical implications and provide insights for the design of CONVCNPs: (i) For the construction of ρ and E, ψ is set to a flexible positive-definite kernel associated with a Reproducing Kernel Hilbert Space (RKHS;), which in desirable properties for E. (ii) Using the work by , we set φ(y) = (y 0, y 1, · · ·, y K) to be the powers of y up to order K. (iii) Theorem 1 requires ρ to be a powerful function approximator of continuous, translation-equivariant maps between functions. In Section 4, we discuss how these theoretical inform our implementations of CONVCNPs. Theorem 1 extends the of discussed in Section 2 by embedding the set into an infinite-dimensional space -the RKHS -instead of a finite-dimensional space. Beyond allowing the model to exhibit translation equivariance, the RKHS formalism allows us to naturally deal with finite sets of varying sizes, which turns out to be challenging with finite-dimensional embeddings. Furthermore, our formalism requires φ(y) = (y 0, y 1, y 2, . . ., y K) to expand up to order no more than the multiplicity of the sets K; if K is bounded, then our hold for sets up to any arbitrarily large finite size M, while fixing φ to be only (K + 1)-dimensional. In this section we discuss the architectures and implementation details for CONVCNPs. Similar to NPs, CONVCNPs model the conditional distribution as where Z is the observed data and Φ a CONVDEEPSET. The key considerations are the design of ρ, φ, and ψ for Φ. We provide separate models for data that lie on-the-grid and data that lie off-the-grid. Form of φ. The applications considered in this work have a single (potentially multi-dimensional) output per input location, so the multiplicity of Z is one (i.e., K = 1). It then suffices to let φ be a power series of order one, which is equivalent to appending a constant to y in all data sets, i.e. φ(y) = [1 y]. The first output φ 1 thus provides the model with information regarding where data has been observed, which is necessary to distinguish between no observed datapoint at x and a datapoint at x with y = 0. Denoting the functional representation as h, we can think of the first channel h as a "density channel". We found it helpful to divide the remaining channels h (1:) by h (Figure 1b, line 5), as this improved performance when there is large variation in the density of input locations. In the image processing literature, this is known as normalized convolution . The normalization operation can be reversed by ρ and is therefore not restrictive. CONVCNPs for off-the-grid data. Having specified φ, it remains to specify the form of ψ and ρ. Our proof of Theorem 1 suggests that ψ should be a stationary, non-negative, positive-definite kernel. The exponentiated-quadratic (EQ) kernel with a learnable length scale parameter is a natural choice. This kernel is multiplied by φ to form the functional representation E(Z) (Figure 1b, line 4; and Figure 1a, arrow 1). Next, Theorem 1 suggests that ρ should be a continuous, translation-equivariant map between function spaces. show that, in deep learning, any translation-equivariant model has a representation as a CNN. However, CNNs operate on discrete (on-the-grid) input spaces and produce discrete outputs. In order to approximate ρ with a CNN, we discretize the input of ρ, apply the CNN, and finally transform the CNN output back to a continuous function X → Y. To do this, for each context and test set, we space points (t i) n i=1 ⊆ X on a uniform grid (at a pre-specified density) over a hyper-cube that covers both the context and target inputs. We then evaluate (E(Z)(t i)) To map the output of the CNN back to a continuous function X → Y, we use the CNN outputs as weights for evenly-spaced basis functions (again employing the EQ kernel), which we denote by ψ ρ (Figure 1b, lines 7-8; Figure 1a, arrow 3). The ing approximation to ρ is not perfectly translation equivariant, but will be approximately so for length scales larger than the spacing of (E(Z)(t i)) n i=1. The ing continuous functions are then used to generate the (Gaussian) predictive mean and variance at any input. This, in turn, can be used to evaluate the log-likelihood. CONVCNP for on-the-grid data. While CONVCNP is readily applicable to many settings where data live on a grid, in this work we focus on the image setting. As such, the following description uses the image completion task as an example, which is often used to benchmark NPs (a;). Compared to the off-the-grid case, the implementation becomes simpler as we can choose the discretization (t i) n i=1 to be the pixel locations. Let I ∈ R H×W ×C be an image -H, W, C denote the height, width, and number of channels, respectively -and let M c be the context mask, which is such that [M c] i,j = 1 if pixel location (i, j) is in the context set, and 0 otherwise. To implement φ, we select all context points, Z c:= M c I, and prepend the context mask: (Figure 1c, line 4). Next, we apply a convolution to the context mask to form the density channel: Figure 1c, line 4). To all other channels, we apply a normalized convolution: (Figure 1c, line 5), where the division is element-wise. The filter of the convolution is analogous to ψ, which means that h is the functional representation, with the convolution performing the role of E (the summation in Figure 1b, line 4). Although the theory suggests using a non-negative, positive-definite kernel, we did not find significant empirical differences between an EQ kernel and using a fully trainable kernel restricted to positive values to enforce non-negativity (see Appendices D.4 and D.5 for details). Lastly, we describe the on-the-grid version of ρ(·), which consists of two stages. First, we apply a CNN to E(Z) (Figure 1c, line 6). Second, we apply a shared, pointwise MLP that maps the output of the CNN at each pixel location in the target set to R 2C, where we absorb MLP into the CNN (MLP can be viewed as an 1×1 convolution). The first C outputs are the means of a Gaussian predictive distribution and the second C the standard deviations, which then pass through a positivity-enforcing function (Figure 1c, . To summarise, the on-the-grid algorithm is given by multiplies by ψ and sums where (µ, σ) are the image mean and standard deviation, ρ is implemented with CNN, and E is implemented with the mask M c and convolution CONV. Training. Denoting the data set D = {Z n} N n=1 ⊆ Z and the parameters by θ, maximum-likelihood training involves (a; b) where we have split Z n into context (Z n,c) and target (Z n,t) sets. This is standard practice in the NP (a; b) and meta-learning settings and relates to neural auto-regressive models . Note that the context set and target set are disjoint (Z n,c ∩ Z n,t = ∅), which differs from the protocol for the NP (a). Practically, stochastic gradient descent methods can be used for optimization. We evaluate the performance of CONVCNPs in both on-the-grid and off-the-grid settings focusing on two central questions: (i) Do translation-equivariant models improve performance in appropriate domains? (ii) Can translation equivariance enable CONVCNPs to generalize to settings outside of those encountered during training? We use several off-the-grid data-sets which are irregularly sampled time series (X = R), comparing to Gaussian processes (GPs;) and ATTNCNP(which is identical to the ANP , but without the latent path in the encoder), the best performing member of the CNP family. We then evaluate on several on-the-grid image data sets (X = Z 2). In all settings we demonstrate substantial improvements over existing neural process models. For the CNN component of our model, we propose a small and large architecture for each experiment (in the experimental sections named CONVCNP and CONVCNPXL, respectively). We note that these architectures are different for off-the-grid and on-the-grid experiments, with full details regarding the architectures given in the appendices. First we consider synthetic regression problems. At each iteration, a function is sampled, followed by context and target sets. Beyond EQ-kernel GPs (as proposed in Garnelo et al. (2018a); ), we consider more complex data arising from Matern-5 2 and weakly-periodic kernels, as well as a challenging, non-Gaussian sawtooth process with random shift and frequency (see Figure 2, for example). CONVCNP is compared to CNP (a) and ATTNCNP. Training and testing procedures are fixed across all models. Full details on models, data generation, and training procedures are provided in Appendix C. Params EQ Weak Periodic Matern Sawtooth CNP 66818 -1.49 ± 3e-3 -2.00 ± 2e-3 -1.61 ± 1e-3 -0.51 ± 1e-5 ATTNCNP 149250 1.08 ± 4e-3 -2.01 ± 2e-3 0.42 ± 2e-3 -0.50 ± 2e-3 CONVCNP 6537 1.03 ± 5e-3 -1.52 ± 2e-3 0.51 ± 4e-3 4.38 ± 4e-3 CONVCNPXL 50617 1.90 ± 4e-3 -1.19 ± 2e-3 0.81 ± 4e-3 6.01 ± 1e-3 Table 1 reports the log-likelihood means and standard errors of the models over 1000 tasks. The context and target points for both training and testing lie within the interval [−2, 2] where training data was observed (marked "training data range" in Figure 2). Table 1 demonstrates that, even when extrapolation is not required, CONVCNP significantly outperforms other models in all cases, despite having fewer parameters. of CONVCNP and ATTNCNP when data is observed outside the range where the models were trained: translation equivariance enables CONVCNP to elegantly generalize to this setting, whereas ATTNCNP is unable to generate reasonable predictions. The PLAsTiCC data set is a simulation of transients observed by the LSST telescope under realistic observational conditions. The data set contains 3,500,734 "light curves", where each measurement is of an object's brightness as a function of time, taken by measuring the photon flux in six different astronomical filters. The data can be treated as a six-dimensional time series. The data set was introduced in a Kaggle competition, 4 where the task was to use these light curves to classify the variable sources. The winning entry modeled the light curves with GPs and used these models to generate features for a gradient boosted decision tree classifier. We compare a multi-input-multi-output CONVCNP with the GP models used in Avocado. 5 CONVCNP accepts six channels as inputs, one for each astronomical filter, and returns 12 outputs -the means and standard deviations of six Gaussians. Full experimental details are given in Appendix C.3. The mean squared error of both approaches is similar, but the held-out log-likelihood from the CONVCNP is far higher (see Table 2). Table 2: Mean and standard errors of log-likelihood and root mean squared error over 1000 test objects from the PLastiCC dataset. Log-likelihood MSE Kaggle GP -0.335 ± 0.09 0.037 ± 4e-3 ConvCP (ours) 1.31 ± 0.30 0.040 ± 5e-3 The CONVCNP model is well suited for applications where simulation data is plentiful, but real world training data is scarce (Sim2Real). The CONVCNP can be trained on a large amount of simulation data and then be deployed with real-world training data as the context set. We consider the Lotka-Volterra model , which is used to describe the evolution of predator-prey populations. This model has been used in the Approximate Bayesian Computation literature where the task is to infer the parameters from samples drawn from the Lotka-Volterra process . These methods do not simply extend to prediction problems such as interpolation or forecasting. In contrast, we train CONVCNP on synthetic data sampled from the Lotka-Volterra model and can then condition on real-world data from the Hudson's Bay lynx-hare data set to perform interpolation (see Figure 3 ; full experimental details are given in Appendix C.4). The CONVCNP performs accurate interpolation as shown in Figure 3. We were unable to successfully train the ATTNCNP for this task. We suspect this is because the simulation data are variable lengthtime series, which requires models to leverage translation equivariance at training time. As shown in Section 5.1, the ATTNCNP struggles to do this (see Appendix C.4 for complete details). To test CONVCNP beyond one-dimensional features, we evaluate our model on on-the-grid image completion tasks and compare it to ATTNCNP. Image completion can be cast as a prediction of pixel intensities y * i (∈ R 3 for RGB, ∈ R for greyscale) given a target 2D pixel location x * i conditioned on an observed (context) set of pixel values Z = (x n, y n) N n=1. In the following experiments, the context set can vary but the target set contains all pixels from the image. Further experimental details are in Appendix D.1. Standard benchmarks. We first evaluate the model on four common benchmarks: MNIST , SVHN , and 32 × 32 and 64 × 64 CelebA . Importantly, these data sets are biased towards images containing a single, well-centered object. As a , perfect translation-equivariance might hinder the performance of the model when the test data are similarly structured. We therefore also evaluated a larger CONVCNP that can learn such non-stationarity, while still sharing parameters across the input space (CONVCNPXL). Table 3 shows that CONVCNP significantly outperforms ATTNCNP when it has a large receptive field size, while being at least as good with a small receptive field size. Qualitative samples for various context sets can be seen in Figure 5. Further qualitative comparisons and ablation studies can be found in Appendix D.3 and Appendix D.4 respectively. Generalization to multiple, non-centered objects. The data sets from the previous paragraphs were centered and contained single objects. Here we test whether CONVCNPs trained on such data can generalize to images containing multiple, non-centered objects. The last column of Table 3 evaluates the models in a zero shot multi-MNIST (ZSMM) setting, where images contain multiple digits at test time (Appendix D.2). CONVCNP significantly outperforms ATTNCNP on such tasks. Figure 4a shows a histogram of the image log-likelihoods for CONVCNP and ATTNCNP, as well as qualitative at different percentiles of the CONVCNP distribution. CONVCNP is able to extrapolate to this out-of-distribution test set, while ATTNCNP appears to model the bias of the training data and predict a centered "mean" digit independently of the context. Interestingly, CONVCNPXL does not perform as well on this task. In particular, we find that, as the receptive field becomes very large, performance on this task decreases. We hypothesize that this has to do with behavior of the model at the edges of the image. CNNs with larger receptive fields -the region of input pixels that affect a particular output pixel -are able to model non-stationary behavior by looking at the distance from any pixel to the image boundary. We expand on this discussion and provide further experimental evidence regarding the effects of receptive field on the ZSMM task in Appendix D.6. Although ZSMM is a contrived task, note that our field of view usually contains multiple independent objects, thereby requiring translation equivariance. As a more realistic example, we took a CONVCNP model trained on CelebA and tested it on a natural image of different shape which contains multiple people (Figure 4b). Even with 95% of the pixels removed, the CONVCNP was able to produce a qualitatively reasonable reconstruction. A comparison with ATTNCNP is given in Appendix D.3. Computational efficiency. Beyond the performance and generalization improvements, a key advantage of the CONVCNP is its computational efficiency. The memory and time complexity of a single self-attention layer grows quadratically with the number of inputs M (the number of pixels for images) but only linearly for a convolutional layer. Empirically, with a batch size of 16 on 32 × 32 MNIST, CONVCNPXL requires 945MB of VRAM, while ATTNCNP requires 5839 MB. For the 56×56 ZSMM CONVCNPXL increases its requirements to 1443 MB, while ATTNCNP could not fit onto a 32GB GPU. Ultimately, ATTNCNP had to be trained with a batch size of 6 (using 19139 MB) and we were not able to fit it for CelebA64. Recently, restricted attention has been proposed to overcome this computational issue , but we leave an investigation of this and its relationship to CONVCNPs to future work. We have introduced CONVCNP, a new member of the CNP family that leverages embedding sets into function space to achieve translation equivariance. The relationship to (i) the NP family, and (ii) representing functions on sets, each imply extensions and avenues for future work. Deep sets. Two key issues in the existing theory on learning with sets (; a;) are (i) the restriction to fixed-size sets, and (ii) that the dimensionality of the embedding space must be no less than the cardinality of the embedded sets. Our work implies that by considering appropriate embeddings into a function space, both issues are alleviated. In future work, we aim to further this analysis and formalize it in a more general context. Point-cloud models. Another line of related research focuses on 3D point-cloud modelling (a; b). While original work focused on permutation invariance (a;), more recent work has considered translation equivariance as well , leading to a model closely resembling CONVDEEPSETS. The key differences with our work are the following: (i) Correlated samples and consistency under marginalization. In the predictive distribution of CON-VCNP (Equation), predicted ys are conditionally independent given the context set. Consequently, samples from the predictive distribution lack correlations and appear noisy. One solution is to instead define the predictive distribution in an autoregressive way, like e.g. PixelCNN++ . Although samples are now correlated, the quality of the samples depends on the order in which the points are sampled. Moreover, the predicted ys are then not consistent under marginalization (b;). Consistency under marginalization is more generally an issue for neural autoregressive models , although consistent variants have been devised . To overcome the consistency issue for CONVCNP, exchangeable neural process models (e.g. ;) may provide an interesting avenue. Another way to introduce dependencies between ys is to employ latent variables as is done in neural processes (b). However, such an approach only achieves conditional consistency: given a context set, the predicted ys will be dependent and consistent under marginalization, but this does not lead to a consistent joint model that also includes the context set itself. For each dataset, an image is randomly sampled, the first row shows the given context points while the second is the mean of the estimated conditional distribution. From left to right the first seven columns correspond to a context set with 3, 1%, 5%, 10%, 20%, 30%, 50%, 100% randomly sampled context points. In the last two columns, the context sets respectively contain all the pixels in the left and top half of the image. CONVCNPXL is shown for all datasets besides ZSMM, for which we show the fully translation equivariant CONVCNP. In this section, we provide the proof of Theorem 1. Our proof strategy is as follows. We first define an appropriate topology for fixed-sized sets (Appendix A.1). With this topology in place, we demonstrate that our proposed embedding into function space is homeomorphic (Lemmas 1 and 2). We then show that the embeddings of fixed-sized sets can be extended to varying-sized sets by "pasting" the embeddings together while maintaining their homeomorphic properties (Lemma 3). Following this, we demonstrate that the ing embedding may be composed with a continuous mapping to our desired target space, ing in a continuous mapping between two metric spaces (Lemma 4). Finally, in Appendix A.3 we combine the above-mentioned to prove Theorem 1. We begin with definitions that we will use throughout the section and then present our . Let X = R d and let Y ⊆ R be compact. Let ψ be a symmetric, positive-definite kernel on X. By the Moore-Aronszajn Theorem, there is a unique Hilbert space (H, ·, · H) of real-valued functions on X for which ψ is a reproducing kernel. This means that (i) ψ(·, x) ∈ H for all x ∈ X and (ii) f, ψ(·, x) H = f (x) for all f ∈ H and x ∈ X (reproducing property). For ψ: X × X → R, X = (x 1, . . ., x n) ∈ X n, and X = (x 1, . . ., x n) ∈ X n, we denote Definition 3 (Interpolating RKHS). Call H interpolating if it interpolates any finite number of points: for every ((x i, y i)) For example, the RKHS induced by any strictly positive-definite kernel, e.g. the exponentiated quadratic (EQ) kernel ψ(x, A.1 THE QUOTIENT SPACE A n / S n Let A be a Banach space. For x = (x 1, . . ., x n) ∈ A n and y = (y 1, . . ., y n) ∈ A n, let x ∼ y if x is a permutation of y; that is, x ∼ y if and only if x = πy for some π ∈ S n where πy = (y π,..., y π(n) ). Let A n / S n be the collection of equivalence classes of ∼. Denote the equivalence class of Proof. We first show that d is well defined on A n / S n. Assume x ∼ x and y ∼ y. Then, x = π x x and y = π y y. Using the group properties of S n and the permutation invariance of · A n: To show the triangle inequality, note that using permutation invariance of · A n. Hence, taking the minimum over π 1, so taking the minimum over π 2 gives the triangle inequality for d. Proposition 2. The canonical map A n → A n / S n is continuous under the metric topology induced by d. Proposition 3. Let A ⊆ A n be topologically closed and closed under permutations. Then [A] is topologically closed in A n / S n under the metric topology. Then there are permutations (π n) ∞ n=1 ⊆ S n such that π n a n → x. Here π n a n ∈ A, because A is closed under permutations. Thus x ∈ A, as A is also topologically closed. We conclude that Proposition 5. The quotient topology on A n / S n induced by the canonical map is metrizable with the metric d. Proof. Since the canonical map is surjective, there exists exactly one topology on A n / S n relative to which the canonical map is a quotient map: the quotient topology . Let p: A n → A n / S n denote the canonical map. It remains to show that p is a quotient map under the metric topology induced by d; that is, we show that U ⊂ A n / S n is open in A n / S n under the metric topology if and only if p Whereas A previously denoted an arbitrary Banach space, in this section we specialize to A = X × Y. We denote an element in A by (x, y) and an element in Z M = A M by ((x 1, y 1),..., (x M, y M)). Alternatively, we denote ((x 1, y 1),..., (x M, y M)) by (X, y) where X = (x 1, . . ., x M) ∈ X M and y = (y 1, . . ., y M) ∈ Y M. We clarify that an element in Z M = A M is permuted as follows: for π ∈ S M, π(X, y) = π ((x 1, y 1),..., (x M, y M)) = ((x π, y π ),..., (x π(n), y π(n) )) = (πX, πy). Note that permutation-invariant functions on Z M are in correspondence to functions on the quotient space induced by the equivalence class of permutations, Z M / S m The latter is a more natural representation. Lemma 3 states that it is possible to homeomorphically embed sets into an RKHS. This is key to proving our main . Before proving Lemma 3, we provide several useful . We begin by demonstrating that an embedding of sets of a fixed size into a RKHS is continuous and injective. Lemma 1. Consider a collection Z M ⊆ Z M that has multiplicity K. Set and let ψ be an interpolating, continuous positive-definite kernel. Define where H K+1 = H × · · · × H is the (K + 1)-dimensional-vector-valued-function Hilbert space constructed from the RKHS H for which ψ is a reproducing kernel and endowed with the inner product f, g is injective, hence invertible, and continuous. Proof. First, we show that E M is injective. Suppose that Denote X = (x 1, . . ., x M) and y = (y 1, . . ., y M), and denote X and y similarly. Taking the inner product with any f ∈ H on both sides and using the reproducing property of ψ, this implies that for all f ∈ H. In particular, since by construction φ 1 (·) = 1, for all f ∈ H. Using that H is interpolating, choose a particularx ∈ X ∪ X, and let f ∈ H be such that f (x) = 1 and f (·) = 0 at all other x i and x i. Then so the number of suchx in X and the number of suchx in X are the same. Since this holds for everyx, X is a permutation of X: X = π(X) for some permutation π ∈ S M. Plugging in the permutation, we can write Then, by a similar argument, for any particularx, Let the number of terms in each sum equal S. Since Z M has multiplicity K, S ≤ K. By Lemma 4 from , the'sum-of-power mapping' from {y i : x i =x} to the first S + 1 elements of i:xi=x φ(y i), i.e. i:xi=x y 0 i,..., i:xi=x y S i, is injective. Therefore, (y i) i:xi=x is a permutation of (y π(i) ) i:xi=x. Note that x i =x for all above y i. Furthermore, note that also x π(i) = x i =x for all above y π(i). We may therefore adjust the permutation π such that y i = y π(i) for all i such that x i =x whilst retaining that x = π(x). Performing this adjustment for allx, we find that y = π(y) and x = π(x). Second, we show that E M is continuous. Compute Having established the injection, we now show that this mapping is a homeomorphism, i.e. that the inverse is continuous. This is formalized in the following lemma. Lemma 2. Consider Lemma 1. Suppose that Z M is also topologically closed in A M and closed under permutations, and that ψ also satisfies (i) M is continuous. Remark 1. To define Z 2 with multiplicity one, one might be tempted to define which indeed has multiplicity one. Unfortunately, Z 2 is not closed: if ⊆ X and ⊆ Y, then (, (1/n, 2)) ∞ n=1 ⊆ Z 2, but (, (1/n, 2)) → (,) / ∈ Z 2, because 0 then has two observations 1 and 2. To get around this issue, one can require an arbitrarily small, but non-zero spacing > 0 between input locations: This construction can be generalized to higher numbers of observations and multiplicities as follows: Remark 2. Before moving on to the proof of Lemma 2, we remark that Lemma 2 would directly follow if Z M were bounded: then Z M is compact, so E M is a continuous, invertible map between a compact space and a Hausdorff space, which means that E −1 M must be continuous. The intuition that the must hold for unbounded Z M is as follows. Since φ 1 (·) = 1, for every f ∈ H M, f 1 is a summation of M "bumps" (imagine the EQ kernel) of the form ψ(·, x i) placed throughout X. If one of these bumps goes off to infinity, then the function cannot uniformly converge pointwise, which means that the function cannot converge in H (if ψ is sufficiently nice). Therefore, if the function does converge in H, (x i) M i=1 must be bounded, which brings us to the compact case. What makes this work is the density channel φ 1 (·) = 1, which forces (x i) M i=1 to be well behaved. The above argument is formalized in the proof of Lemma 2. First, we demonstrate that, assuming the claim, H M is closed. Note that by boundedness of is compact and therefore closed, since every compact subset of a metric space is closed. Therefore, the image of E M | [Z J] contains the limit f. Since the image of E M | [Z J] is included in H M, we have that f ∈ H M, which shows that H M is closed. Next, we prove that, assuming the claim, E −1 −1 is continuous, because a continuous bijection from a compact space to a metric space is a homeomorphism. Therefore M is continuous. It remains to show the claim. Let f 1 denote the first element of f, i.e. the density channel. Using the reproducing property of ψ, → f 1 in H means that it does so uniformly pointwise (over x). Hence, we can let N ∈ N be such that n ≥ N implies that |f At the same time, by pointwise non-negativity of ψ, we have that Towards contradiction, suppose that (which is a contradiction. The following lemma states that we may construct an encoding for sets containing no more than M elements into a function space, where the encoding is injective and every restriction to a fixed set size is a homeomorphism. and let ψ be an interpolating, continuous positive-definite kernel that satisfies where H K+1 = H × · · · × H is the (K + 1)-dimensional-vector-valued-function Hilbert space constructed from the RKHS H for which ψ is a reproducing kernel and endowed with the inner product f, g is injective, hence invertible. Denote this inverse by E −1, where Proof. Recall that E m is injective for every m ∈ [M]. Hence, to demonstrate that E is injective it remains to show that (H m) M m=1 are pairwise disjoint. To this end, suppose that for m = m. Then, by arguments like in the proof of Lemma 1, Since φ 1 (·) = 1, this gives m = m, which is a contradiction. Finally, by repeated application of Lemma 2, E −1, the space of continuous bounded functions from X to Y, such that every restriction Φ| [Z m] is continuous, and let E be from Lemma 3. Then is continuous. is continuous. From here on, we let ψ be a stationary kernel, which means that it only depends on the difference of its arguments and can be seen as a function X → R. With the above in place, we are finally ready to prove our central , Theorem 1. is closed under permutations, and (iv) is closed under translations. Set and let ψ be an interpolating, continuous positive-definite kernel that satisfies where H K+1 = H × · · · × H is the (K + 1)-dimensional-vector-valued-function Hilbert space constructed from the RKHS H for which ψ is a reproducing kernel and endowed with the inner product f, g Then a function Φ: Z ≤M → C b (X, Y) satisfies (i) continuity of the restriction Φ| Zm for every m ∈ [M], (ii) permutation invariance (Property 1), and (iii) translation equivariance (Property 2) if and only if it has a representation of the form is continuous and translation equivariant. Proof of sufficiency. To begin with, note that permutation invariance (Property 1) and translation equivariance (Property 2) for Φ are well defined, because Z ≤M is closed under permutations and translations by assumption. First, Φ is permutation invariant, because addition is commutative and associative. Second, that Φ is translation equivariant (Property 2) follows from a direct verification and that ρ is also translation equivariant: Proof of necessity. Our proof follows the strategy used by;. To begin with, since Φ is permutation invariant (Property 1), we may define for which we verify that every restriction Φ| [Z m] is continuous. By invertibility of E from Lemma 3, we have is translation equivariant, because ψ is stationary. Also, by assumption Φ is translation equivariant (Property 2). Thus, their composition ρ is also translation equivariant. Remark 3. The function ρ: H ≤M → C b (X, Y) may be continuously extended to the entirety of H K+1 using a generalisation of the Tietze Extension Theorem by. There are variants of Dugundji's Theorem that also preserve translation equivariance. In both our 1d and image experiments, our main comparison is to conditional neural process models. In particular, we compare to a vanilla CNP (1d only; Garnelo et al. (2018a) ) and an ATTNCNP . Our architectures largely follow the details given in the relevant publications. CNP baseline. Our baseline CNP follows the implementation provided by the authors. 6 The encoder is a 3-layer MLP with 128 hidden units in each layer, and RELU non-linearities. The encoder embeds every context point into a representation, and the representations are then averaged across each context set. Target inputs are then concatenated with the latent representations, and passed to the decoder. The decoder follows the same architecture, outputting mean and standard deviation channels for each input. Attentive CNP baseline. The ATTNCNP we use corresponds to the deterministic path of the model described by for image experiments. Namely, an encoder first embeds each context point c to a latent representation (x (c), xy ∈ R 128. For the image experiments, this is achieved using a 2-hidden layer MLP of hidden dimensions 128. For the 1d experiments, we use the same encoder as the CNP above. Every context point then goes through two stacked self-attention layers. Each self-attention layer is implemented with an 8-headed attention, a skip connection, and two layer normalizations (as described in , modulo the dropout layer). To predict values at each target point t, we embed x (t) → r (t) x and x (c) → r (c) x using the same single hidden layer MLP of dimensions 128. A target representation r (t) xy is then estimated by applying cross-attention (using an 8-headed attention described above) with keys K:= {r, and query q := r (t) x. Given the estimated target representationr xy, the conditional predictive posterior is given by a Gaussian pdf with diagonal covariance parametrised by pre ∈ R 3 and decoder is a 4 hidden layer MLP with 64 hidden units per layer for the images, and the same decoder as the CNP for the 1d experiments. , we enforce we set a minimum standard deviation σ (t) min = [0.1; 0.1; 0.1] to avoid infinite log-likelihoods by using the following post-processed standard deviation: In this section, we give details regarding our experiments for the 1d data. We begin by detailing model architectures, and then provide details for the data generating processes and training procedures. The density at which we evaluate the grid differs from experiment to experiment, and so the values are given in the relevant subsections. In all experiments, the weights are optimized using Adam and weight decay of 10 −5 is applied to all model parameters. The learning rates are specified in the following subsections. Throughout the experiments (Sections 5.1 to 5.3), we consider two models: CONVCNP (which utilizes a smaller architecture), and CONVCNPXL (with a larger architecture). For all architectures, the input kernel ψ was an EQ (exponentiated quadratic) kernel with a learnable length scale parameter, as detailed in Section 4, as was the kernel for the final output layer ψ ρ. When dividing by the density channel, we add ε = 10 −8 to avoid numerical issues. The length scales for the EQ kernels are initialized to twice the spacing 1/γ 1/d between the discretization points (t i), where γ is the density of these points and d is the dimensionality of the input space X. Moreover, we emphasize that the size of the receptive field is a product of the width of the CNN filters and the spacing between the discretization points. Consequently, for a fixed width kernel of the CNN, as the number of discretization points increases, the receptive field size decreases. One potential improvement that was not employed in our experiments, is the use of depthwise-separable convolutions . These dramatically reduce the number of parameters in a convolutional layer, and can be used to increase the CNN filter widths, thus allowing one to increase the number of discretization points without reducing the receptive field. The architectures for CONVCNP and CONVCNPXL are described below. CONVCNP. For the 1d experiments, we use a simple, 4-layer convolutional architecture, with RELU nonlinearities. The kernel size of the convolutional layers was chosen to be 5, and all employed a stride of length 1 and zero padding of 2 units. The number of channels per layer was set to, where the final channels where then processed by the final, EQ-based layer of ρ as mean and standard deviation channels. We employ a SOFTPLUS nonlinearity on the standard deviation channel to enforce positivity. This model has 6,537 parameters. CONVCNPXL. Our large architecture takes inspiration from UNet . We employ a 12-layer architecture with skip connections. The number of channels is doubled every layer for the first 6 layers, and halved every layer for the final 6 layers. We use concatenation for the skip connections. The following describes which layers are concatenated, where L i ← [L j, L k] means that the input to layer i is the concatenation of the activations of layers j and k: Like for the smaller architecture, we use RELU nonlinearities, kernels of size 5, stride 1, and zero padding for two units on all layers. The kernels used for the Gaussian Processes which generate the data in this experiment are defined as follows: • EQ: CNP when trained on an EQ kernel (with length scale parameter 1). "True function" refers to the sample from the GP prior from which the context and target sets were sub-sampled. " Ground Truth GP" refers to the GP posterior distribution when using the exact kernel and performing posterior inference based on the context set. The left column shows the predictive posterior of the models when data is presented in same range as training. The centre column shows the model predicting outside the training data range when no data is observed there. The right-most column shows the model predictive posteriors when presented with data outside the training data range. • weakly periodic: with f 1 (x) = cos(8πx) and f 2 (x) = sin(8πx), and with During the training procedure, the number of context points and target points for a training batch are each selected randomly from a uniform distribution over the integers between 3 and 50. This number of context and target points are randomly sampled from a function sampled from the process (a Gaussian process with one of the above kernels or the sawtooth process), where input locations are uniformly sampled from the interval [−2, 2]. All models in this experiment were trained for 200 epochs using 256 batches per epoch of batch size 16. We discretize E(Z) by evaluating 64 points per unit in this setting. We use a learning rate of 3e−4 for all models, except for CONVCNPXL on the sawtooth data, where we use a learning rate of 1e−3 (this learning rate was too large for the other models). The random sawtooth samples are generated from the following function: where A is the amplitude, f is the frequency, and t is "time". Throughout training, we fix the amplitude to be one. We truncate the series at an integer K. At every iteration, we sample a frequency uniformly in, K in, and a random shift in [−5, 5]. As the task is much harder, we sample context and target set sizes over. Here the CNP and ATTNCNP employ learning rates of 10 −3. All other hyperparameters remain unchanged. The CONVCNP was trained for 200 epochs using 1024 batches of batch size 4 per epoch. For training and testing, the number of context points for a batch are each selected randomly from a uniform distribution over the integers between 1 and the number of points available in the series (usually between 10-30 per bandwidth). The remaining points in the series are used as the target set. For testing, a batch size of 1 was used and statistics were computed over 1000 evaluations. We compare CONVCNP to the GP models used in using the implementation in https:// github.com/kboone/avocado. The data used for training and testing is normalized according to t(v) = (v − m)/s with the values in Table 4. These values are estimated from a batch sampled from the training data. To remove outliers in the GP , log-likelihood values less than −10 are removed from the evaluation. These same datapoints were removed from the CONVCNP as well. For this dataset, we only used the CONVCNPXL, as we found the CONVCNP to underfit. The learning rate was set to 10 −3, and we discretize E(Z) by evaluating 256 points per unit. We describe the way simulated training data for the experiment in Section 5.3 was generated from the Lotka-Volterra model. The description is borrowed from . Let X be the number of predators and Y the number of prey at any point in our simulation. According to the model, one of the following four events can occur: A: A single predator is born according to rate θ 1 XY, increasing X by one. B: A single predator dies according to rate θ 2 X, decreasing X by one. C: A single prey is born according to rate θ 3 Y, increasing Y by one. A single prey dies (is eaten) according to rate θ 4 XY, decreasing Y by one. The parameter values θ 1, θ 2, θ 3, and θ 4, as well as the initial values of X and Y govern the behavior of the simulation. We choose θ 1 = 0.01, θ 2 = 0.5, θ 3 = 1, and θ 4 = 0.01, which are also used in and generate reasonable time series. Note that these are likely not the parameter values that would be estimated from the Hudson's Bay lynx-hare data set , but they are used because they yield reasonably oscillating time series. Obtaining oscillating time series from the simulation is sensitive to the choice of parameters and many parametrizations in populations that simply die out. Time series are simulated using Gillespie's algorithm : 1. Draw the time to the next event from an exponential distribution with rate equal to the total rate θ 1 XY + θ 2 X + θ 3 Y + θ 4 XY. 2. Select one of the above events A, B, C, or D at random with probability proportional to its rate. 3. Adjust the appropriate population according to the selected event, and go to 1. The simulations using these parameter settings can yield a maximum population of approximately 300 while the context set in the lynx-hare data set has an approximate maximum population of about 80 so we scaled our simulation population by a factor of 2/7. We also remove time series which are longer than 100 units of time, which have more than 10000 events, or where one of the populations is entirely zero. The number of context points n for a training batch are each selected randomly from a uniform distribution between 3 and 80, and the number of target points is 150 − n. These target and context points are then sampled from the simulated series. The Hudson's Bay lynx-hare data set has time values that range from 1845 to 1935. However, the values supplied to the model range from 0 to 90 to remain consistent with the simulated data. For evaluation, an interval of 18 points is removed from the the Hudson's Bay lynx-hare data set to act as a target set, while the remaining 72 points act as the context set. This construction highlights the model's interpolation as well as its uncertainty in the presence of missing data. Models in this setting were trained for 200 epochs with 256 batches per epoch, each batch containing 50 tasks. For this data set, we only used the CONVCNP, as we found the CONVCNPXL to overfit. The learning rate was set to 10 −3, and we discretize E(Z) by evaluating 100 points per unit. We attempted to train an ATTNCNP for comparison, but due to the nature of the synthetic data generation, many of the training series end before 90 time units, the length of the Hudson's Bay lynx-hare series. Effectively, this means that the ATTNCNP was asked to predict outside of its training interval, a task that it struggles with, as shown in Section 5.1. The plots in Figure 9 show that the ATTNCNP is able to learn the first part of the time series but is unable to model data outside of the first 20 or so time units. Perhaps with more capacity and training epochs the ATTNCNP training would be more successful. Note from Figure 3 that our model does better on the synthetic data than on the real data. This could be due to the parameters of the Lotka-Volterra model used being a poor estimate for the real data. Training details. In all experiments, we sample the number of context points uniformly from U(ntotal 100, ntotal 2), and the number of target points is set to n total. The context and target points are sampled randomly from each of the 16 images per batch. The weights are optimised using Adam with learning rate 5 × 10 −4. We use a maximum of 100 epochs, with early stopping of 15 epochs patience. All pixel values are divided by 255 to rescale them to the range. In the following discussion, we assume that images are RGB, but very similar models can be used for greyscale images or other gridded inputs (e.g. 1d time series sampled at uniform intervals). Proposed convolutional CNP. Unlike ATTNCNP and off-the-grid CONVCNP, on-the-grid CON-VCNP takes advantage of the gridded structure. Namely, the target and context points can be specified in terms of the image, a context mask M c, and a target mask M t instead of sets of input-value pairs. Although this is an equivalent formulation, it makes it more natural and simpler to implement in standard deep learning libraries. In the following, we dissect the architecture and algorithmic steps succinctly summarized in Section 4. Note that all the convolutional layers are actually depthwise separable ; this enables a large kernel size (i.e. receptive fields) while being parameter and computationally efficient. 1. Let I denote the image. Select all context points signal:= M c I and append a density channel density:= M c, which intuitively says that "there is a point at this position": [signal, density]. Each pixel value will now have 4 channels: 3 RGB channels and 1 density channel M c. Note that the mask will set the pixel value to 0 at a location where the density channel is 0, indicating there are no points at this position (a missing value). 2. Apply a convolution to the density channel density = CONV θ (density) and a normalized convolution to the signal signal:= CONV θ (signal)/density. The normalized convolution makes sure that the output mostly depends on the scale of the signal rather than the number of observed points. The output channel size is 128 dimensional. The kernel size of CONV θ depends on the image shape and model used (Table 5). We also enforce element-wise positivity of the trainable filter by taking the absolute value of the kernel weights θ before applying the convolution. As discussed in Appendix D.4, the normalization and positivity constraints do not empirically lead to improvements for on-the-grid data. Note that in this setting, E(Z) is [signal, density]. 3. Now we describe the on-the-grid version of ρ(·), which we decompose into two stages. In the first stage, we apply a CNN to [signal, density]. This CNN is composed of residual blocks , each consisting of 2 convolutional layers with ReLU activations and no batch normalization. The number of output channels in each layer is 128. The kernel size is the same across the whole network, but depends on the image shape and model used (Table 5). 4. In the second stage of ρ(·), we apply a shared pointwise MLP: R 128 → R 2C (we use the same architecture as used for the ATTNCNP decoder) to the output of the first stage at each pixel location in the target set. Here C denotes the number of channels in the image. The first C outputs of the MLP are treated as the means of a Gaussian predictive distribution, and the last C outputs are treated as the standard deviations. These then pass through the positivity-enforcing function shown in Equation. Figure 10: Samples from our generated Zero Shot Multi MNIST (ZSMM) data set. In the real world, it is very common to have multiple objects in our field of view which do not interact with each other. Yet, many image data sets in machine learning contain only a single, well-centered object. To evaluate the translation equivariance and generalization capabilities of our model, we introduce the zero-shot multi-MNIST setting. The training set contains all 60000 28 × 28 MNIST training digits centered on a black 56 × 56 . (Figure 10a). For the test set, we randomly sample with replacement 10000 pairs of digits from the MNIST test set, place them on a black 56 × 56 , and translate the digits in such a way that the digits can be arbitrarily close but cannot overlap (Figure 10b). Importantly, the scale of the digits and the image size are the same during training and testing. D.3 ATTNCNP AND CONVCNP QUALITATIVE COMPARISON Figure 11 shows the test log-likelihood distributions of an ATTNCNP and CONVCNP model as well as some qualitative comparisons between the two. Although most mean predictions of both models look relatively similar for SVHN and CelebA32, the real advantage of CONVCNP becomes apparent when testing the generalization capacity of both models. Figure 12 shows CONVCNP and ATTNCNP trained on CelebA32 and tested on a downscaled version of Ellen's famous Oscar selfie. We see that CONVCNP generalizes better in this setting. To understand the importance of the different components of the first layer, we performed an ablation study by removing the density normalization (CONVCNP no norm.), removing the density channel (CONVCNP no dens.), removing the positivity constraints (CONVCNP no abs.), removing the positivity constraints and the normalization (CONVCNP no abs. norm.), and replacing the fully trainable first layer by an EQ kernel similar to the continuous case (CONVCNP EQ). Table 6 shows the following: (i) Appending a density channel helps. (ii) Enforcing the positivity constraint is only important when using a normalized convolution. (iii) Using a less expressive EQ filter does not significantly decrease performance, suggesting that the model might be learning similar filters (Appendix D.5). Figure 13: First filter learned by CONVCNPXL, CONVCNP, and CONVCNP EQ for all our datasets. In the case of RGB images, the plotted filters are for the first channel (red). Note that not all filters are of the same size. As discussed in Appendix D.4, using a less expressive EQ filter does not significantly decrease performance. Figure 13 shows that this happens because the fully trainable kernel learns to approximate the EQ filter. As seen in Table 3, a CONVCNPXL with large receptive field performs significantly worse on the ZSMM task than CONVCNP, which has a smaller receptive field. Figure 14 shows a more detailed comparison of the models, and suggests that CONVCNPXL learns to model non-stationary behaviour, namely that digits in the training set are centred. We hypothesize that this issue stems from the the treatment of the image boundaries. Indeed, if the receptive field is large enough and the padding values are significantly different than the inputs to each convolutional layer, the model can learn position-dependent behaviour by "looking" at the distance from the padded boundaries. For ZSMM, Figure 15 suggests that "circular" padding, where the padding is implied by tiling the image, helps prevent the model from learning non-stationarities, even as the size of the receptive field becomes larger. We hypothesize that this is due to the fact that "circularly" padded values are harder to distinguish from actual values than zeros. We have not tested the effect of padding on other datasets, and note that "circular" padding could in other issues. Figure 15: Effect of the receptive field size on ZSMM's log-likelihood. The line plot shows the mean and standard deviation over 6 runs. The blue curve corresponds to a model with zero padding, while the orange one corresponds to "circular" padding.

Represent the TLDR for finding the corresponding scientific article

We extend deep sets to functional embeddings and Neural Processes to include translation equivariant members

This paper 1) characterizes functions representable by ReLU DNNs, 2) formally studies the benefit of depth in such architectures,  3) gives an algorithm to implement empirical risk minimization to global optimality for two layer ReLU nets.


Question: Represent the article for finding a single-sentence TLDR.

Question: Classical models describe primary visual cortex (V1) as a filter bank of orientation-selective linear-nonlinear (LN) or energy models, but these models fail to predict neural responses to natural stimuli accurately. Recent work shows that convolutional neural networks (CNNs) can be trained to predict V1 activity more accurately, but it remains unclear which features are extracted by V1 neurons beyond orientation selectivity and phase invariance. Here we work towards systematically studying V1 computations by categorizing neurons into groups that perform similar computations. We present a framework for identifying common features independent of individual neurons' orientation selectivity by using a rotation-equivariant convolutional neural network, which automatically extracts every feature at multiple different orientations. We fit this rotation-equivariant CNN to responses of a population of 6000 neurons to natural images recorded in mouse primary visual cortex using two-photon imaging. We show that our rotation-equivariant network outperforms a regular CNN with the same number of feature maps and reveals a number of common features, which are shared by many V1 neurons and are pooled sparsely to predict neural activity. Our findings are a first step towards a powerful new tool to study the nonlinear functional organization of visual cortex. The mammalian retina processes image information using a number of distinct parallel channels consisting of functionally, anatomically, and transcriptomically defined distinct cell types. In the mouse, there are 14 types of bipolar cells BID8, which provide input to 30-50 types of ganglion cells BID2 BID23. In visual cortex, in contrast, it is currently unknown whether excitatory neurons are similarly organized into functionally distinct cell types. A functional classification of V1 neurons would greatly facilitate understanding its computations just like it has for the retina, because we could focus our efforts on identifying the function of a small number of cell types instead of characterizing thousands of anonymous neurons. Recent work proposed a framework for learning functional cell types from data in an unsupervised fashion while optimizing predictive performance of a model that employs a common feature space shared among many neurons BID16. The key insight in this work is that all neurons that perform the same computation but have their receptive fields at different locations, can be represented by a feature map in a convolutional network. Unfortunately, this approach cannot be applied directly to neocortical areas. Neurons in area V1 extract local oriented features such as edges at different orientations, and most image features can appear at arbitrary orientations -just like they can appear at arbitrary locations. Thus, to define functional cell types in V1, we would like to treat orientation as a nuisance parameter (like receptive field location) and learn features independent of orientation. In the present paper, we work towards this goal. While we do not answer the biological question whether there are indeed well-defined clusters of functional cell types in V1, we provide the technical foundation by extending the work of Klindt and colleagues BID16 and introducing a rotation-equivariant convolutional neural network model of V1. We train this model directly on the responses of 6000 mouse V1 neurons to learn a shared feature space, whose features are independent of orientation. We show that this model outperforms state-of-the-art CNNs for system identification and allows predicting V1 responses of thousands of neurons with only 16 learned features. Moreover, for most neurons, pooling from only a small number of features is sufficient for accurate predictions. Functional classification of cell types. Characterizing neurons according to some identified, potentially nonlinear response properties has a long history. For instance, researchers have identified simple and complex cells BID12, pattern-and component-selective cells BID10, end-stopping cells BID24, gain control mechanisms BID5 and many more. However, these properties have been identified using simple stimuli and the models that have been built for them do not apply to natural stimuli in a straightforward way. In addition, most of these studies were done using single-cell electrophysiology, which does not sample cells in an unbiased fashion. Thus, we currently do not know how important certain known features of V1 computation are under natural stimulation conditions and what fraction of neurons express them. Recent work using two-photon imaging in the monkey BID26 started closing this gap by systematically investigating response properties of V1 neurons to an array of complex stimuli and found, for instance, that about half of the neurons in V1 are selective to complex patterns such as curvature, corners and junctions rather than just oriented line segments. Recent work in the retina showed that different types of Ganglion cells BID2 ) -and to some extend also bipolar cells BID9 ) -can be identified based on functional properties. Thus, in the retina there is a relatively clear correspondence of anatomical and genetic cell types and their functional output, and we are getting closer to understanding each retinal cell type's function. Learning shared feature spaces for neural populations. Antolík and colleagues pioneered the idea of learning a shared nonlinear feature representation for a large population of neurons BID1, which others have also used in both retina and V1 BID3 BID14 BID4. BID16 proposed a framework to learn functional cell types in an unsupervised fashion as a by-product of performing system identification. They propose a structurally constrained readout layer, which provides a compact characterization of each neuron's computation. By enforcing this representation to be sparse, they suggest that the learned features may correspond to distinct functional cell types. Our present work builds upon this idea and extends it to be applicable to V1.Rotation equivariance in CNNs. There is a large body of work on equivariant representations BID20. Here we review only the most closely related approches. BID6 introduce group-equivariant CNNs, which include rotation equivariance as a special case and form the basis of our approach. BID31 use a steerable filter basis to learn rotation-equivariant CNNs. We essentially use the same approach, but with a different set of steerable filters (2d Hermite functions instead of circular harmonics). RotEqNet BID18 ) is related to our work in the sense that it also applies multiple rotated versions of each filter to the input. However, instead of maintaining all ing feature maps as inputs for the next layer, they apply an orientation pooling operation, which reduces each set of feature maps to a two-dimensional vector field. Harmonic networks BID32 are an alternative approach that achieves full 360• rotation equivariance by limiting the structure of the convolutional filters. Finally, BID7 The network consists of a rotation-equivariant convolutional core common for all neurons (blue box) and a neuron-specific readout [red box BID16]. Inputs are static images from ImageNet. Prediction targets are responses of 6005 V1 neurons to these images. Rotation equivariance is achieved by using weight sharing across filter orientations. Therefore, eight rotated versions of each filter exist, ing in eight groups of feature maps (depicted by rainbow colored boxes). B. Illustration of weight sharing across orientations for the second and third layer. The previous layer's output consists of eight groups of feature maps (one for each rotation angle). The output is generated by applying a learned filter to each of the input feature maps (here 8 × 16 kernels shown in the first row). Rotated versions (2nd and following rows) are generated by rotating each kernel and cyclically permuting the feature maps. C. Filters are represented in a steerable basis (2d Hermite functions). Functions up to rank 5 are shown here. Network architecture. Our architecture follows that of BID16. The image is first processed by a multi-layer convolutional core shared by all neurons FIG0, blue), which we describe in detail below. The ing feature representation is then turned into a response prediction for each neuron by applying a sparse linear combination across space, followed by a sparse linear combination across features and a pointwise, soft-thresholding nonlinearity FIG0, red). The entire model is trained end-to-end directly on predicting neural responses, without employing transfer learning or pre-training on auxiliary tasks BID33 BID4.This model has a number of desirable properties in terms of data efficiency and interpretability. The separation into convolutional core and readout, together with the strong structural constraints on the readout pushes all the'heavy lifting' into the core while the readout weights (spatial mask and feature weights) provide a relatively low-dimensional characterization of each neuron's function. Because the core is shared among all neurons (here: thousands), many of which implement similar functions, we can learn complex non-linear functions very accurately. Equivariance. A function f: X → Y is called equivariant with respect to a group of transformations Π if such transformations of the input lead to predictable transformations of the output. Formally, for every π ∈ Π there is a transformation ψ ∈ Ψ such that for every DISPLAYFORM0 CNNs are shift-equivariant by construction, meaning that every translation of the image leads to a matching translation in the feature maps (i. e. π and ψ are both translation by a number of pixels horizontally and vertically). Shift equivariance is a useful property for neural system identification, because it allows us to represent many neurons that perform similar computations but in different locations in space by a single convolutional feature map instead of learning each neurons' nonlinear input-output function'from scratch.'Rotation-equivariant CNN. Neurons in V1 do not only perform similar functions at different locations, but also extract similar features with different orientations. Thus, for modeling populations of V1 neurons, learning a representation that is equivariant to rotation in addition to translation would be desirable. To achieve rotation equivariance, we use group convolutions BID6 ).Here we use the group of all rotations by multiples of 45 •. That is, for each convolutional filter in the first layer, we have eight rotated copies, each of which produces one feature map FIG0. Thus, if we learn 16 different filters in the first layer, it will have a total of 8 × 16 = 128 feature maps. Formally, if π is a rotation of the image by, say, 45•, then ψ is a rotation of the feature maps by also 45• combined with a cyclic permutation of the feature maps by one rotation step. For the second (and all subsequent) layers, the procedure becomes a bit more involved. Sticking with the numbers from above, for every feature in the second layer we now learn 128 filters -one for each input feature map FIG0. To preserve rotation equivariance, we need to create all rotated copies of these filters and cyclically permute the feature maps such that each rotated filter receives the rotated version of the input (depicted by the colored boxed in FIG0, second and following rows).To implement weight sharing across filter orientation without aliasing artifacts, we represent the filters in a steerable basis BID31. We use the two-dimensional Hermite functions in polar coordinates, which form a steerable, orthonormal basis (FIG0 ; see also BID28 BID11 . For filters of size k, we use all 2d Hermite functions up to rank k, which means we have k(k + 1)/2 basis functions. We sample each filter at twice the resolution and then downsample by 2×2 average pooling to reduce aliasing. Neural data. We recorded the responses of 6005 excitatory neurons in primary visual cortex (layers 2/3 and 4) from one mouse by taking two consecutive scans with a large-field-of-view twophoton mesoscope BID25. Activity was measured using the genetically encoded calcium indicator GCaMP6f. V1 was targeted based on anatomical location as verified by numerous previous experiments performing retinotopic mapping using intrinsic imaging. We selected cells based on a classifier for somata on the segmented cell masks and deconvolved their fluorescence traces BID21. We did not filter cells according to visual responsiveness. The aquisition frame rate was 4.8 Hz in both scans. We monitored pupil position, dilation, and absolute running speed of the animal. However, because eye movements were rare and we are primarily interested in the average response of neurons given the visual stimulus, we did not further take into account eye position and running speed. Visual stimuli. Stimuli consisted of 5500 images taken from ImageNet BID22, cropped to fit a 16:9 monitor, and converted to gray-scale. The screen was 55 × 31 cm at a distance of 15 cm, covering roughly 120 DISPLAYFORM0 In each scan, we showed 5400 of these images once (training and validation set) and the remaining 100 images 20 times each (test set). Each image was presented for 500ms followed by a blank screen lasting between 500ms and 1s. For each neuron, we extract the accumulated activity between 50ms and 550ms after stimulus onset using a Hamming window. Preprocessing. We rescaled the images to 64 × 36 pixels and standardized them by subtracting the mean over all images in the training set and all pixels from each image and dividing by the standard deviation (also taken over all images in the training set and all pixels). We divided the responses of each neuron by its standard deviation over time. We did not center the neural responses, because they are non-negative after deconvolution and zero has a clear meaning. Model fitting and evaluation. We initialized all weights randomly from a truncated normal distribution with mean zero and standard deviation 0.01. The biases of the batch normalization layers were initially set to zero. In contrast, we set the biases in each neuron's readout to a non-zero initial value, since the neural responses are not centered on zero. We initialized these biases such that the initial model prediction was on average half the average response of each neuron. To fit the models, we used the Adam Optimizer BID15 with an initial learning rate of 0.002, a single learning rate decay and early stopping. We monitored the validation loss every 50 iterations and decreased the learning rate once by a factor of 10 when the validation loss had not decreased for five validation steps in a row. We then further optimized the model until the same criterion is reached again. To evaluate the models, we use the following procedure. For each neuron we compute the Pearson correlation coefficient between the model prediction and the average response over the 20 repetitions of the 100 test images. We then average the correlations over all neurons. This approach tells us how well the model predicts the average response of neurons to a given stimulus, ignoring trial-to-trial variability (which is interesting in itself, but not the focus of the present work).Architecture details. Our architecture consists of three convolutional layers with filter sizes of 13, 5 and 5 pixels, respectively. Thus, the receptive fields of the CNN's last layer's units were 21 px, corresponding to ∼ 60• and covering both classical and extra-classical receptive field. We use zero padding such that the feature maps maintain the same size across layers. We use 16 filter sets (i. e. 128 feature maps) in the first two layers and 8 to 48 filter sets in the third layer (number cross-validated, see below). After each layer, we apply batch normalization followed by a learned scale and bias term for each feature map.1 After the first and second layer, but not after the third, we apply a soft-thresholding nonlinearity f (x) = log(1 + exp(x)). The feature maps of the third layer then provide the input for each neuron's readout, which consists of a linear combination first over space and then over features, followed by an added bias term and a final soft-thresholding nonlinearity. Thus, each neuron implements a cascade of three LN operations. Regularization. We use the same three regularizers as BID16. For smoothness of convolution kernels we set the relative weight of the first layer to twice as strong as in the second and third layer to account for the larger kernels. We apply group sparsity to the convolution kernels of the second and third layer. We regularize the spatial masks and feature weights in the readout to be sparse by applying the L1 penalty on the 3d tensor that from taking their outer tensor product. The weights of all three regularizers are cross-validated as described in the next paragraph. Model selection. We cross-validated over the number of filter sets in the third layer, ranging from 8 to 48 (i. e. 64 to 384 feature maps 2) and the strength of the regularizers. For each architecture, we fit 32 models with different initial conditions and randomly drawn hyperparameters (smoothness of filters: 0.001-0.03, group sparsity of filters: 0.001-0.1, sparsity of readout: 0.005-0.03) and chose the best one according to its loss on the validation set (i. e. not using the test set). For all baseline and control models (see below), we also cross-validated over 32 randomly sampled sets of hyperparameters drawn from the same range of values. Baseline and control experiments. As a baseline, we fit a number of regular CNNs without rotation equivariance. These models are completely identical in terms of number of layers, filter sizes, readout layer and fitting procedure, except for the weight sharing constraint across orientation. We fit models with the same number of feature maps as their rotation-equivariant counterparts as well as smaller models with fewer feature maps (see Table 1).Previous work has enforced the feature weights in the readout to be positive BID16 ). Because we do not enforce such constraint in the present work, we ran a control experiment, in which we enforce the readout weights (both masks and feature weights) to be positive. Sparsity of the spatial masks is well justified, because we know that receptive fields are localized. However, it is unclear whether sparsity on the feature weights is desirable, since the function each neuron implements may not be well aligned with the coordinate axes spanned by the small number of features we learn. Thus, we also fit a model without the sparsity regularizer for the feature weights. 1 In our experiments the network did not implement exact rotation equivariance, because batch normalization was applied to each feature map individually instead of jointly to all rotated versions. We therefore re-ran a subset of experiments where we corrected this issue and verified that model performance was indistinguishable.2 384 feature maps was the maximum we could fit into 16 GB of available GPU memory. Number of feature maps in last layer Table 1: Average correlation on test set of our rotation-equivariant CNN, two controls, and several regular CNNs as baselines. The three numbers for the CNNs are the number of feature maps in each layer; other parameters are identical to the rotation-equivariant CNN. Standard deviations are across the top-five models of each type. Tools. We performed our model fitting and analyses using DataJoint BID34, Numpy/Scipy BID29 ), Matplotlib (, Seaborn BID30, Jupyter BID17, Tensorflow BID0, and Docker BID19 .Availability of code and models. Code to reproduce all experiments and models as well as pretrained models are available at https://github.com/aecker/cnn-sys-ident. Architecture search and model selection. We start by evaluating the predictive performance of our rotation-equivariant CNN on a dataset of 6005 neurons from mouse primary visual cortex. First, we performed an initial architecture search, drawing inspiration from earlier work BID16 and exploring different numbers of layers and feature maps and different filter sizes. We settled on an architecture with three convolutional layers with filter sizes of 13, 5 and 5 pixels, respectively, and used 16 filter sets (i. e. 128 feature maps) in the first two layers. For the third layer, which provides the nonlinear feature space from which all neurons pool linearly, we cross-validated over the number of features, ranging from 8 to 48 (i. e. 64 to 384 feature maps).Our model achieved an average correlation on the test set of 0.47. The performance improved with the number of features in the third layer, but the improvement was small beyond 16 features (Fig. 2). For the following analyses we therefore focus on this model with 16 features as a compromise between model complexity and performance. It is encouraging that a model with only 16 features is sufficient to accurately model a population as large as 6000 neurons. Earlier work modeling V1 responses used a shared feature space of 10-20 BID1 or 48 features BID16 ) for populations of 50-100 neurons, which reduced the dimensionality of the feature space by a factor of 2-5 compared to the number of neurons. Here, we reduce the dimensionality by a factor of 375 while achieving a similar level of performance. Rotation-equivariant CNN outperforms regular CNN. To compare our new model to existing approaches, we also fit a regular CNN with identical architecture except for the weight sharing constraint across orientations. In addition, we fit a number of smaller CNNs with fewer feature maps in each layer, which are more similar in terms of number of parameters, but potentially have less expressive power. Our rotation-equivariant CNN outperforms all baselines (Table 1) and generally requires less data FIG1, showing that enforcing weight sharing across orientations is not only potentially useful for interpreting the model (as we show below), but also serves as a good regularizer to fit a larger, more expressive model. Feature space generalizes to unseen neurons. To show that our network learns common features of V1 neurons, we excluded half of the neurons when fitting the network. We then fixed the rotationequivariant convolutional core and trained only the readout (spatial mask and feature weights) for the other half of the neurons. The ing test correlation for these neurons (0.46) was indistinguishable from that of the full model (0.47), showing that the learned features transfer to neurons not used to train the feature space. Feature weights are sparse. The intuition behind the sparse, factorized readout layer is that the spatial mask encodes the receptive field location of each neuron while the feature weights parameterize the neuron's nonlinear computation. We now ask how the neurons are organized within this 16-dimensional function space. On the one extreme of the spectrum, each neuron could pick a random direction, in which case sparsity would not be the right prior and the feature weights should be dense. On the other end of the spectrum, there could be a discrete number of functional cell types. In this case, each cell type would be represented by a single feature and the feature weights should be maximally sparse, i. e. one-hot encodings of the cell type. To analyze the feature weights, we first marginalize them over orientation. To do so, we take the sum of squares over all 8 orientations for each of the 16 features and then normalize them such that the energy of the 16 features sums to one. We find that the feature weights are indeed quite sparse •.(FIG2 : most of the neurons use only 1-5 features FIG2) and the strongest feature captures more than 50% of the energy for 63% of the neurons. To ensure that the sparsity of the feature weights is a property of the data and not a trivial of our L 1 penalty, we performed an ablation study, where we fit a model that applies sparsity only on the spatial masks, but uses L 2 regularization for the feature weights. This model performed worse than the original model (Table 1) and produced a significantly denser weight distribution FIG2, suggesting that sparsity is indeed a justified assumption. Feature weights have consistent sign. There is a second line of evidence that the features learned by our model are meaningful beyond just providing an arbitrary basis of a 16-dimensional feature space in which neurons are placed at random: the weights that different neurons assign to any given feature have remarkably consistent signs FIG4 ). For 11 of 16 features, more than 90% of the non-zero weights have the same sign. 3 Thus, the negation of one feature that drives a large group of neurons is not a feature that drives many neurons. Visualization of the different groups of neurons. Having presented evidence that the features identified by our model do indeed represent a simple and compact, yet accurate description of a large population of neurons, we now ask what these features are. As a first step, we group all neurons into one of 16 groups based on their strongest feature weight. This approach is obviously to some extent arbitrary. By no means to we want to argue that there are exactly 16 types of excitatory neurons in V1 or that excitatory neurons in V1 can even be classified into distinct functional types. Nevertheless, we believe that the 16 groups defined in this way are useful in a practical sense, because they represent features that cover more than 50% of the variance of a large fraction of neurons and therefore yield a very compact description of the most important features of V1 computation across the entire excitatory neuronal population. To visualize what each of these features compute, we pick the 16 most representative examples 4 of each group and use the model to approximate their linear receptive field (RF). To this end, we compute the model gradient at a gray image FIG5. The crucial point of this plot is that it shows that the premise of rotation equivariance holds: cells are clustered according to similar spatial patterns, but independent of their preferred orientation. As expected, most linear RFs resemble Gabor filters, which show differences in symmetry (odd: #2, #11, #15 vs even: #3, #6, #12) as well polarity (#3 vs. #6, #12), while some groups exhibit center-surround structure (#4, #14, #16). We developed a rotation-equivariant convolutional neural network model of V1 that allows us to characterize and study V1 computation independent of orientation preference. Although the visual system is not equivariant to rotation -there are known biases in the distribution of preferred orientations -, enforcing weight sharing across orientations allowed us to fit larger, more expressive models given a limited dataset. While our work lays out the technical foundation, we only scratched the surface of the many biological questions that can now be addressed. Future work will have to investigate the learned features in much more detail, test to what extent they generalize across recording sessions and animals, whether they are consistent across changes in the architecture andmost importantly -whether neurons in V1 indeed cluster into distinct, well-defined functional types and this organization finds any resemblance in anatomical or genetic properties BID27 of the neurons recorded.

represent this text to find a much much longer article about it with say hundreds-thousands of words

A rotation-equivariant CNN model of V1 that outperforms previous models and suggest functional groupings of V1 neurons.

Represent the short summary for finding the corresponding scientific article (long)

We detect statistical interactions captured by a feedforward multilayer neural network by directly interpreting its learned weights.


Question: Represent the article for finding a single-sentence TLDR

Question: In classic papers, Zellner demonstrated that Bayesian inference could be derived as the solution to an information theoretic functional. Below we derive a generalized form of this functional as a variational lower bound of a predictive information bottleneck objective. This generalized functional encompasses most modern inference procedures and suggests novel ones. Consider a data generating process φ ∼ p(φ) from which we have some N draws that constitute our training set, x P = {x 1, x 2, . . ., x N} ∼ p(x|φ). We can also imagine (potentially infinitely many) future draws from this same process x F = {x N +1, . . .} ∼ p(x|φ). The predictive information I(x P ; x F) 1 gives a unique measure of the complexity of a data generating process . The goal of learning is to capture this complexity. To perform learning, we form a global representation of the dataset p(θ|x P). This can be thought of as a learning algorithm, that, given a set of observations, produces a summary statistic of the dataset that we hope is useful for predicting future draws from the same process. This algorithm could be deterministic or more generally, stochastic. For example, imagine training a neural network on some data with stochastic gradient descent. Here the training data would be x P, the test data x F and the neural network parameters would be θ. Our training procedure implicitly samples from the distribution p(θ|x P). How do we judge the utility of this learned global representation? The mutual information I(θ; x F) quantifies the amount of information our representation captures about future draws. 2 To maximize learning we therefore aim to maximize this quantity. 1. We use I(x; y) for the mutual information between two random variables: 2. It is interesting to note that in the limit of an infinite number of future draws, I(θ; x F) approaches I(θ; φ). Therefore, the amount of information we have about an infinite number of future draws from This is, of course, only interesting if we constrain how expressive our global representation is, for otherwise we could simply retain the full dataset. The amount of information retained about the observed data: I(θ; x P) is a direct measure of our representation's complexity. The bits a learner extracts from data provides upper bounds on generalization . Combined, these motivate the predictive information bottleneck objective, a generalized information bottleneck : We can turn this into an unconstrained optimization problem with the use of a Lagrange multiplier β: max While this objective seems wholly out of reach, we can make progress by noting that our random variables satisfy the Markov chain: x F ← φ → x P → θ, in which θ and x F are conditionally independent given x P: This implies: and the equivalent unconstrained optimization problem: 3 The first term here: I(θ; x P |x F) is the residual information between our global representation and the dataset after we condition on full knowledge of the data generating procedure. This is a direct measure of the inefficiency of our proposed representation. Simple variational bounds can be derived for this objective, just as was done for the (local) information bottleneck objective in. First, we demonstrate a variational upper bound on I(θ; the process is the same as the amount of information we have about the nature and identity of the data generating process itself. 3. A similar transformation for the (local) variational information bottleneck appeared in. 4. · is used to denote expectations, and unless denoted otherwise with respect to the full joint density Here we upper bound the residual information by using a variational approximation to p(θ|x F), the marginal of our global representation over all datasets drawn from the same data generating procedure. Any distribution q(θ) independent of x F suffices. Next we variationally lower bound I(θ; x P) with: The entropy of the training data H(x P) is a constant outside of our control that can be ignored. Here we variationally approximate the "posterior" of our global representation with a factorized "likelihood": i q(x i |θ) = q(x P |θ) ∼ p(x P |θ). Notice that while p(x P |θ) will not factorize in general, we can certainly consider a family of variational approximations that do. Combining these variational bounds, we generate the objective: We have thus derived, as a variational lower bound on the predictive information bottleneck, the objective postulates (with β = 1) is satisfied for inference procedures that optimally process information. demonstrates, this encompasses a wide array of modern inference procedures, including Generalized Bayesian Inference and a generalized Variational Inference, dubbed Gibbs VI . 5 Below we highlight some of these and other connections. If, in Equation, we identity q(θ) with a fixed prior and q(x|θ) with a fixed likelihood of a generative model, optimizing this objective for p(θ|x P) in the space of all probability densities gives the generalized Boltzmann distribution : where Z is the partition function. 6 This is a generalized form of Bayesian Inference called the power likelihood . Here the inverse temperature β acts as a Lagrange multiplier controlling the trade-off between the amount of information we retain about our observed data (I(θ; x P)) and how much predictive information we capture (I(θ; x F)). As β → ∞ (temperature goes to zero), we recover the maximum likelihood solution. At β = 1 (temperature = 1) we recover ordinary Bayesian inference. As β → 0 (temperature goes to infinity), we recover just prior predictive inference that ignores the data entirely. These limits are summarized in Table 1. 5. To incorporate the Generalized VI with divergence measures other than KL, we need only replace our mutual informations (which are KL based) with their corresponding generalizations. Table 1: Power Bayes can be recovered as a variational lower bound on the predictive information bottleneck objective (Equation). More generally, notice that in Equation the densities q(x|θ) and q(θ) are not literally the likelihood and prior of a generative model, they are variational approximations that we have complete freedom to specify. This allows us to describe other more generalized forms of Bayesian inference such as Divergence Bayes or the full Generalized Bayes provided we can interpret the chosen loss function as a conditional distribution. If we limit the domain of p(θ|x P) to a restricted family of parametric distributions, we immediately recover not only standard variational inference, but a broad generalization known as Gibbs Variational Inference (; ;). Furthermore, nothing prevents us from making q(x|θ) or q(θ) themselves parametric and simultaneously optimizing those. Optimizing the prior with a fixed likelihood, unconstrained p(θ|x P), and β = 1 the objective mirrors Empirical Bayesian approaches, including the notion of reference priors . Alternatively, optimizing a parametric likelihood with a parametric representation p(θ|x P), fixed prior, and β = 1 equates to a Neural Process . Consider next data augmentation, where we have some stochastic process that modifies our data with implicit conditional density t(x |x). If the augmentation procedure is centered about zero so that x t(x |x) = x and our chosen likelihood function is concave, then we have: which maintains our bound. For example, for an exponential family likelihood and any centered augmentation procedure (like additive mean zero noise), doing generalized Bayesian inference on an augmented dataset is also a lower bound on the predictive information bottleneck objective. We have shown that a wide range of existing inference techniques are variational lower bounds on a single predictive information bottleneck objective. This connection highlights the drawbacks of these traditional forms of inference. In all cases considered in the previous section, we made two choices that loosened our variational bounds. First, we approximated p(x P |θ), with a factorized approximation q(x P |θ) = i q(x i |θ). Second, we approximated the future conditional marginal p(θ|x F) = dx P p(θ|x P)p(x P |x F) as an unconditional "prior". Neither of these approximations is necessary. For example, consider the following tighter "prior": q(θ|x F) ∼ dx P p(θ|x P)q(x P |x F). Here we reuse a tractable global representation p(θ|x P) and instead create a variational approximation to the density of alternative datasets drawn from the same process: q(x P |x F). We believe this information-theoretic, representation-first perspective on learning has the potential to motivate new and better forms of inference. 7

Rederive a wide class of inference procedures from an global information bottleneck objective.

represent the short summary for finding the corresponding scientific article (long).

We show that posterior collapse in linear VAEs is caused entirely by marginal log-likelihood (not ELBO). Experiments on deep VAEs suggest a similar phenomenon is at play.