Source-linked AI summary
Convolutional Conditional Neural Processes
Jonathan Gordon, Wessel P. Bruinsma, Andrew Y. K. Foong, James Requeima, Yann Dubois, Richard E. Turner
TL;DR
Neural Processes need translation equivariance for domains such as time series, spatial data, and images, but existing models generally learn this structure from data. ConvCNP embeds sets into function space, establishes a convolutional-deep-set representation for translation-equivariant mappings, and reports strong performance across several benchmarks, including challenging out-of-domain tasks.
Problem
Existing Neural Processes do not directly encode translation equivariance, although it is an important inductive bias for time series, spatial data, and images.
Method
ConvCNP extends Neural Processes with functional set representations and a convolutional-deep-set construction for translation-equivariant mappings.
Results
ConvCNP demonstrates excellent performance on several synthetic and real-world benchmarks and zero-shot generalization to challenging out-of-domain tasks.
Takeaways & Limitations
Embedding translation equivariance enables ConvCNP to generalize zero-shot to challenging out-of-domain tasks.
Takeaways & Limitations
The predictive distribution assumes conditional independence given the context set, producing samples that lack correlations and appear noisy.
Abstract
from arXiv · showhide
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 a finite-dimensional vector space. 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.
1 INTRODUCTION
Neural Processes support flexible prediction from context sets, but existing models must learn translation equivariance from data. ConvCNP addresses this by extending neural-process representations to function spaces and evaluating the resulting model on multiple benchmarks.
- Neural Processes define conditional predictions from observed context sets and support meta-learning and transfer learning.CNPs can be trained by maximum likelihood in a way that matches test-time use.
- Time series, spatial data, and images often require predictions to translate correspondingly when input locations are translated.Existing NP models must learn this structure from data, which is sample- and parameter-inefficient.
- CNN-style translation equivariance is difficult to transfer directly because NP data may be off-grid, partially observed, and embedded in finite-dimensional vectors.These differences separate NP settings from standard CNN assumptions.
- ConvCNP extends the NP family by using functional representations to express translation-equivariant models.The paper presents this as a new NP member accounting for translation equivariance.
- The paper provides a representation theorem, extends NPs with translation equivariance, and evaluates ConvCNP on synthetic and real-world benchmarks.The reported evaluation covers several benchmarks and is described as demonstrating excellent performance.
2 BACKGROUND AND FORMAL PROBLEM STATEMENT
The paper formalizes neural-process inputs as finite sets of input-output pairs and distinguishes permutation invariance from translation equivariance. CNPs encode such sets into finite-dimensional vectors, motivating a functional representation for translation-aware modeling.
- Notation: The input space is X = R^d, outputs lie in compact Y ⊆ R^d′, and Z contains finite collections of input-output pairs.Z_M denotes sets of size M, Z_≤M sets of at most M, and Z all finite sets.
- Conditional Neural Processes: CNPs model p(y|x, Z) through an encoder E: Z → R^e and decoder ρ that produces continuous bounded functions.Unlike latent NPs, the described CNP models do not use latent variables.
- Set representations: Because set elements are unordered, functions on Z are required to be permutation invariant.The paper defines invariance separately for fixed-size sets and for all finite sets.
- Translation equivariance: Translation equivariance requires predictions to shift correspondingly when all input locations are translated by τ.The paper formalizes the transformation of context sets and mappings under translations.
3 CONVOLUTIONAL DEEP SETS
ConvDeepSets represent sets as functions, allowing permutation-invariant and translation-equivariant mappings to coexist. Theorem 1 establishes that the proposed functional form is equivalent to any continuous mapping with both properties under stated domain conditions.
- Functional representations: Functional mappings send sets Z to an appropriate function space H, with E(Z) serving as the set’s functional representation.This replaces ordinary vector-valued set embeddings with function-valued embeddings.
- Representation theorem: Any continuous mapping satisfying permutation invariance and translation equivariance has a ConvDeepSet representation using continuous φ, ψ, and translation-equivariant ρ.The output map ρ operates between function spaces.
- Representation theorem: Theorem 1 extends the finite-dimensional set-representation result to infinite-dimensional RKHS embeddings and varying-size finite sets.The RKHS formalism supports translation equivariance and naturally handles finite sets of different sizes.
- Construction: For real-world time series and image data, one observation per input location commonly corresponds to multiplicity one.This is the motivating practical case for the theorem’s varying-size set treatment.
- Construction: The construction uses a positive-definite kernel ψ, while φ encodes output values through a finite power series determined by multiplicity.For multiplicity K, φ expands through order K; practical choices are informed by the theorem and its proof.
4 CONVOLUTIONAL CONDITIONAL NEURAL PROCESSES
ConvCNP implements the functional representation with kernels and convolutional architectures for both off-grid and on-grid data. It normalizes context information by density and produces predictive means and variances through translation-aware neural mappings.
- Model architecture: ConvCNP models the conditional distribution using a CONVDEEPSET representation, with separate implementations for on-grid and off-grid data.The observed data form the context set Z used by the representation Φ.
- Off-the-grid data: For multiplicity-one data, φ(y) = [1 y]^⊤ adds an observation indicator alongside the output value.The first functional channel distinguishes an unobserved location from an observed location whose value is zero.
- Off-the-grid data: Normalized convolution divides non-density channels by the density channel, improving performance when input-location density varies substantially.The normalization can be reversed by ρ and is therefore not restrictive.
- Off-the-grid data: Off-grid ConvCNPs use a learnable exponentiated-quadratic kernel to form a functional representation, then discretize it for a CNN and reconstruct a continuous function.The reconstruction uses kernel-weighted basis functions and is approximately translation equivariant for sufficiently large length scales.
- On-the-grid data: On-grid image ConvCNPs implement E with a context mask and convolution, followed by a CNN and shared pointwise MLP for Gaussian means and standard deviations.The image implementation uses pixel locations as the discretization grid.
5 EXPERIMENTS AND RESULTS
Experiments evaluate CONVCNP on irregular time series, simulations, and image completion, testing both in-domain performance and translation-equivariant generalization. Across these settings, CONVCNP often outperforms ATTNCNP and extrapolates beyond training distributions, while also reducing computational costs.
- Synthetic 1D experiments: CONVCNP significantly outperforms CNP and ATTNCNP on synthetic one-dimensional tasks, even without extrapolation and despite using fewer parameters.The comparison covers EQ-kernel Gaussian-process-like data, Matérn processes, weakly periodic kernels, and sawtooth functions over 1000 tasks.
- Synthetic 1D experiments: Translation equivariance enables CONVCNP to generalize to observations outside the training range, whereas ATTNCNP produces unreasonable predictions.The comparison includes challenging Matérn and sawtooth functions with data observed beyond the training interval.
- PLAsTiCC experiments: On PLAsTiCC, CONVCNP and the comparison approach have similar mean squared error, but CONVCNP achieves a much higher held-out log-likelihood.PLAsTiCC contains six-dimensional astronomical light curves generated under realistic observational conditions.
- Predator-prey models: Sim2Real: CONVCNP accurately interpolates the Hudson’s Bay lynx–hare series after training on Lotka–Volterra simulations, while ATTNCNP could not be successfully trained for this task.The experiment tests deployment with real-world context data after simulation training.
- 2D image completion experiments: On standard image benchmarks, CONVCNP significantly outperforms ATTNCNP with a large receptive field and is at least as good with a small receptive field.The benchmarks include MNIST, SVHN, and 32×32 and 64×64 CelebA image-completion tasks.
- 2D image completion experiments: CONVCNP significantly outperforms ATTNCNP on zero-shot multi-MNIST and reconstructs a different-shaped natural image with 95% of pixels removed.Very large receptive fields can reduce zero-shot multi-MNIST performance, while CONVCNP still gives a qualitatively reasonable CelebA reconstruction.
- Computational efficiency: A single self-attention layer has quadratic memory and time complexity in input count, whereas a convolutional layer scales linearly.On 32×32 MNIST, CONVCNPXL uses 945MB of VRAM versus 5839MB for ATTNCNP at batch size 16.
6 RELATED WORK AND DISCUSSION
The paper positions ConvCNP as a functional, translation-equivariant extension of neural-process representations and relates it to set embeddings, point-cloud models, and predictive-distribution design choices.
- Neural Processes: ConvCNP extends neural-process models by embedding sets into function space to achieve translation equivariance.The paper presents this as a new CNP-family member and identifies extensions to neural processes and function representation on sets.
- Deep sets: Functional embeddings alleviate fixed-size and dimensionality restrictions identified in existing deep-set theory.The paper states that appropriate function-space embeddings address both restrictions.
- Point-cloud models: ConvCNP resembles recent translation-equivariant point-cloud models but differs in how its ψ component is implemented.The supplied discussion begins the comparison by contrasting ConvCNP with an MLP-based ψ implementation.
- Correlated samples and consistency under marginalization: ConvCNP predictions are conditionally independent given the context, so generated samples lack correlations and can appear noisy.The discussion identifies autoregressive models and latent-variable or exchangeable neural processes as possible routes toward dependent samples.
- Set representations: The permutation-based distance d([x], [y]) = minπ∈Sn ∥x −πy∥An is established as a metric, and the quotient topology is shown to be metrizable with d.Supporting propositions establish continuity, openness, and closed-set properties for the canonical quotient map.
- Set representations: The quotient space of sets under permutations provides a natural representation for permutation-invariant functions.The paper defines permutation classes and studies their metric and quotient-topology properties.
A.3 PROOF OF THEOREM 1
Theorem 1 characterizes permutation-invariant and translation-equivariant set functions through a functional embedding followed by a continuous translation-equivariant map. The proof establishes both necessity and sufficiency of this representation.
- Construction: The embedding uses a positive-definite RKHS kernel and a feature map φ(y) containing powers of y through order K.These choices provide the functional representation used to construct E and ρ.
- Theorem statement: Theorem 1 states that continuity, permutation invariance, and translation equivariance are equivalent to the representation Φ(Z) = ρ(E(Z)).The encoder and outer map are defined in the theorem’s representational form.
- Sufficiency: The sufficiency proof derives permutation invariance from commutative and associative addition and verifies translation equivariance directly.The domain is assumed closed under permutations and translations, making both properties well defined.
- Necessity: The proof defines ρ as Φ composed with the inverse embedding E−1, then uses continuity and translation equivariance of both components.Stationarity of ψ makes E−1 translation equivariant, so their composition inherits equivariance.
- Extension: The resulting map ρ may be continuously extended to the full function space, with variants of the extension theorem preserving translation equivariance.This extension is described as a remark following the proof.
C.1 CNN ARCHITECTURES
The ConvCNP experiments use EQ-kernel functional representations with small and large convolutional architectures. The receptive field depends jointly on filter width and discretization spacing.
- Kernel setup: All architectures use learnable EQ kernels for the input embedding and final output layer, with length scales initialized from discretization spacing.The initialization is twice the spacing between discretization points.
- Receptive field: The receptive field equals the product of CNN filter width and discretization spacing, so increasing point density shrinks it for fixed filter width.Depthwise-separable convolutions are proposed as an untested way to increase widths without as many parameters.
- CONVCNP: The CONVCNP uses a 4-layer, stride-1 CNN with kernel size 5, zero padding of 2, and channel widths [16, 32, 16, 2].The final channels are processed by the EQ-based output layer as mean and standard deviation channels.
- CONVCNPXL: CONVCNPXL uses a 12-layer UNet-inspired architecture with skip connections, doubling channels for six layers and halving them for six layers.Skip connections concatenate activations from earlier and later layers.
C.2 SYNTHETIC 1D EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
The synthetic experiments evaluate ConvCNPs on Gaussian-process and sawtooth regression with randomly sampled contexts and targets. Additional figures visualize predictions within and beyond the training range.
- Training setup: Context and target counts are uniformly sampled between 3 and 50 for the general synthetic experiments, with input locations sampled from [−2, 2].Models are trained for 200 epochs in this setting.
- Sawtooth task: Sawtooth functions fix amplitude to one while sampling frequency from [3, 5], truncation order K from [10, 20], and shifts from [−5, 5].Because the task is harder, context and target sizes are sampled over [3, 100].
C.3 PLASTICC EXPERIMENTAL DETAILS
The PLAsTiCC experiments train and test ConvCNPs by randomly partitioning available series points into context and target sets. The larger ConvCNPXL model is used when the smaller model underfits.
- Evaluation procedure: Training and testing randomly select the number of context points from 1 to the number available in each series, assigning remaining points as targets.Series usually contain 10–30 points per bandwidth.
- Training procedure: Testing uses batch size 1 and computes statistics over 1000 evaluations after 200 training epochs with 1024 batches of size 4 per epoch.The same passage describes the context-target partitioning used during training and testing.
- Prediction settings: Figure 8 illustrates model predictions when data are presented within the training range, outside it without observations, and outside it with observations.The figure caption defines these three prediction settings across its columns.
- Model selection: Only CONVCNPXL is used for this dataset because the smaller CONVCNP was found to underfit.The functional embedding is discretized at 256 points per unit.
C.4 PREDATOR–PREY EXPERIMENTAL DETAILS
The experiments combine simulated predator–prey data, real lynx–hare interpolation, and image tasks to assess ConvCNP prediction and generalization. The image model represents masked observations with signal and density channels, then uses convolutions and a pointwise decoder for target predictions.
- Predator–prey simulation: The simulated populations are scaled by 2/7, while series longer than 100 time units, exceeding 10000 events, or containing an entirely zero population are removed.These processing choices align the simulations more closely with the lynx–hare data and exclude unsuitable series.
- Predator–prey evaluation: Evaluation on the Hudson’s Bay lynx–hare data removes 18 points for targets and uses the remaining 72 as context, testing interpolation and uncertainty with missing data.Only the CONVCNP is used in this setting because CONVCNPXL was found to overfit.
- Image formulation: For images, on-the-grid CONVCNP replaces context point sets with an image, context mask, and target mask, making the formulation natural to standard deep-learning libraries.The context mask selects observed pixels and is appended as a density channel alongside the masked signal.
- Image architecture: The image encoder applies convolution to density and normalized convolution to signal, producing a 128-dimensional representation whose scale depends mainly on signal rather than observation count.The decoder first applies a residual CNN, then a shared pointwise MLP at target pixels; its outputs parameterize Gaussian means and standard deviations.
- Image evaluation: In the zero-shot multi-MNIST setting, models train on centered single digits but test on translated non-overlapping digit pairs; CONVCNP generalizes better than ATTNCNP on a downscaled Oscar selfie.Training and testing retain the same digit scale and image size in ZSMM, while the selfie test changes resolution and content.
D.4 ABLATION STUDY: FIRST LAYER
The first-layer ablation study tests density normalization, the density channel, positivity constraints, and an EQ filter. The reported findings identify the density channel as beneficial and show that an EQ filter can retain performance.
- Evaluation constraint: The CelebA32-to-Oscar-selfie comparison uses a downscaled test image and low-resolution training data because ATTNCNP has a large memory footprint.The reconstruction is worse than the corresponding in-domain result despite the larger context set.
- Ablation findings: Appending a density channel helps image-model performance in the first-layer ablation study.The density channel distinguishes unobserved locations from observed zero-valued observations.
- Ablation findings: Positivity constraints are important only when the first layer uses a normalized convolution.The ablation separately removes positivity constraints and normalization to isolate their interaction.
- Ablation findings: Replacing the fully trainable first-layer filter with an EQ filter does not significantly decrease performance.The ablation compares CONVCNP EQ with variants removing normalization, density, or positivity constraints.
D.5 QUALITATIVE ANALYSIS OF THE FIRST FILTER
The qualitative analysis examines receptive fields, learned first filters, and padding on zero-shot multi-MNIST. Large receptive fields can hurt performance, while circular padding may reduce learned positional non-stationarity.
- Learned filters: Figure 13 compares the first filters learned by CONVCNPXL, CONVCNP, and CONVCNP EQ across datasets, using the red channel for RGB images.The plotted filters are not necessarily the same size.
- Learned filters: The fully trainable kernel learns to approximate the EQ filter, explaining why the less expressive EQ filter does not significantly reduce performance.This links the qualitative filter comparison to the ablation result.
- Receptive fields: CONVCNPXL with a large receptive field performs significantly worse than the smaller-receptive-field CONVCNP on ZSMM.The comparison suggests that CONVCNPXL learns non-stationary behavior because training digits are centered.
- Padding: Circular padding appears to prevent learned non-stationarities as the receptive field becomes larger on ZSMM.The authors hypothesize that tiled values are harder to distinguish from actual values than zero padding.
- Padding: The effect of padding was not tested on other datasets, and circular padding could introduce other issues.This limits the scope of the padding conclusion beyond ZSMM.