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 authors propose KITT, a novel method exploiting a transformer-based architecture to generate kernel recommendations. KITT has two main features: (1) agnostic to the length and dimensionality of the data; (2) it offers double permutation invariance. They demonstrate the effect of KITT on synthetic data compared with other baselines.