PaperlayerView exact revision on arXiv

Source-linked AI summary

Self-Attention Between Datapoints: Going Beyond Individual Input-Output Pairs in Deep Learning

Jannik Kossen, Neil Band, Clare Lyle, Aidan N. Gomez, Tom Rainforth, Yarin Gal

arXiv:2106.02584v2cs.LGstat.ML

TL;DR

The paper challenges models that predict from parameters and one input by introducing NPTs, which take entire datasets as input and learn cross-datapoint interactions with self-attention. NPTs achieve competitive benchmark results and complex reasoning capabilities, while their quadratic attention cost limits scalability.

  • Problem

    Most supervised deep learning predicts from model parameters and a single input, leaving direct dependencies between datapoints outside the prediction mechanism.

  • Method

    NPTs take the entire dataset and masking matrix as input, alternating attention between datapoints and attributes to learn end-to-end how other datapoints inform predictions.

  • Results

    NPTs achieve highly competitive benchmark performance, solve complex reasoning tasks over datapoints, and learn to rely on datapoint attention for prediction.

  • Takeaways & Limitations

    NPTs provide a general treatment of dataset-level dependencies in deep-learning prediction, extending parametric modeling with learned direct attention to other datapoints.

  • Takeaways & Limitations

    NPTs inherit quadratic self-attention scaling limitations, reaching 24 GB of GPU memory at about 8000 datapoints and requiring approximations for larger datasets.

Abstract

from arXiv · show

We challenge a common assumption underlying most supervised deep learning: that a model makes a prediction depending only on its parameters and the features of a single input. To this end, we introduce a general-purpose deep learning architecture that takes as input the entire dataset instead of processing one datapoint at a time. Our approach uses self-attention to reason about relationships between datapoints explicitly, which can be seen as realizing non-parametric models using parametric attention mechanisms. However, unlike conventional non-parametric models, we let the model learn end-to-end from the data how to make use of other datapoints for prediction. Empirically, our models solve cross-datapoint lookup and complex reasoning tasks unsolvable by traditional deep learning models. We show highly competitive results on tabular data, early results on CIFAR-10, and give insight into how the model makes use of the interactions between points.

1 Introduction

NPTs challenge single-input parametric prediction by using the training dataset directly and learning interactions between datapoints with self-attention. They combine this mechanism with end-to-end learning, achieving competitive benchmark performance and complex cross-datapoint reasoning.

  • Motivation: Parametric models predict from learned parameters and one input, whereas NPTs can condition predictions on the training dataset directly.NPTs retain parametric mechanisms while adding learned dependence on other datapoints.
  • Approach: NPTs take the entire dataset as input and explicitly learn relationships between datapoints using multi-head self-attention.The model learns how to balance parametric prediction with direct use of other datapoints.
  • Relation to non-parametric models: Conventional non-parametric models use explicit training-data dependence but typically rely on fixed interaction procedures that cannot learn interactions from data.NPTs differ by learning these interactions end-to-end.
  • Empirical scope: NPTs achieve highly competitive performance on tabular datasets and early image-classification results.The introduction reports these outcomes without giving a quantitative value in the supplied passages.
  • Empirical scope: NPTs can solve complex reasoning tasks by combining representation learning with cross-datapoint lookup.The paper contrasts this capability with conventional deep learning and conventional non-parametric models.

2 Non-Parametric Transformers

NPTs process datasets and masks jointly, alternating attention between datapoints and attributes to predict masked values. This design supports learned cross-datapoint relationships but faces quadratic self-attention costs on large datasets.

  • Inputs and masking: NPTs receive the entire dataset and a masking matrix, predicting masked entries from observed values.The framework can represent classification, regression, imputation, self-supervision, and semi-supervision through mask placement.
  • Architecture: The architecture embeds datapoints independently, then alternates Attention Between Datapoints and Attention Between Attributes before producing masked-entry predictions.Attention Between Datapoints operates across samples, while Attention Between Attributes operates within each datapoint.
  • Properties: NPTs are permutation-equivariant: shuffling input datapoints produces correspondingly shuffled predictions.This encodes that learned relationships should not depend on datapoint ordering.
  • Attention Between Datapoints: Attention Between Datapoints transforms representations using pairwise relationships among all datapoints and can build higher-order interactions across layers.Unlike usual Transformer attention over features or tokens within one datapoint, ABD attends across datapoints.
  • Scalability: Self-attention has O(n^2) time and space complexity, so large datasets require approximations such as random minibatching.Standard NPT model sizes reach 24 GB of GPU memory at about 8000 datapoints; minibatches include training and test data at prediction time.

3 Related Work

NPTs relate to deep non-parametric models, Neural Processes, attention architectures, meta-learning, transductive learning, and metric learning, but differ in dataset setup and the scope of learned interactions.

  • Deep non-parametric models: Deep Gaussian Processes and Deep Kernel Learning combine neural representation learning with non-parametric ideas, but use different inference or interaction mechanisms than NPTs.Deep GPs require complex approximate inference, while DKL applies a neural network independently to each datapoint before its non-parametric component.
  • Neural Processes: Neural Processes use latent-variable models and architectural constraints for approximate consistency, whereas NPTs directly learn relationships between datapoints in their dataset input.The supplied passage emphasizes differing training regimes between Gaussian Processes and Neural Processes.
  • Attention architectures: Set Transformer attends to input sets but relies on multiple independent sets for training and makes one prediction per set, unlike NPTs’ broader setting.Other attention models cited process individual images, protein sequences, or tabular features without general cross-datapoint reasoning.
  • Few-shot learning and meta-learning: NPT evaluation tasks share motivations with meta-learning because relational structure learned from training data supports generalization to novel test inputs.The supplied passage frames this as a shared motivation rather than an identity between the approaches.
  • Semi-supervised learning: NPTs natively support semi-supervised and transductive settings by including unlabeled datapoints with masked targets in the input matrix.This connects dataset-level attention to methods that use unlabeled inputs during training.
  • Metric learning: NPT attention can implicitly learn similarity and dissimilarity relations like metric-learning and retrieval models, while using them for prediction.The passage characterizes attention between datapoints as implicitly learning such relations.

4 Experiments

Experiments evaluate NPTs on standard tabular and image benchmarks, controlled lookup tasks, and real-data interaction tests. NPTs perform competitively, learn cross-datapoint lookup and interventions, and rely on similar datapoints when useful.

  • Tabular Data Results: NPT achieves the highest average ranking on binary and multi-class classification, ties XGBoost on regression, and is best on 4 of 10 tabular datasets.The baselines include tree-based methods, MLPs, k-NN, and TabNet.
  • Image Data Results: 93.7% test accuracy on CIFAR-10 and 98.3% accuracy on MNIST demonstrate competitive image-classification results.The CIFAR-10 model uses a CNN followed by ABD layers, while MNIST uses linear patching.
  • Lookup and Intervention: 99.9% Pearson correlation with duplicate targets and RMSE 0.44 show that NPT learns the semi-synthetic cross-datapoint lookup.The model attends to duplicate datapoints and achieves an error about an order of magnitude below the original Protein task.
  • Lookup and Intervention: Near-perfect correlation after intervening on duplicate targets shows that predictions adjust to altered values without retraining.The interventions use target values highly unlikely under the training distribution.
  • Real-Data Interactions: Corrupting other rows significantly reduces performance on most datasets, while some datasets allow NPT to ignore other inputs and behave like a parametric model.On Protein, corruption makes NPT perform worse than all considered baselines; on Forest Cover, Kick, and Breast Cancer, performance is not significantly affected.
  • Real-Data Interactions: Attention maps and data deletion tests indicate that NPT relies more on datapoints similar in feature space when predicting on real data.Kept datapoints have significantly lower average feature-space distance than deleted datapoints, with Wilcoxon signed-rank p ≈ 8.77 · 10^-130.

5 Limitations, Future Work, and Conclusions

NPTs extend deep learning by taking entire datasets as input and learning relationships between datapoints through self-attention. The paper highlights competitive benchmark performance and cross-datapoint reasoning, while identifying scaling as a central limitation and several directions for future work.

  • Future Work: Future work proposes applying NPTs to continual learning, multi-task learning, few-shot generalization, and domain adaptation.The paper also suggests studying whether relations between datapoints remain useful under distribution shift.
  • Conclusions: NPTs use the entire dataset and self-attention to learn complex relationships between datapoints for prediction.The model learns end-to-end when and how to attend to other datapoints.
  • Conclusions: NPTs achieve highly competitive benchmark performance and solve complex reasoning tasks over datapoints.The conclusions also report that NPTs learn to rely on attention between datapoints on real data.

A Proof – NPT Is Equivariant over Datapoints

NPT is row-equivariant: permuting datapoints permutes the corresponding output rows without changing the computation’s substance. This follows from the equivariance of its component operations and their composition.

  • Definition: Row-equivariance means that permuting input datapoints produces the correspondingly permuted output rows.The definition applies a permutation to the input rows and matches each output at the permuted index.
  • Supporting lemmas: Row-wise operations, linear maps, and compositions preserve row-equivariance.The supporting lemmas establish these closure properties for functions applied across datapoint rows.
  • Self-attention: Self-attention is row-equivariant even though its intermediate query-key product also permutes attention columns.The final matrix multiplication restores the desired row-equivariance.
  • NPT components: Multihead self-attention, residual blocks, and the ABD and ABA modules are row-equivariant.These components inherit equivariance from self-attention, row-wise operations, and addition.
  • NPT: Because every NPT layer is row-equivariant and compositions preserve that property, NPT itself is row-equivariant.The proof concludes the property for the complete architecture.

B.1.1 Attention Maps for the Semi-Synthetic Experiments

NPT’s attention maps show that lookup behavior is distributed across heads and depths rather than required uniformly. Some heads attend from originals to duplicates, while others exhibit query-independent patterns or additional duplicate-to-duplicate attention.

  • Attention patterns: Some, but not all, attention heads display the diagonal lookup pattern needed for the semi-synthetic task.A single head with the appropriate pattern would suffice to implement lookup and solve the task perfectly.
  • Lookup behavior: Original datapoints attend to duplicate points, enabling NPT to look up their target values.This attention behavior minimizes the prediction loss for the original datapoints.
  • Attention patterns: Some shallow heads appear to learn almost query-independent behavior rather than performing a lookup.Heads 0, 1, and 2 at depth 0 are examples shown in the attention visualizations.
  • Attention patterns: Duplicate datapoints also attend to duplicate datapoints, although this interaction is irrelevant to loss because duplicates have no predicted targets.The resulting double-diagonal structure arises because the learned queries ignore the masked label column.

B.1.2 Modified Semi-Synthetic Experiments

Modified semi-synthetic tasks require relationships beyond exact lookup, and NPTs learn those relationships while fixed non-parametric baselines do not. NPT performance remains near-perfect across randomized features, target offsets, and their combination.

  • Task modifications: The modified tasks test interactions more complex than simple lookups, including randomized attributes and shifted duplicate targets.Random Features requires matching on a subset of features, while Add One changes only duplicate targets.
  • Interpretation: NPTs can learn interaction rules because they are not restricted to fixed distance-based non-parametric behavior.Unlike fixed kernel methods, NPTs can learn k-NN-style predictions or more complicated relations and can fall back to parametric prediction.
  • Results: NPTs maintain near-perfect performance across the original, Random Features, Add One, and combined modified tasks.Their losses are more than an order of magnitude lower than on the original data and correspond to Pearson correlation r > 99.9%.
  • Results: NPTs learn to match using non-randomized features and subtract 1 from duplicate targets when required.They can perform either modification separately or both simultaneously.
  • Baselines: k-NN and DKL solve the original duplication task under concessions but fail to accommodate the modifications.For k-NN, even k > 1 achieves little better than guessing on the modified tasks; DKL also fails on the presented modifications.
  • Baseline limitations: The DKL comparison is limited by training and implementation constraints, including the inability to condition the GP on new test data during each minibatch.The authors state that additional adaptations might make DKL solve the Random Features task, while Add One remains incompatible with its interpolation mechanism.

B.3.1 Attention Maps on Real Data

On the Protein regression dataset, attention maps reveal stronger attention between inputs that are nearby in feature space. The diagonal pattern appears in some, but not all, heads across model depths.

  • Visualization setup: Protein datapoints are sorted by feature-space distance to an arbitrary test datapoint so nearby attention-map entries represent similar inputs.This ordering gives the global structure of the maps an interpretable meaning.
  • Attention pattern: Some, but not all, attention heads display the similarity-based diagonal pattern.The figure covers all depths and a selection of heads for the Protein regression dataset.
  • Attention pattern: A diagonal attention pattern indicates that NPT attends more strongly to inputs similar in feature space.The visualizations plot the logarithm of the attention values.

B.3.2 Data Deletion Experiment

Data deletion tests indicate that NPT predictions depend on interactions with similar datapoints. Kept datapoints are closer to the active test point than deleted or randomly selected datapoints.

  • NPT iteratively deletes datapoints that do not significantly change a test prediction, then compares kept and deleted sets in input feature space.The experiment targets a single test sample at a time and measures average distances from the active datapoint.
  • p ≈0: Kept datapoints are significantly closer to the active point than deleted datapoints under a Wilcoxon signed-rank test.The raw Wilcoxon statistic is 3125889.5.
  • p ≈8.77 · 10^-130: Kept datapoints are also significantly closer than randomly selected datapoints after controlling deletion counts.The random subset deletes as many datapoints as the targeted procedure.
  • 79.5% of active test datapoints can be reduced to 1% of the original n, and 99.5% can be reduced to 10%.These results use n = 2048 datapoints, including 398 test datapoints, and the more computationally demanding deletion settings.
  • The results strongly suggest that NPT relies on interactions between similar datapoints for prediction.Figure B.3 visualizes average feature-space differences for kept, random, and deleted datapoints.

B.4 Ablation Study 1: NPT Hyperparameters

Ablations show that NPT remains effective across substantially different datasets and many hyperparameter changes, while masking and architectural choices can matter for particular tasks. The benchmark results also expose computational costs and evaluation caveats.

  • Setup: NPT uses nearly identical base models for Boston Housing (n = 506) and Protein (n = 45730), despite differing dataset sizes and feature types.Boston Housing includes categorical and continuous variables, whereas Protein is much larger.
  • Protein: RMSE = 3.40 ± 0.05: Base NPT ranks first on Protein, ahead of CatBoost at RMSE = 3.51.Large optimization changes, including removing or greatly increasing feature masking, reduce Protein performance.
  • Boston: None of the Boston ablations significantly changes base-model performance; NPT remains top-ranked unless semi-supervision or feature masking is substantially altered.The second-ranked Boston performer is MLP at RMSE = 3.32.
  • Robustness: Ablations support successful application with little tuning across datasets of different sizes and feature types, while depth and head-count changes appear nonsignificant.A reasonably low feature-masking probability, such as 15%, may be important for stable training.
  • Computational cost: NPT is computationally heavier: Protein training takes 11 hours versus less than 1 hour for boosting methods, and Higgs peak GPU memory reaches 19.18 GB.Prediction-time behavior was not investigated in detail, and reported costs depend on hardware and unoptimized implementations.
  • Evaluation caveats: The benchmark comparisons have split-related caveats because some datasets lack fixed test sets, so performance differences may reflect different train-test partitions.The authors also report a discrepancy between their TabNet Poker Hand result and the original report.

B.8 Image Classification Results

NPT achieves competitive image-classification results while showing measurable reliance on interactions between images. Results vary substantially by encoder and dataset, and the experiments are intended to demonstrate interaction learning rather than match pretrained Transformer performance.

  • Architectures: The image experiments use either a ResNet-18 encoder followed by ABD layers or a linear patching encoder followed by ABD and, in one setup, ABA layers.The CNN encoder setup is used for the main CIFAR-10 results, while linear patching is applied to MNIST and additional CIFAR-10 experiments.
  • Scope: The experiments target reasonable performance and evidence of datapoint interactions, not parity with pretrained Transformer image classifiers.The authors note that prior linear-patching work often uses large datasets or pretraining, unlike this setup.
  • CNN encoder: 93.7% test accuracy: ResNet-18 plus NPT on CIFAR-10 loses 1.2% under data corruption, while ResNet-18 alone reaches 93.9%.The corruption result suggests meaningful reliance on other datapoints for prediction.
  • CNN encoder: 94.7% test accuracy: ImageNet-pretrained ResNet-18 plus NPT exceeds the pretrained ResNet-18 alone at 94.2%.The combined model loses 0.7% in the data corruption experiment.
  • Linear patching encoder: 98.3% on MNIST and 68.2% on CIFAR-10: Linear patching NPT performance differs sharply across the two image datasets.Data corruption changes accuracy by -0.4% on MNIST and -5.1% on CIFAR-10.

C.1.2 NPT Training on Small Data

NPT’s small-data training uses a largely shared configuration with limited dataset-specific adjustments and evaluates a small set of variants. Target masking performs strongly across small datasets, while the architecture also supports missing-value prediction through its masking mechanism.

  • Configuration: The small-data NPT-Small configuration increases hidden dimensions to e = 128, uses a 50% flat learning-rate phase, and processes each full dataset without minibatching.Most architectural and optimization hyperparameters remain shared with NPT-Base.
  • Configuration: Eight variants are tested per dataset, including changes to layers, heads, target masking, feature masking, and learning-rate scheduling.The sweep contains eight distinct models rather than a full grid over modifications.
  • Results: A target-masking variant ranks first on each of the four small datasets.The authors attribute this potentially to label masking, regularization, or more stable convergence, while noting that target masking can increase training time.
  • Results: NPT is reported as robust to hyperparameter changes, with no obvious performance trends across many configurations.For medium and large datasets, only minor manual adaptations are made to account for dataset size and attribute count.
  • Missing values: NPT accommodates missing values end-to-end by marking missing entries with mask tokens rather than requiring imputation preprocessing.Stochastic feature masking trains the model to predict masked values while ignoring their input entries.

E Classification and Regression Benchmark Details

The benchmark suite spans varied UCI datasets, tasks, feature types, and dataset sizes, using consistent evaluation splits and broad baseline comparisons. Hyperparameter searches are documented separately for NPT and each baseline model class.

  • Dataset coverage: Datasets vary in instance count, feature count, categorical or continuous composition, and task type.
  • Evaluation setup: Evaluation uses canonical test sets when available, otherwise 10-fold cross-validation or fixed 0.7/0.2/0.1 and 0.7/0.1/0.2 splits depending on dataset size.
  • Hyperparameter searches: Hyperparameter documentation covers sweep counts, search spaces, and model-specific settings for tree models, MLPs, and k-NN.

F Societal Impacts of NPT

The paper identifies unresolved societal and operational concerns around deploying NPTs, especially explanation, fairness, privacy, and computational cost. It therefore advises against using NPT in applications where explanations, fairness, or privacy are desired or legally required at this stage.

  • Deployment scope: The authors do not recommend careless production deployment because NPTs and their properties remain insufficiently understood.
  • Explanations: Individual NPT decisions have not been sufficiently investigated for explanation, despite Transformer-based architectures being amenable to explanations.
  • Fairness: NPTs have not yet been investigated with common calibration techniques for satisfying definitions of fairness.
  • Privacy: Because NPT directly uses training data for prediction, it is not a private model by definition.
  • Computational cost: Training one NPT is expensive relative to one tree-based baseline such as XGBoost, although the authors note similar costs in some settings.
  • Deployment scope: The paper advises against NPT use where explanations, fairness, or privacy are desired or legally required at this stage.
Loading 2106.02584v2…