PaperlayerView exact revision on arXiv

Source-linked AI summary

Attentive Neural Processes

Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, Yee Whye Teh

arXiv:1901.05761v2cs.LGstat.ML

TL;DR

Neural Processes underfit observed context points, producing inaccurate predictions. The paper adds differentiable attention to focus on relevant contexts, improving prediction accuracy and training speed while expanding modelled function ranges.

  • Problem

    Neural Processes can underfit their context sets, yielding inaccurate predictive means and overestimated variances at observed inputs.

  • Method

    Attentive Neural Processes incorporate differentiable attention that learns to attend to context points relevant to each target while preserving permutation invariance.

  • Results

    ANPs greatly improve prediction accuracy, train faster than NPs, and expand the range of functions that can be modelled.

  • Takeaways & Limitations

    Attention helps ANPs reconstruct context points accurately, predict targets, preserve increased uncertainty away from contexts, and model a wider range of functions.

  • Takeaways & Limitations

    Attention raises computational complexity from O(n + m) to O(n(n + m)) and makes prediction slower, despite comparable or faster training in practice.

Abstract

from arXiv · show

Neural Processes (NPs) (Garnelo et al 2018a;b) approach regression by learning to map a context set of observed input-output pairs to a distribution over regression functions. Each function models the distribution of the output given an input, conditioned on the context. NPs have the benefit of fitting observed data efficiently with linear complexity in the number of context input-output pairs, and can learn a wide family of conditional distributions; they learn predictive distributions conditioned on context sets of arbitrary size. Nonetheless, we show that NPs suffer a fundamental drawback of underfitting, giving inaccurate predictions at the inputs of the observed data they condition on. We address this issue by incorporating attention into NPs, allowing each input location to attend to the relevant context points for the prediction. We show that this greatly improves the accuracy of predictions, results in noticeably faster training, and expands the range of functions that can be modelled.

1 INTRODUCTION

Neural Processes efficiently model distributions over regression functions conditioned on arbitrary context sets, but their fixed mean-aggregated context representation can cause underfitting at observed inputs. Attentive Neural Processes address this by learning which context points matter for each target prediction.

  • Neural Processes: Neural Processes map context input-output pairs to conditional distributions over regression functions with prediction complexity linear in context-set size.They can condition on arbitrary numbers of contexts and predict arbitrary target outputs.
  • Problem: NPs underfit context points, producing inaccurate predictive means and overestimated variances at observed input locations.In image regression, reconstructions can remain globally coherent while failing to reproduce the observed top-half pixels accurately.
  • Problem: Mean aggregation gives every context representation equal weight, making it difficult for the decoder to identify context points relevant to a specific target.Increasing representation dimensionality is theoretically possible but practically insufficient in the reported experiments.
  • Approach: ANPs use differentiable attention so each target input can attend to relevant context points while preserving permutation invariance.The approach is motivated by Gaussian-process kernels, which indicate context-point relevance through input similarity.
  • Reported outcomes: ANPs improve context reconstruction and training speed relative to NPs and model a wider range of functions.The paper evaluates them on 1D function regression and 2D image regression.

2 BACKGROUND

The NP represents permutation-invariant context information and predicts target distributions through deterministic and optional latent paths. Attention provides permutation-invariant query-dependent aggregation mechanisms, including dot-product and multihead variants.

  • Neural Processes: An NP models functions from xi ∈ R^dx to yi ∈ R^dy and supports arbitrary context and target set sizes with ordering-invariant predictions.In practice, the context set is a subset of the target set.
  • Deterministic path: The deterministic encoder passes each context pair through an MLP and mean-aggregates the representations into a finite-dimensional permutation-invariant summary rC.The decoder uses xi and rC to parameterize a factorized Gaussian likelihood across targets.
  • Latent path: The latent NP path adds a global latent z, modeled by a factorized Gaussian parameterized by a context summary sC, to represent uncertainty across process realizations.The paper investigates using both deterministic and latent paths because this gives the most expressive model and supports attention.
  • Training: Training reconstructs targets with a likelihood term and regularizes the latent path through a KL term between target- and context-conditioned distributions.The encoder and decoder parameters are learned by maximizing the ELBO.
  • Attention: Attention computes query-specific weighted aggregates of key-value pairs while remaining invariant to the ordering of those pairs.Dot-product attention uses query-key similarity, and multihead attention applies learned linear transformations before combining multiple attention heads.
  • Attention: Dot-product attention can be computed with two matrix multiplications and a softmax, while multihead attention concatenates and linearly transforms head-specific values.These mechanisms provide parameterized alternatives to locality-based attention such as the Laplace kernel.

3 ATTENTIVE NEURAL PROCESSES

Attentive Neural Processes replace NP aggregation with attention mechanisms that model context interactions and produce target-specific context representations. This improves expressivity and accuracy while increasing computational complexity, with parallelizable attention keeping training practically competitive.

  • Architecture: ANPs apply self-attention to context points before aggregation and cross-attention from each target input to context representations.Self-attention models interactions among context points in both deterministic and latent paths.
  • Architecture: Cross-attention replaces deterministic mean aggregation with a query-specific representation r* that weights context points according to their relevance to each target.The decoder remains unchanged except that it receives r* instead of the shared representation rC.
  • Relationship to NP: Uniform attention recovers the original NP, while ANP retains permutation invariance in the contexts and uses the same training loss.The likelihood remains Gaussian, with a diagonal Gaussian latent posterior.
  • Cost: Attention raises computational complexity from O(n + m) to O(n(n + m)) for n contexts and m targets.Self-attention operates across contexts and target queries compute weights over all contexts.
  • Cost: Despite higher complexity, ANP training remains comparable in practice and is reported to learn significantly faster than NP in iterations and wall-clock time.The trade-off is slower prediction time.

4 EXPERIMENTAL RESULTS

Experiments on synthetic GP and image regression show that attention substantially improves NP context reconstruction and predictive quality, while preserving flexible function modeling and enabling resolution changes. Multihead and stacked attention also support faster learning and richer image reconstructions, though the paper presents these applications as demonstrations rather than state-of-the-art replacements.

  • 1D function regression: ANPs learn GP context reconstructions more rapidly and reach lower reconstruction error than NPs, including when measured against wall-clock time.Multihead ANP can outperform an NP with a larger bottleneck while using 10% of the wall-clock time.
  • 1D function regression: Increasing the NP bottleneck improves reconstruction only up to a limit, whereas multihead attention achieves better results without simply enlarging the representation.Beyond a certain bottleneck size, NP learning becomes too slow and convergence remains worse than multihead ANP.
  • 1D function regression: Multihead attention reconstructs context points accurately, while multiple heads smooth interpolations and preserve greater uncertainty away from contexts.Dot-product attention also reconstructs nearly all contexts accurately but produces less smooth predictions than multihead attention.
  • 2D image regression: Stacked Multihead ANP produces accurate full-image reconstructions from random pixels or the top half, and different latent samples yield diverse global image structures.The model generalises from training with at most 200 context points to prediction with 512 context points in the top-half reconstruction task.
  • 2D image regression: On CelebA, Multihead and Stacked Multihead ANP improve context reconstruction error over NP, while stacked self-attention adds qualitative gains in crispness and global coherence.Target NLL improves with multihead cross-attention, with smaller additional gains from stacked self-attention.
  • Resolution mapping: A Stacked Multihead ANP trained on 32 × 32 images maps lower resolutions to realistic outputs and generates sharper 256 × 256 images than linear and cubic interpolation baselines.The model has not seen images beyond the original resolution and sometimes reconstructs facial features that the coarse input does not separate.

5 RELATED WORK

The related work situates ANPs alongside Gaussian Processes, meta-learning, conditional latent-variable models, and spatial prediction systems, emphasizing parallels in attention and regression-function modeling.

  • Gaussian Processes: ANPs relate to Gaussian Processes because both attention and GP kernels measure similarity between points in the input domain.
  • Gaussian Processes: Unlike GPs, ANPs learn predictive uncertainties directly from data and use a different training regime, making direct comparisons difficult.
  • Meta-Learning: ANPs can be viewed as few-shot learning models that condition predictions for a new function on input-output pairs observed at test time.
  • Meta-Learning: Related regression work includes Vfunc, while GP research also addresses multi-task learning using several related approaches.
  • Generative Query Networks: Generative Query Networks correspond to a special case of NPs in which inputs are viewpoints and outputs are scene frames.

6 CONCLUSION AND DISCUSSION

The paper concludes that ANPs augment NPs with attention to address underfitting, improving prediction accuracy and training speed while broadening the functions they can model.

  • ANPs augment NPs with attention to resolve the fundamental problem of underfitting.
  • The authors report improved context and target NLL, faster training, and a wider range of modelled functions.
  • Future work includes cross-attention in the latent path, dependencies among local latents, stochastic text completion, and closer connections to Image Transformers.

A ARCHITECTURAL DETAILS FOR (A)NP

The appendix specifies architectural choices for NP and ANP experiments in one- and two-dimensional regression, including attention variants and self-attention depth.

  • Attention architecture: The 1D experiments use basic Multihead1 cross-attention, while the 2D experiments use Image Transformer-style multihead cross-attention without dropout.
  • Attention architecture: The 2D Stacked Multihead ANP uses two self-attention layers in both deterministic and latent paths.
  • Attention architecture: Adding more than two self-attention layers produced no noticeable qualitative or quantitative gains.

B EXPERIMENTAL DETAILS OF 1D FUNCTION REGRESSION EXPERIMENT

The 1D regression experiments use GP-generated curves with fixed or randomly varying kernel hyperparameters, and compare model architectures under specified implementation settings.

  • Data generation: The experiments use a squared-exponential GP with length scale l = 0.6 and fixed kernel scale σ2.
  • Data generation: For random hyperparameters, l is sampled from U[0.1, 0.6] and σf from U[0.1, 1], with likelihood noise σn = 0.02.
  • Training setup: Training uses batch size 16, with either 16 curves from one GP or 16 curves generated from independently sampled hyperparameters.
  • Model settings: NP latent dimensionality d varies over {128, 256, 512, 1024}, whereas ANP always uses d = 128.
  • Model settings: Figure 8 presents the NP and ANP architectures used for both 1D and 2D regression.

C ADDITIONAL FIGURES FOR 1D REGRESSION ON GP DATA

Additional 1D GP experiments compare attention-based NPs with fixed-kernel settings, KL behaviour, and Bayesian optimization performance. Multihead ANP approaches the oracle GP in predictions and regret, while dot-product attention can produce non-smooth interpolations.

  • Oracle GP comparison: Multihead ANP predictions are notably closer to the oracle GP than NP predictions, although predictive variance remains underestimated.The authors suggest variational inference as one possible explanation for the underestimated variance.
  • Fixed hyperparameters: Fixed kernel hyperparameters still produce highly non-smooth conditional distributions for dot-product attention.The same behaviour appears in the random-kernel-hyperparameter case.
  • Attention behaviour: Dot-product attention can collapse to a nearest-neighbour predictor, yielding good reconstructions but poor interpolation between context points.This occurs when one softmax entry becomes saturated.
  • KL behaviour: In fixed-hyperparameter GP data, the multihead ANP KL term quickly reaches 0, indicating that the deterministic path is sufficient for accurate predictions.With random hyperparameters, added data variation produces a non-zero KL and greater use of latent variables.
  • Bayesian optimization: In Bayesian optimization, multihead attention gives the smallest simple regret consistently and approaches the oracle GP.Its cumulative-regret slope also decreases most rapidly among the NPs, indicating effective use of previous evaluations.

D EXPERIMENTAL DETAILS OF 2D IMAGE REGRESSION EXPERIMENT

The 2D image-regression experiments train the same architecture on MNIST and CelebA using randomly sampled pixel contexts and targets. Training uses Adam, separate learning rates, and minimal architectural tuning.

  • Data sampling: Training samples randomly selected image pixels as targets and subsets of those pixels as contexts, with n drawn from 3 to 200.The number of targets is sampled as m ∼ n + U[0, 200 − n].
  • Optimisation: MNIST and CelebA use batch size 16, with learning rates of 5e-5 and 5e-4 respectively.The Adam optimiser is used with TensorFlow defaults for other hyperparameters.
  • Architecture: The stacked self-attention architecture follows the Image Transformer design but omits dropout and positional embeddings.The same architecture is used for both datasets, with little tuning of architectural hyperparameters.

E ADDITIONAL FIGURES FOR 2D IMAGE REGRESSION ON MNIST AND CELEBA

Additional image-regression figures show that attention reduces uncertainty and improves reconstruction quality. Stacked multihead attention produces sharper, more globally coherent images, while multihead heads contribute when targets are outside the context.

  • Predictive uncertainty: Attention substantially reduces predictive uncertainty around image-reconstruction edges as the number of contexts increases.For the attention-based NP, uncertainty nearly disappears with the full context, unlike the original NP.
  • Additional reconstructions: Figures 14 and 15 provide additional MNIST and CelebA reconstructions of full images from top-half contexts.The figures extend the image-regression comparisons across more examples.
  • Reconstruction quality: Stacked Multihead ANP significantly improves over Multihead ANP, producing sharper images with better global coherence.The improvement also holds when faces are not axis-aligned.
  • Attention heads: When the target pixel is disjoint from the context, all attention heads become useful for prediction.This contrasts with the case where the target is included in the context, where one head can suffice.
  • Attention visualisation: Figure 16 visualizes the pixels attended by each multihead-attention head for a target pixel and 100 random context pixels.Each head uses a distinct colour and the target is marked by a cross.
Loading 1901.05761v2…