Dimension Agnostic Neural Processes

TL;DR

This work introduces Dimension Agnostic Neural Processes (DANP), a general-purpose regressor capable of handling tasks with varying input and output dimensions while providing predictive uncertainty. It leverages a Dimension Aggregator Block (DAB) to map heterogeneous features into a fixed-dimensional space and augments the architecture with a Transformer-based latent path to capture broader, task-generalizable features. Through extensive experiments on Gaussian Process regression, image and video completion, and Bayesian optimization, DANP consistently outperforms prior Neural Process variants in target likelihood and calibration, while demonstrating strong zero-shot and few-shot generalization to unseen dimensionalities. The approach offers a practical step toward a single foundation regressor that can adapt across diverse data structures and real-world scenarios, with reproducibility and ethical considerations addressed. The mathematical formulation rests on a shared predictive density $p(\mathbf{Y}|\mathbf{X},{\mathcal{D}}_{c})$ that integrates over latent variables, enabling principled uncertainty estimation across tasks.

Abstract

Meta-learning aims to train models that can generalize to new tasks with limited labeled data by extracting shared features across diverse task datasets. Additionally, it accounts for prediction uncertainty during both training and evaluation, a concept known as uncertainty-aware meta-learning. Neural Process(NP) is a well-known uncertainty-aware meta-learning method that constructs implicit stochastic processes using parametric neural networks, enabling rapid adaptation to new tasks. However, existing NP methods face challenges in accommodating diverse input dimensions and learned features, limiting their broad applicability across regression tasks. To address these limitations and advance the utility of NP models as general regressors, we introduce Dimension Agnostic Neural Processes(DANP). DANP incorporates Dimension Aggregator Block(DAB) to transform input features into a fixed-dimensional space, enhancing the model's ability to handle diverse datasets. Furthermore, leveraging the Transformer architecture and latent encoding layers, DANP learns a wider range of features that are generalizable across various tasks. Through comprehensive experimentation on various synthetic and practical regression tasks, we empirically show that DANP outperforms previous NP variations, showcasing its effectiveness in overcoming the limitations of traditional NP models and its potential for broader applicability in diverse regression scenarios.

Figures (11)

• Figure 1: Model comparison between and . While nguyen2022transformer solely employs a deterministic pathway with Masked Transformer layers, incorporates both and an extra latent pathway alongside Transformer layers and a Self-Attention layer.
• Figure 2: The overview of module. A can encode and decode inputs and outputs of varying dimensions.
• Figure 3: Posterior samples of in (Left) the Zero-shot scenario with a 1-dimensional dataset and (Right) the Image completion task using the EMNIST and CelebA datasets. (Left) Black stars represent the context points and the dashed line indicates the ground truth for the target points. Each color represents different posterior samples generated from the latent path. (Right) Displays the full image, context points, predictive mean, and standard deviation of for both the EMNIST and CelebA datasets. Outputs for both images are produced by a single model.
• Figure 4: Results for on 6-, 9-, 10-, and 16-dimensional hyperparameter tuning tasks in the HPO-B benchmark. Note that is pre-trained on 2-, 3-, 8-dimensional tasks.
• Figure 5: Results for with various tasks. These four figures, from left to right, show the regret results for 1-dimensional with RBF kernel, 2-dimensional cosine, 3-dimensional Ackley, and the experiments.