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

1 Introduction . From CNNs [ 57 ] to Transformers [ 90 ] , most of supervised deep learning relies on parametric modeling : models learn parameters θ from a set of training data Dtrain = { ( x1 , y1 ) , . . . , ( xn , yn ) } to maximize training likelihoods p ( y | x ; θ ) mapping from features x ∈ X to target values y ∈ Y . At test time , they then make a prediction p ( y∗ | x∗ ; θ ) that depends only on those parameters θ and the test input x∗ . That is , parametric models do not consider direct dependencies between datapoints . This paper challenges parametric modeling as the dominant paradigm in deep learning . Based on the same end-to-end learning motivations that underpin deep learning itself , we consider giving models the additional flexibility of using training data directly when making predictions p ( y∗ | x∗ , Dtrain ; θ ) . Concretely , we introduce Non-Parametric Transformers ( NPTs ) : a general deep learning architecture that takes the entire dataset as input and predicts by explicitly learning interactions between datapoints ( Fig . 1 ) . NPTs leverage both parametric and non-parametric predictive mechanisms , with the use of end-to-end training allowing the model to naturally learn from the data how to balance the two . Namely , instead of just learning predictive functions from the features to the targets of independent datapoints , NPTs can also learn to reason about general relationships between inputs . We use multi-head self-attention [ 4 , 59 , 90 ] to model relationships between datapoints and construct ∗Equal Contribution . Correspondence to { jannik.kossen , neil.band } @ cs.ox.ac.uk . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . a training objective for NPTs with a stochastic masking mechanism inspired by self-supervised reconstruction tasks in natural language processing [ 24 ] . We show that these models learn to look up information from other datapoints and capture the causal mechanism generating the data in semi-synthetic settings . However , unlike conventional non-parametric models , NPTs are not forced to only make predictions in this way : they can also use the power of ordinary parametric deep learning . Background . While questioning parametric modeling assumptions is unconventional in deep learning , in statistics , so-called non-parametric models are a well-known and long-established field of study . Non-parametric models make predictions in explicit dependence of the training data p ( y∗ | x∗ , Dtrain ) . The most popular example of such models in the machine learning community are perhaps Gaussian Processes [ 74 ] . Non-parametric models typically do not require any training of parameters , and instead often directly interpolate between training points according to a fixed procedure , e.g. , [ 74 , p.17 ] . The interactions between inputs are fully defined by architectural choices and a small set of hyperparameters that must be carefully chosen . Conventional non-parametric models can not learn – in the sense familiar to deep learning practitioners – interactions from the data , limiting the flexibility these models have in adapting to the data at hand . Approaches such as Deep Gaussian Processes [ 22 ] , Deep Kernel Learning [ 95 ] , and Neural Processes [ 36 , 37 , 49 ] have all sought to apply ideas from deep neural networks to non-parametrics . Compared to NPTs , these approaches rely heavily on motivations from stochastic processes . This leads to them being either less flexible than NPTs or requiring strong assumptions on the data , making them inapplicable to the practical scenarios considered in this paper ( cf . §3 ) . Unlike previous work , NPTs explicitly learn to predict from interactions between datapoints , and they can be applied to general supervised machine learning tasks . We refer to §3 for an overview of these and other related approaches . A key contribution of this paper is opening the door to a more general treatment of how deep learning models can make use of dependencies between datapoints for predictions . Our results demonstrate that NPTs make use of interactions between datapoints in practice , and we show highly competitive performance on several established tabular datasets as well as early image classification results . Additionally , we show that NPTs can solve complex reasoning tasks by combining representation learning and cross-datapoint lookup ; something that is impossible for conventional deep learning or non-parametric models due to their inability to learn relations between datapoints . We next discuss the specifics of our model ( §2 ) , before moving on to related work ( §3 ) , empirical results ( §4 ) , and finally , limitations , future work , and conclusions ( §5 ) . 2 Non-Parametric Transformers . Non-Parametric Transformers ( NPTs ) explicitly learn relationships between datapoints to improve predictions . To accomplish this , they rely on three main ingredients : ( 1 ) We provide the model with the entire dataset – all datapoints – as input . We approximate this with minibatches where necessary for large data . At test time , both training and test data are input to the model ; during training , the model learns to predict targets from the training data ( §2.6 ) . ( 2 ) We use self-attention between datapoints to explicitly model relationships between datapoints . For example , at test time , the attention mechanism models relationships amongst training points , amongst test points , and between the two . ( 3 ) NPT ’ s training objective is to reconstruct a corrupted version of the input dataset . Similar to BERT [ 24 ] , we apply stochastic masking to inputs and minimize a loss on predictions at entries masked out in the input . Next , we introduce the three components in detail . 2.1 Datasets as Inputs . NPTs take as input the entire datasetX ∈ Rn×d . The datapoints are stacked as the rows of this matrix { Xi , : ∈ Rd | i ∈ 1 . . . n } , and we refer to the columns as attributes { X : ,j ∈ Rn | j ∈ 1 . . . d } . Each attribute is assumed to share a semantic meaning among all datapoints . In single-target classification and regression , we assume that the targets ( labels ) are the final attribute X : ,d , and the other attributes { X : ,j | j 6= d } are input features , e.g. , the pixels of an image . Each Xi , j is an entry or value . In addition to tabular data , many modalities such as images , graphs , or timeseries can be reshaped to fit this format . Note that this is a departure from common notation for supervised learning as introduced in §1 , as the inputX now includes both features and targets ( collectively , attributes ) . In masked language modeling [ 24 ] , mask tokens denote which words in a sentence are unknown and where , at training time , model predictions will have a loss backpropagated . Analogously , we use a binary matrix M ∈ Rn×d to specify which entries are masked in the input X . This matrix is also passed to NPT as input . The task is to predict the masked values XM = { Xi , j |Mi , j = 1 } from the observed valuesXO = { Xi , j |Mi , j = 0 } , i.e. , to predict p ( XM |XO ) . In summary , NPT takes as input the entire dataset and masking matrix ( X , M ) , and makes predictions X̂ ∈ Rn×d for values masked at input . This general setup accommodates many machine learning settings simply by adjusting the placement of the binary masks in M . We focus on single-target classification and regression – corresponding to a masking matrixM with 1s at all entries of the label columnX : ,d – but outline multi-target settings , imputation , self-supervision using input features , and semi-supervision in Appendix C.4 . Next , we describe the NPT architecture . 2.2 NPT Architecture . An overview of the Non-Parametric Transformer ( NPT ) is depicted in Fig . 2 . NPT receives the dataset and masking matrix ( X , M ) as input ( Fig . 2a ) . We stack these and apply an identical linear embedding to each of n datapoints , obtaining an input representationH ( 0 ) ∈ Rn×d×e ( Fig . 2b ) . Next , we apply a sequence of multi-head self-attention layers [ 4 , 24 , 90 ] . Crucially , we alternatingly apply attention between datapoints and attention between attributes of individual datapoints ( Figs . 2c-d ) . These operations allow our model to learn both relationships between datapoints as well as transformations of individual datapoints . Finally , an output embedding gives the prediction X̂ ∈ Rn×d , which now has predicted values at entries that were masked at input . We refer to Appendix C.3 for details , such as treatment of categorical and continuous variables . Importantly : Property 1 . NPTs are equivariant to a permutation of the datapoints . ( cf . Appendix A for proof . ) In other words , if the set of input datapoints is shuffled , NPTs produce the same prediction but shuffled in an analogous manner . This explicitly encodes the assumption that the learned relations between datapoints should not depend on their ordering . At a high level , permutation-equivariance holds because all components of NPTs are permutation-equivariant , and the composition of permutationequivariant functions is itself permutation-equivariant . We now briefly recap multi-head self-attention which plays an important role throughout the NPT architecture . 2.3 Multi-Head Self-Attention . Multi-head self-attention ( MHSA ) is a powerful mechanism for learning complex interactions between elements in an input sequence . Popularized in natural language processing [ 4 , 24 , 90 ] , MHSA-based models have since been successfully applied to many areas of machine learning ( cf . §3 ) . Dot-product attention computes attention weights by comparing queries { Qi ∈ R1×hk | i ∈ 1 . . . n } with keys { Ki ∈ R1×hk | i ∈ 1 . . .m } , ultimately updating the representation of the queries by aggregating over values { Vi ∈ R1×hv | i ∈ 1 . . .m } via the attention weights . We stack the queries , keys , and values into matrices Q ∈ Rn×hk , K ∈ Rm×hk , and V ∈ Rm×hv and , as is commonly done for convenience , assume hk = hv = h. Then , we compute dot-product attention as Att ( Q , K , V ) = softmax ( QKT / √ h ) V . ( 1 ) Multi-head dot-product attention concatenates a series of k independent attention heads MHAtt ( Q , K , V ) = concat axis=h ( O1 , . . . , Ok ) W O , where ( 2 ) Oj = Att ( QW Q j , KW K j , VW V j ) . ( 3 ) We learn embedding matrices WQj , W K j , W V j ∈ Rh×h/k , j ∈ { 1 , . . . , k } for each head j , and WO ∈ Rh×h mixes outputs from different heads . Here , we focus on multi-head self -attention , MHSelfAtt ( H ) = MHAtt ( Q = H , K = H , V = H ) , which uses the same inputs for queries , keys , and values . Following Transformer best practices to improve performance [ 16 , 24 , 59 , 66 , 90 ] , we first add a residual branch and apply Layer Normalization ( LN ) [ 3 ] followed by MHSelfAtt ( · ) , Res ( H ) = HW res + MHSelfAtt ( LN ( H ) ) , ( 4 ) with learnable weight matrix W res ∈ Rh×h . Then , we add another residual branch with LN and a row-wise feed-forward network ( rFF ) , finally giving the full multi-head self-attention layer as MHSA ( H ) = Res ( H ) + rFF ( LN ( Res ( H ) ) ∈ Rn×h . ( 5 )

Summary: This paper introduces a neural network architecture termed Non-Parametric Transformer (NPT) that uses self-attention (i) between data points and (ii) between features within each data point. The core conceptual contribution is to treat supervised deep learning problems as tasks where the model receives the entire dataset (incl. both labeled and unlabeled data points) as input — in contrast to the common approach of providing data samples independently to a deep learning model. The proposed NPT model is evaluated on tabular datasets where the authors report competitive results compared to boosted decision trees (which represent the state of the art). Applied on MNIST and CIFAR10, the NPT model performs reasonably well. The authors further perform a series of experiments to analyze the learned attention patterns and the importance of self-attention between data points for generalization performance.