PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformer Neural Processes: Uncertainty-Aware Meta Learning Via Sequence Modeling

Tung Nguyen, Aditya Grover

arXiv:2207.04179v2cs.LGcs.AI

TL;DR

Neural Processes aim to provide accurate, uncertainty-aware predictions for new tasks but can underfit and require intractable likelihood approximations. Transformer Neural Processes address this gap through autoregressive transformer sequence modeling with structure-aware invariance and equivariance, outperforming prior NP variants across several benchmarks. Future work targets higher-dimensional scaling and cleaner separation of functional and point uncertainties.

  • Problem

    Neural Processes can suffer from underfitting and intractable likelihoods, limiting uncertainty-aware meta learning for sequential decision making.

  • Method

    TNPs formulate uncertainty-aware meta learning as autoregressive sequence modeling with a transformer architecture designed for context invariance and target equivariance.

  • Results

    TNPs outperform attention- and non-attention-based Neural Process variants by a large margin across meta regression, image completion, contextual bandits, and Bayesian optimization.

  • Takeaways & Limitations

    Tractable and equivariant TNP variants retain the empirical benefits of the framework to a good degree, rather than requiring only the most expressive decoding distribution.

  • Takeaways & Limitations

    The authors identify scaling to high-dimensional problems and cleanly separating functional from point uncertainties as future work.

Abstract

from arXiv · show

Neural Processes (NPs) are a popular class of approaches for meta-learning. Similar to Gaussian Processes (GPs), NPs define distributions over functions and can estimate uncertainty in their predictions. However, unlike GPs, NPs and their variants suffer from underfitting and often have intractable likelihoods, which limit their applications in sequential decision making. We propose Transformer Neural Processes (TNPs), a new member of the NP family that casts uncertainty-aware meta learning as a sequence modeling problem. We learn TNPs via an autoregressive likelihood-based objective and instantiate it with a novel transformer-based architecture. The model architecture respects the inductive biases inherent to the problem structure, such as invariance to the observed data points and equivariance to the unobserved points. We further investigate knobs within the TNP framework that tradeoff expressivity of the decoding distribution with extra computation. Empirically, we show that TNPs achieve state-of-the-art performance on various benchmark problems, outperforming all previous NP variants on meta regression, image completion, contextual multi-armed bandits, and Bayesian optimization.

1. Introduction

Uncertainty-aware meta learning seeks models that adapt to unseen tasks from few labeled examples while providing accurate, uncertainty-quantified predictions for sequential decision making. TNPs address limitations of Neural Processes with autoregressive sequence modeling and a transformer architecture, achieving strong performance across benchmark tasks.

  • Meta learning aims to adapt models quickly to unseen tasks using only a few labeled examples.
  • Uncertainty-aware predictions are especially important in Bayesian optimization and multi-armed bandits, where uncertainty can guide data acquisition.
  • Neural Processes model distributions over functions and enable rapid adaptation by conditioning predictions for unlabeled points on labeled context points.
  • NP variants often have intractable marginal likelihoods, requiring surrogate variational lower bounds instead of direct likelihood optimization.
  • TNPs cast uncertainty-aware meta learning as sequence modeling and train by autoregressively maximizing target-point conditional log-likelihoods.
  • TNPs use a transformer architecture and variants that trade off decoding expressivity, computational tractability, and exact equivariance, outperforming prior NP variants across four benchmark areas.

2. Background

Meta learning evaluates predictions on unseen functions after observing a small labeled training set, while uncertainty-aware models output joint predictive distributions over test points. Neural Processes implement this through latent-variable modeling and variational training, whereas transformers provide the sequential modeling machinery studied here.

  • Meta learning assumes an unknown function distribution and tests generalization on unseen functions after providing a small labeled training set.
  • Uncertainty-aware meta learning outputs a joint predictive distribution over the test set, with predictive standard deviations quantifying uncertainty.
  • Neural Processes are stochastic processes that predict unlabelled target points from labelled context points while using latent variable z to represent functional uncertainty.
  • Because the NP likelihood is intractable, training maximizes an evidence lower bound of the log-likelihood.
  • NPs can be viewed as VAEs over target labels, using a permutation-invariant context encoder and a decoder that predicts target labels independently given z.
  • Transformers use self-attention to model relationships among input tokens, with positional encoders incorporating sequential information.

3. Transformer Neural Processes

TNPs cast uncertainty-aware meta-learning as autoregressive sequence modeling while enforcing context invariance and target equivariance. The framework includes autoregressive, diagonal, and non-diagonal decoders that trade predictive expressivity against computational tractability.

  • Sequence-modeling objective: TNPs treat each function’s evaluations as an ordered sequence, use randomly selected context pairs for conditioning, and autoregressively model the remaining target likelihood.The training objective is optimized over batches of functions with randomly sampled context sizes.
  • Architectural desiderata: Vanilla autoregressive transformers violate context invariance and target equivariance because positional encodings make predictions depend on point ordering.The architecture must preserve pair associations without allowing arbitrary permutations to change the predictive distribution.
  • Autoregressive TNP-A: TNP-A uses concatenated input-output tokens, auxiliary padded target tokens, and a custom attention mask to preserve autoregressive prediction order.Contexts attend only to themselves, while targets and padded targets attend to contexts and the appropriate previous targets.
  • Autoregressive TNP-A: TNP-A approximates target equivariance by averaging predictions over randomly sampled target permutations, making exact computation tractable during training and evaluation.The permutation group is too large to enumerate directly, so Monte Carlo symmetrization is used.
  • Diagonal TNP-D: TNP-D assumes conditional independence among target points, yielding a diagonal covariance and exact target equivariance, but this assumption can limit accurate function modeling.TNP-D removes target pairs from the input and predicts targets from contexts and padded target inputs.
  • Non-diagonal TNP-ND: TNP-ND uses a non-diagonal covariance to balance TNP-D’s tractability with TNP-A’s expressivity while satisfying context invariance and target equivariance.The main experiments use a Cholesky parameterization; its structured representation supports arbitrary target counts with O(n) space instead of O(n^2).

4. Experiments

Experiments evaluate TNPs against NP baselines across regression, image completion, contextual bandits, and Bayesian optimization. TNPs generally improve predictive likelihood, image completion, and decision-making performance, while decoder variants expose expressivity–tractability trade-offs.

  • 4. Experiments: The experiments compare TNPs with CNPs, NPs, BNPs, and attentive variants across regression, image completion, Bayesian optimization, and contextual bandits.These tasks have been used extensively to benchmark NP-based models.
  • 4.1. 1-D Regression: TNPs outperform other methods on 2/3 regression kernels by a large margin, although they underperform on the periodic kernel.Among TNP variants, TNP-A achieves the best likelihood, followed by TNP-ND and TNP-D.
  • 4.2. Image completion: TNPs significantly improve image-completion likelihood over baselines on EMNIST and CelebA.TNP-A achieves the best likelihood, followed by TNP-ND and TNP-D.
  • 4.2. Image completion: TNPs produce noticeably better completed images than the best baseline, with only slightly worse performance on EMNIST unseen classes than seen classes.The figure uses 100 context points and randomly drawn EMNIST unseen classes.
  • 4.3. Contextual bandits: TNPs outperform all baselines by a large margin across contextual-bandit settings, especially at higher δ values.Their performance drops only slightly as problem difficulty increases, whereas baseline methods barely work on hard problems.
  • 4.4. Bayesian Optimization: TNPs significantly outperform baselines on all three 1D Bayesian-optimization kernels, including the periodic kernel.In multi-dimensional tasks, TNP-ND beats baselines in 5/6 tasks, while TNP-D does so in 4/6 tasks.

5. Related Work

Related work situates TNPs within neural processes and transformer-based modeling. Prior work established NP variants for function distributions and uncertainty, while transformers model interactions across sequence elements and data points.

  • Neural Processes: CNPs use deterministic representations and therefore lack functional uncertainty, while NPs introduce latent variables to address this limitation.The related-work discussion places TNPs within the progression from CNPs to stochastic NPs.
  • Transformers: Transformers were introduced for flexible language modeling and later achieved breakthroughs in language and vision.This motivates studying transformer architectures for uncertainty-aware meta-learning.
  • Transformers: Set Transformers apply transformer architectures to set-input problems by modeling interactions between data points.This work is described as closely related to TNPs.

6. Discussion

The discussion frames TNPs as a sequence-modeling alternative for uncertainty-aware meta-learning and attributes gains to replacing the full architecture stack, not only the decoder. It also identifies scaling to higher-dimensional problems and separating functional from point uncertainty as future directions.

  • TNPs reframe uncertainty-aware meta-learning through sequence modeling, using an autoregressive objective and transformer backbone.
  • Replacing the entire architecture stack improves performance across benchmark tasks beyond the benefits of a more expressive decoding distribution alone.
  • Tractable and equivariant TNP-D and TNP-ND retain much of the empirical benefit of the more expressive TNP-A parameterization.
  • Future work targets scaling TNPs to high-dimensional problems beyond current benchmarks.
  • Future work also considers separating functional and point uncertainties, potentially using stochastic transformers.

A. Implementation details

The implementation details describe TNP decoder computation, experiment configurations, and a training procedure that addresses the contextual-bandit mismatch between training and evaluation.

  • TNP-NP computes the lower-triangular factor of the multivariate normal output distribution from the final transformer-decoder representations.The lower-triangular operator removes the upper-triangular parts, and n = N − m is the number of target points.
  • Contextual-bandit training randomly drops some context-arm rewards and regresses their ground truth to close the train–test mismatch.The mismatch arises because training observes rewards for all arms, whereas testing observes only the selected arm.
  • The reported configurations use four embedding layers, feed-forward dimension 128, four attention heads, and six transformer layers in one setup.
  • The TNP-ND configuration uses two self-attention layers on target representations, projection dimension 20, and four projection layers.
  • Training schedules vary across experiments, including 100000 steps or 200 epochs, with a 0.5 reward-drop rate reported for contextual bandits.

B. Additional results

Additional experiments report strong TNP performance across regression, image completion, and sequential decision-making settings, alongside qualitative image and sampling analyses.

  • B.1. 1-D regression with additional metrics: Three TNP variants outperform baselines on all reported regression metrics for 2/3 kernels, while TNP-A performs best on log-likelihood.The variants have similar accuracy and calibration, whereas attention often improves accuracy but worsens calibration among other NP methods.
  • B.2. Image completion: Three TNP variants outperform baselines on all 3 evaluation metrics for both CelebA and EMNIST, with TNP-A best on log-likelihood.The variants are similar in accuracy and sharpness; calibration error was omitted because it was too time-consuming.
  • B.2. Image completion: On full-image completion, TNP-D achieves better performance than the strongest baseline BANP when the target is the entire image.
  • B.2. Image completion: TNPs outperform Conv(C)NP by a large margin on EMNIST image completion, while Conv(C)NP requires one-dimensional off-grid inputs after discretization.
  • B.2. Image completion: Qualitative results show more accurate images with fewer artifacts, including interpretable letters on unseen EMNIST classes produced only by TNPs.
  • C. TNPs vs other NP variants on sampling multiple functions: TNP-A samples target outputs sequentially from predictive distributions conditioned on context points and previously sampled targets.

C.2. Sampling results

Sampling experiments compare TNP-A with other NP variants on regression and image completion, showing stronger sample diversity with sparse context and consistency with denser context.

  • 1-D regression: With 10 context points, TNP-A produces more diverse regression samples, whereas other NP methods represent uncertainty mainly through larger predictive-variance regions.With 30 context points, all methods produce consistent functions.
  • Image completion: With 20 image context points, TNP-A produces three samples representing different digits, while diversity from other NP methods is unclear.With 50 context points, the methods produce consistent samples.

D. Architectural ablation analysis

The ablations compare alternative input representations, autoregressive architectures, covariance parameterizations, and training regimes for TNPs. Coupling context inputs and labels is crucial, while meta-trained TNPs outperform the autoregressively pretrained variant, especially on sequential decision-making tasks.

  • Input representations: TNP-D is slightly better than separate embeddings, whereas cross-attention performs poorly because it cannot jointly inspect context pairs (x_i, y_i).The comparison attributes the failure to the model’s inability to infer the underlying function distribution from uncoupled inputs and labels.
  • Input representations: Coupling each context input x_i with its label y_i is crucial for TNP performance.Without access to context pairs, the model makes a random guess with very large variance.
  • Autoregressive architectures: TNP-A, TNP-A-1, and TNP-A-2 are compared as alternative autoregressive architectures using different attention and masking arrangements.TNP-A-1 separates context and target encodings, while TNP-A-2 feeds both into one masked encoder that blocks context-to-target attention.
  • Covariance parameterization: Low-rank covariance parameterization adds an MLP for diagonal terms and guarantees a positive definite predictive covariance matrix.The ablation compares this version with Cholesky decomposition on one-dimensional regression across Gaussian-process kernels and evaluation metrics.
  • Training regime: TNP-A-Pretrained performs reasonably well across tasks but still lags behind meta-trained TNPs, especially on contextual bandits and Bayesian optimization.The pretrained model uses an autoregressive objective and proper masking at evaluation, whereas the main TNPs are trained in the meta-learning regime.

E.4. Bayesian Optimization

The Bayesian optimization evaluation compares TNP-A-Pretrained with three TNP variants on one-, two-, and three-dimensional tasks. TNP-A-Pretrained performs reasonably well overall but remains behind meta-trained TNPs, particularly in sequential decision making.

  • 1D Bayesian optimization: TNP-A-Pretrained and three TNP variants are compared using regret on one-dimensional Bayesian optimization tasks across kernels.Each kernel is evaluated over 100 runs with different seeds, reporting the mean and standard deviation.
  • 2D and 3D Bayesian optimization: TNP-A-Pretrained and three TNP variants are also evaluated on two- and three-dimensional Bayesian optimization tasks.Each function is run 100 times with different seeds, with mean and standard deviation reported.
  • Overall comparison: TNP-A-Pretrained performs reasonably well overall but still lags behind the meta-trained TNPs, especially in Bayesian optimization.The authors identify improving pretraining to match meta-trained performance as future work.
Loading 2207.04179v2…