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 has the following contributions: * Propose Non-Parametric Transformers (NPT), a trained transformer applying attention across training data points to make predictions. * Demonstrate that attention across datapoints is crucial for NPTs. * Display good results of transformers on tabular datasets. * Apply a self-supervised objective to tabular data by masking some features.