Kernel Identification Through Transformers

1 Introduction . In recent years deep parametric models have become a prominent class of model for supervised learning and have delivered impressive empirical performance over a wide range of tasks . An important limitation , however , is that in their conventional form deep models do not provide prediction uncertainty . While their Bayesian counterparts try to achieve this , they require significant modifications to the training procedure and are computationally expensive . Uncertainty quantification in deep models is widely considered to be an open problem , the large array of research proposing alternative Bayesian neural networks underscores this [ 10 , 12 , 17 ] . On the other hand , kernel driven methods within the Bayesian framework like Gaussian processes ( GPs ) account for prediction uncertainty by design . While GPs provide a flexible framework for inferring distributions over functions , the inductive biases are controlled by the kernel function1 . A well chosen kernel will typically yield dramatically better performance than a poorly chosen one . How should we learn expressive kernels for high-dimensional tasks ? This has frequently been highlighted as a central question for the continued relevance of GP methods [ 11 ] . This work uses representations generated by a deep neural network to identify suitably expressive kernels for highdimensional GP regression tasks . Kernel recommendation is performed by a decoder with access to a large vocabulary of primitive kernels and products of primitive kernels . The decoder maps an ∗Work undertaken while at Secondmind 1also called covariance function or covariance kernel 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . encoded representation of a dataset to a kernel that can be used to model it . The representation is attained by encoding the dataset , treated as a sequence of ( input , output ) pairs D = { xi , yi } Ni=1 , utilising the permutation-equivariant nature of self-attention networks . By training KITT with a sufficiently rich vocabulary of kernels , it can predict suitable kernels for a diverse array of real datasets . This work presents the following novel contributions : • Inspired by the successes of image captioning networks , we develop a novel framework named KITT for amortised kernel search . KITT takes raw datasets for predictive modelling as input and proposes kernels composed from a large vocabulary of kernel functions . • KITT ’ s architecture introduces two novel features : it is entirely agnostic to the length and dimensionality of the data we wish to perform inference on , and it offers double permutation invariance ( this ensures its outputs are invariant to permutations in either input dimensions or data points ) . • We show that KITT can deliver kernel predictions in under 0.1 seconds . • We introduce a novel variant of the linear kernel which forms a key component of KITT ’ s vocabulary . • We demonstrate that the kernels identified by KITT offer strong performance against other baselines which deal with kernel engineering in the context of GPs . 2 Background . This section offers a brief review of the two topics which are central to this work , namely Gaussian Processes and Transformers . Gaussian Processes . GPs offer a highly versatile framework for predictive modelling [ 23 ] with generalisation properties controlled by a kernel function parameterised by hyperparameters . The functional form of the kernel governs the global attributes of the supported functions , such as smoothness and periodicity . However , a suitable kernel function is unknown a priori and the choice of the kernel function is a fundamental model selection problem . Once a kernel has been chosen , training conventionally proceeds by learning a point estimate of the hyperparameters that maximise the GP log marginal likelihood θ∗ = argmaxθ log p ( y|θ ) . The marginal likelihood is available in closed-form for models with Gaussian likelihoods . Below we briefly summarise the standard GP framework . GPs are distributions over functions from which one can sample realised function values for given inputs . Concretely , for observations X = { xi } Ni=1 , and positive definite kernel function kθ ( · , · ) , with hyperparameters θ , f ( · ) ∼ GP ( 0 , kθ ) . Typically , we observe noisy realisations of the latent function which are corrupted with Gaussian noise , yi = f ( xi ) + i , i ∼ N ( 0 , σ2n ) , and infer the kernel hyperparameters through maximising the likelihood of the model . The GP marginal likelihood objective , p ( y|θ ) , is obtained by marginalising the likelihood y|f ∼ N ( f , σ2nI ) over the prior f |X , θ ∼ N ( 0 , Kθ ) , p ( y|θ ) = ∫ p ( y|f ) p ( f |θ ) df = N ( 0 , Kθ + σ2nI ) , ( 1 ) where f = f ( X ) denotes a vector of realised function values , andKθ denotes theN×N covariance matrix corresponding to evaluations of the covariance function at the N training inputs , ( Kθ ) i , j = kθ ( xi , xj ) . A long standing question is how best to select an appropriate kernel function for a given task . One approach is to search over a discrete space of kernels , defined by combining a selection of primitive kernels with a predefined grammar [ 2 , 6 ] . Typically , a greedy search is performed to identify the kernel offering the best representation of the data . Ideally , the quality of the kernel is quantified by the Bayesian model evidence , which can be computed by marginalising the marginal likelihood over the hyperparameters . However , since the integral is challenging to compute , each kernel ’ s suitability is instead usually determined via a proxy for the model evidence , such as the Bayesian Information Criterion ( BIC ) [ 25 ] . For a principled , Bayesian approach to kernel design , instead of selecting a single kernel , one ought to consider multiple candidates . In other words , it is desirable to marginalise over the space of kernels , not only over the space of functions for a single kernel . p ( y|D ) = ∑ i p ( y|Ki , D ) p ( D|Ki ) p ( Ki ) . ( 2 ) This yields a rich posterior distribution comprised of a mixture of Gaussians . A conventional kernel search makes three key approximations . First , that contributions from all but the single chosen kernel K∗ can be neglected ; second , that contributions from all but the maximum likelihood hyperparameters θ∗ can be neglected ; third , that the proxy for the model evidence is a reliable one . There are several regimes , for example where the data is sparse or noisy , where all three of these assumptions do not hold . We shall aim to improve upon all three of these issues . Transformers Transformers are a form of deep neural network which rely upon the attention mechanism [ 32 ] to capture global context . While they were originally proposed to tackle machine translation tasks [ 29 ] , they have rapidly attained state-of-the-art performance in a number of other areas of machine learning [ 14 , 21 , 22 ] . Of particular relevance to this work , they have been successfully applied to image captioning [ 5 , 33 ] , a task which involves summarising the key characteristics of rich data in a grammatical form . This has a striking parallel to the challenge of selecting an appropriate kernel , especially since a form of grammar can be used to construct a broad selection of GP kernels . The self-attention mechanism of transformers naturally lends itself to the permutation-invariant setting , as demonstrated by Zaheer et al . [ 34 ] . This invariance to permutations in the ordering of inputs has been exploited to create Set Transformers , introduced in Lee et al . [ 18 ] . The AHGP model [ 19 ] uses a variant of the Set Transformer to infer the hyperparameters of the spectral mixture kernel . Like all deep networks , transformers thrive when presented with an abundance of training data . Fortunately , for the task we have at hand , the training set is unlimited in size as we may sample training data from GPs with known kernels . As stressed by Liu et al . [ 19 ] , selecting an architecture which reflects the appropriate invariances is vital . The AHGP model is designed to be invariant to permutations in the ordering of the datapoints . We go one step further , and introduce a model which is doubly permutation invariant : its output is also invariant to permutations in the input dimension ( whereas AHGP is equivariant to the input dimensions ) . 3 KITT . In this section we describe and motivate KITT , a network which takes as inputs a set of datapoints { xi , yi } , and outputs a kernel recommendation in the form of a ‘ caption ’ , by utilising a kernel grammar . The code is available at https : //github.com/frgsimpson/kitt . Kernel Grammar and Vocabulary : Throughout the GP literature , the most commonly used kernels belong to a limited set of primitive functions . In this work , we utilise eight primitive kernels : the squared exponential ; periodic ; white noise ; three variants of the Matern kernel ( 1 2 , 3 2 , 5 2 ) ; the cosine kernel , and our novel variant of the linear kernel . This list comprises six stationary isotropic kernels , one stationary anisotropic kernel ( cosine ) , and one non-stationary kernel ( linear ) . From this small set of primitive kernels , we wish to construct a larger array of more expressive kernels . This can be achieved by leveraging the closure properties of kernel functions [ 27 ] . Permitted operations include addition , multiplication , convolution , composition , and affine transformations . For simplicity , this work shall only consider two operators : addition and multiplication . Whilst this grammar of addition and multiplication appears superficially simple , it has some idiosyncrasies that would make it challenging for a network to learn . For example , if a network is unaware that multiplication is commutative , it would have to learn to recognise k1 ∗ k2 and k2 ∗ k1 separately , for all combinations of k1 and k2 . It would also need to learn that multiplying a noise kernel with another stationary kernel yields another noise kernel . Encoding this information a priori greatly facilitates the learning process . To achieve this , we enlarge the vocabulary by defining product kernels as single ‘ words ’ , rather than incorporating products as part of the grammar ( the full list is provided in the Appendix ) . This allows us to exclude redundant combinations from consideration . For example , only a single token is used to represent both k2 ∗ k1 and k1 ∗ k2 . While this enlarges the vocabulary , we can be confident that the network will be able to cope , since there will still be far fewer ‘ words ’ than can be found in natural language tasks where transformers are known to excel . We use products of two primitive kernels as part of the base vocabulary and found that incorporating higher order products do not significantly change performance . A further advantage of defining product kernels at the vocabulary level is that we need no longer include operators inside the vocabulary . The multiplication is already baked into the expanded set of kernels , while the addition takes place implicitly , much like the white space between words in a natural language task . This precludes the construction of nested structures of operators , which would allow for even richer kernels . However , since this work represents a first attempt at performing a kernel search with a neural network , we choose to keep the captioning task to be a relatively simple one , and leave more complex grammatical compositions as an opportunity for future exploration . Even with this relatively simple grammar , if we permit a caption of four words from a base vocabulary of 32 , we are effectively searching across a space of around 36 000 kernels . Priors . While the priors we impose upon the hyperparameters will have some impact during optimisation , it is their influence upon the generation of KITT ’ s training data that is of central importance to this work . Random samples are drawn from the hyperparameter priors p ( θ ) and input locations x ∼ U ( −2.5 , 2.5 ) before each random sample of y is generated . The variances and lengthscale parameters of all kernels ( including product kernels ) are assigned lognormal priors , such that log θ ∼ N ( 0 , 1 ) . The cosine kernel is unique among KITT ’ s vocabulary , in that it is inherently anisotropic , which can be important if a preferred direction exists within the data . Samples drawn from the kernel manifest as plane waves which propagate along a characteristic direction of the kernel . If we were to impose a lognormal prior ( or any other positively constrained prior ) on its lengthscales , this would restrict the direction of the kernel to a small fraction 21−D of its permitted parameter space . We therefore adopt a different approach in this case , assigning ` ∼ Cauchy ( 0 , 5 ) for the lengthscales . Training . KITT is trained entirely on synthetic data . Each training example is generated from a randomly selected kernel , with a randomly drawn set of hyperparameters . This data could be generated on the fly , but since this can be a computationally expensive process , we generated a training set in advance which comprised of 200 , 000 labeled examples . While the model can be constructed for an arbitrary number of inputs points and input dimensions , during training we restrict ourselves to the case of 4 input dimensions and 64 input points per sample . The loss function corresponds to − log p ( k|D ) , the negative log probability the network assigned to the correct term in the vocabulary . The Adam optimiser was used with an initial learning rate of 10−4 , and a decay schedule with a decay rate of 0.1 every 50 , 000 iterations . Due to the relatively noisy nature of the classification task , a relatively large batch size of 128 was found to be beneficial . The vocabulary included product kernels of at most two terms in addition to the primitive kernels , yielding a final vocabulary of size 34 . Heteroscedastic noise . When taking the product between the white noise and any stationary kernel , we recover another noise kernel . These redundant expressions are omitted from KITT ’ s vocabulary . However , the result of the product between the noise kernel and the ( non-stationary ) linear kernel merits special attention . The linear kernel is defined as KLIN ( x , x ′ ) = σ2 ( x− c ) ( x′ − c ) , ( 3 ) where σ2 denotes the vector of variances for the linear kernel , and c represents the shift parameter . Unlike the other primitive kernels , the linear kernel possesses independent variance terms for each input dimension . The product between the linear kernel and the noise kernel generates a form of noise whose variance changes linearly with respect to the inputs . This presents an opportunity to model heteroscedastic noise , within the conventionally homoscedastic domain of GP regression models . When modelling real world tasks , this potentially offers a major advantage , since the noise variance often changes across the input space . Note that if we simply set c = 0 , as is often assumed when working with the linear kernel , then the linearly varying noise term is extremely limited . The noise variance could only ever increase as we move away from the origin . As with the cosine kernel , this reduces us to a small fraction ( 21−D ) of the viable parameter space . In order to lift this restriction , we introduce a ‘ shift ’ vector in the linear kernel , such that the origin is free to move along each input dimension . This naturally leads us to ask what an appropriate prior for this shift vector would be . We consider it equally likely that the noise amplitude increases or decreases with x . To reflect this belief , and accounting for our normalised inputs , we seek a prior of the form dσ2/dx ∼ N ( 0 , 0.1 ) . For large displacements , we note that the shift parameter can be approximately expressed as the ratio of two normally distributed variables : the gradient of the noise and the noise amplitude at the origin . This observation suggests a suitable prior on c is given by the Cauchy distribution . We adopt c ∼ Cauchy ( 0 , 5 ) throughout . Model Architecture . The kernel ’ s likelihood p ( D|K ) is invariant to permutations of the ordering of the datapoints , and to permutations of the ordering of the input dimensions . KITT ’ s output of kernel recommendations should therefore exhibit these two important properties . Zaheer et al . [ 34 ] demonstrated that a function f ( X ) which is invariant to permutations in the elements of X can be expressed in the form ρ ( ∑ i φ ( xi ) ) ) , where ρ and φ are differentiable functions . This was exploited by Lee et al . [ 18 ] in constructing the Set Transformer . The scenario we encounter is slightly more complex in that there is a two-tiered hierarchy of invariances . We seek a function over the training set D , which can be expressed as a collection of high dimensional data vectors Di : f ( { D } ) = ρ [ ∑ i φ ( Di ) ] . Here the function φ ( Di ) must be invariant over the permutations of the different input dimensions . This can therefore be decomposed in a similar manner , φ ( Di ) = ρ′ [ ∑ j φ ′ ( Dij ) ] . Combining these two equations leaves a final expression of the form f ( { D } ) = ρ ∑ i ρ′ ∑ j φ′ ( Dij ) , ( 4 ) where i and j can be either dimensions or datapoints . This can be interpreted as : encode over j ; pool over j ; encode over i ; pool over i ; decode . This formalism sets the foundations for our choice of architecture , as shown in Figure 1 . At a lower level , KITT is comprised of the following components whose acronyms we define here : • rFF : Row-wise feed-forward layer with ReLU •MP : Mean pooling function • SAB : Set Attention Block [ 18 ] • LayerNorm : Layer Normalisation [ 1 ] •Multihead : Multi-Headed Attention Mechanism [ 29 ] . Encoder : The encoder has the architecture of a transformer with self attention blocks . Our goal is to encode datasets Dj = { xi , yi } Ni=1 of shape N × ( D + 1 ) where xi ∈ RD and yi ∈ R ; we need to incorporate both invariance to ordering of the data points ( row-wise shuffle ) and equivariance to a re-ordering of the dimensions ( column-wise shuffle ) . In order to achieve this , our encoder has two sub-components responsible for encoding along the sequence and dimension axes respectively : 1 . SEQ ENC : RN× ( D+1 ) → RD×E . A sequence encoding component which acts on input datasets and outputs dimension-level representations G ≡ { gd } Dd=1 , gd ∈ RE where E is the embedding dimension . The sequence encoder forward pass is formulated as : SEQ ENC ( D ) = MP ( SAB×6 ( rFF ( D ) ) ) Where mean pooling is applied over the sequence . Our implementation of the SAB component [ 18 ] differs slightly from that of Lee et al . [ 18 ] , SAB ( Z ) = LayerNorm ( C + Z ) , where we have defined C = Dropout ( rFF [ Multihead ( Z , Z , Z ) ] ) [ 29 ] . 2 . DIM ENC : RD×E → RD×E A dimension encoding component : DIM ENC which acts on dimension level encodings to generate final representations { hd } Dd=1 of dimension E. The dimension encoder forward pass is formulated as : DIM ENC ( G ) = rFF×2 ( SAB×6 ( G ) ) The encoder forward pass entails passing each input data set to the sequence encoding component followed by the dimension encoding component . ENCODER ( D ) = DIM ENC ( SEQ ENC ( D ) ) Proposition 1 : Let SN denote the set of all permutations of the row-wise indices { 1 , 2 , . . . , N } and Dπ denote an input dataset with the ordering of indices given by π ∈ SN The sequence encoding component SEQ ENC is invariant to a permutation of the indices within a dataset D. SEQ ENC ( D ) = SEQ ENC ( Dπ ) ∀ π ∈ SN Proposition 2 : Let QD denote the set of all permutations of the column-wise indices ( dimensions ) { 1 , 2 , . . . , D } and Dν denote an input dataset with the ordering of indices given by an arbitrary ν ∈ QD . The sequence encoder SEQ ENC and dimension encoder DIM ENC components are equivariant to a permutation of the dimensions within a dataset D. ENCODER ( Dν ) = ν ( DIM ENC ( SEQ ENC ( D ) ) ) ∀ ν ∈ QD We include proofs for propositions 1 and 2 in the supplementary . Decoder : KITT ’ s decoder iteratively builds a caption from the encoded dataset representations , generated by the encoder described above , and a prompt which consists of the kernel expression thus far ( see Fig . 1 ) . Our decoder closely resembles the one proposed in Vaswani et al . [ 29 ] except that we remove the positional encodings and adjust the number of layers . The decoder uses selfattention blocks to first attend to the prompt and then to attend to the representations from the encoder using the processed dataset representations as query values . These two applications of attention are alternated in several layers until a new component kernel is proposed from a distribution generated from the final representations . We note that this component kernel proposition is invariant to the ordering of the dataset representations and thus to a shuffling of the dimensions of the original dataset . Hence , the model is fully invariant . The end-to-end process is depicted in Figure 1 and a detailed schematic of the decoder is included in the supplementary material . Scalability . Training of KITT is only performed once . Once training of the network is completed , all inference procedures require only a single forward pass through the network . As a result , instead of the O ( N3 ) cost commonly associated with explicit marginal likelihood evaluations , the cost of a forward pass through KITT is O ( DN2 + D2 ) . Due to effective parallelisation of the attention mechanism on GPUs , noted by Vaswani et al . [ 29 ] , and the modest size of the KITT network , we experience near constant wall clock time in practice . Inference . Given some new dataset D , inference proceeds as below , closely mimicking the procedure used to generate a caption for images . Note that the best caption is not necessarily constructed by choosing the best kernel at each step of the decoder . 1 . Pass the data D into the encoder ; pass the resulting encodings and an ( initially empty ) kernel expression E into the decoder . 2 . Retrieve the output probabilities , selecting the kernel k with the highest probability and append the chosen kernel ( or operator ) to our full kernel expression E . 3 . Repeat steps 1 & 2 until either the < STOP > token is selected or the maximum caption length is reached . 4 . Repeat steps 1-3 several times but now select kernels stochastically , weighted by their probability , to construct a set of high-ranking kernel expressions . 5 . For each of the top three candidate kernels expressions , as ranked by their total probability assigned by the network , we optimise the associated hyperparameters with BFGS [ 9 ] . 6 . Combine the posterior distribution of the forecasts based on either the overall probability assigned by the network , or another proxy for the model evidence such as the Bayesian Information Criterion . In summary , we select a kernel based upon the output of the pretrained KITT network , before optimising the hyperparameters in a conventional manner .

Summary: The paper presents a new approach to select kernel functions via transformers. The proposed approach mimics image captioning tasks where the model takes the pair of data covariates and targets as image inputs,and then produces kernel grammars as captions. The paper designs a transformer architecture in the way that it is invariant to permutations across input indices and input dimensions.