PaperlayerView exact revision on arXiv

Source-linked AI summary

Kernel Identification Through Transformers

Fergus Simpson, Ian Davies, Vidhi Lalchand, Alessandro Vullo, Nicolas Durrande, Carl Rasmussen

arXiv:2106.08185v2stat.MLcs.LG

TL;DR

Choosing expressive kernels for high-dimensional GP regression is challenging, although kernel choice controls GP inductive biases and function support. KITT uses a transformer to recommend kernels quickly, and its selections perform strongly against kernel-engineering baselines, including consistent advantages over AHGP and competitiveness with slower tree search.

  • Problem

    Constructing suitable, expressive kernels for high-dimensional GP regression is challenging, and kernel choice determines GP inductive biases and prior function support.

    Source: §1 · p. 1

  • Method

    KITT uses a transformer-based model and a vocabulary of kernels to map datasets to kernel recommendations, with inference under 0.1 seconds.

    Source: §1 · pp. 1–2

  • Results

    KITT consistently outperforms AHGP and is competitive with the far slower tree-search method in NLPD comparisons.

    Source: §4 · p. 8

  • Takeaways & Limitations

    KITT can select suitable kernels for previously unseen real-world datasets despite being trained solely on synthetic data.

    Source: §6 · p. 10

    For the linear kernel, setting c = 0 restricts the noise variance to increasing away from the origin, leaving only a small fraction of the viable parameter space.

    Source: §3 · p. 5

Abstract

from arXiv · show

Kernel selection plays a central role in determining the performance of Gaussian Process (GP) models, as the chosen kernel determines both the inductive biases and prior support of functions under the GP prior. This work addresses the challenge of constructing custom kernel functions for high-dimensional GP regression models. Drawing inspiration from recent progress in deep learning, we introduce a novel approach named KITT: Kernel Identification Through Transformers. KITT exploits a transformer-based architecture to generate kernel recommendations in under 0.1 seconds, which is several orders of magnitude faster than conventional kernel search algorithms. We train our model using synthetic data generated from priors over a vocabulary of known kernels. By exploiting the nature of the self-attention mechanism, KITT is able to process datasets with inputs of arbitrary dimension. We demonstrate that kernels chosen by KITT yield strong performance over a diverse collection of regression benchmarks.

1 Introduction

Conventional deep models lack prediction uncertainty, while Gaussian Processes provide it but depend on choosing a suitable kernel. KITT uses neural representations and a transformer-based decoder to recommend kernels from a vocabulary, with fast inference and strong performance against kernel-engineering baselines.

1 Introduction

  • Deep models conventionally lack prediction uncertainty, whereas Gaussian Processes account for it by design and depend on their kernel to set inductive biases.The introduction frames uncertainty quantification in deep models as an open problem and notes research on alternative Bayesian neural networks [10] [17].Source: §1 · p. 1
  • KITT addresses the challenge of learning expressive kernels for high-dimensional tasks by using deep neural network representations to identify kernels.A decoder with access to a large vocabulary maps an encoded dataset representation to a kernel; the supplied description of this pipeline continues in p10.Source: §1 · pp. 1–2
  • The dataset is encoded as a sequence of input-output pairs, and self-attention is used to obtain a representation that can be mapped to a kernel.The paper describes the approach as agnostic to dataset length and dimensionality.Source: §1 · p. 2
  • KITT’s recommended kernels offer strong performance against other baselines for GP kernel engineering, with recommendations produced in under 0.1 seconds.The framework is designed for high-dimensional regression tasks and uses a vocabulary of primitive kernels and their products.Source: §1 · pp. 1–2

2 Background

Gaussian Processes support predictive modelling, but their function priors and generalisation properties depend on a kernel that is not known in advance. The section reviews kernel search and model averaging, then motivates transformer components through attention, permutation invariance, and the availability of synthetic training data.

Gaussian Processes

  • A GP kernel controls the supported functions’ global attributes, such as smoothness and periodicity, making kernel choice a fundamental model-selection problem.After selecting a kernel, conventional training learns a point estimate of its hyperparameters by maximizing the GP log marginal likelihood.Source: §2 · p. 2

Kernel search

  • Conventional kernel search combines primitive kernels with a grammar, typically uses greedy search, and relies on proxies such as BIC because marginalizing over hyperparameters is difficult.The ideal criterion described is Bayesian model evidence marginalized over hyperparameters.Source: §2 · p. 2

Bayesian kernel selection

  • Considering multiple kernel candidates yields a mixture-of-Gaussians posterior, while conventional search approximates away all but one kernel, maximum-likelihood hyperparameters, and proxy evidence.The text notes that these assumptions may fail in sparse or noisy data regimes.Source: §2 · pp. 2–3

Transformers

  • Transformer attention captures global context and has been used for image captioning, whose grammar-based summaries parallel constructing GP kernels from a grammar.The paper also notes the relevance of permutation-invariant self-attention and Set Transformers; AHGP uses a Set Transformer variant for spectral-mixture-kernel hyperparameters.Source: §2 · p. 3

Training data

  • Synthetic GP samples with known kernels provide an unlimited-size training source for transformers, which generally benefit from abundant training data.The paper states that training examples can be sampled from GPs with known kernels.Source: §2 · p. 3

Permutation invariance

  • KITT is designed to be invariant to permutations of both datapoints and input dimensions, extending AHGP’s datapoint invariance while changing its input-dimension behavior.AHGP is described as equivariant to input dimensions, whereas KITT’s output is invariant to their permutations.Source: §2 · p. 3

3 KITT

KITT takes a dataset and generates kernel recommendations as expressions from a predefined vocabulary. Its architecture handles permutations of datapoints and input dimensions, while training and inference use a synthetic-data-trained transformer to recommend and rank candidate kernels.

Kernel Grammar and Vocabulary

  • KITT uses eight primitive kernels and addition and multiplication to construct more expressive kernel expressions.Products of two primitive kernels are treated as vocabulary items, including a single token for commutative equivalents; higher-order products did not significantly change performance.Source: §3 · pp. 3–4

Training and Priors

  • Training uses 200,000 synthetic examples generated from randomly selected kernels and hyperparameters, with 4 input dimensions and 64 datapoints per sample.The vocabulary has 34 kernels, including products of at most two terms; kernel variances and lengthscales use lognormal priors, while cosine lengthscales use a Cauchy prior.Source: §3 · p. 4

Heteroscedastic noise

  • Multiplying the linear kernel by the noise kernel models noise whose variance changes linearly with the inputs.A shift vector frees the linear kernel's origin to move along each input dimension, allowing the noise variance to increase or decrease with the inputs.Source: §3 · pp. 4–5

Model Architecture

  • KITT represents datasets with a two-tiered attention architecture that is invariant to datapoint order and produces kernel recommendations invariant to input-dimension order.The sequence encoder pools across datapoints, followed by a dimension encoder; the decoder builds a kernel expression from encoded data representations and the expression so far.Source: §3 · pp. 5–6

Scalability

  • KITT is trained only once; afterward, all inference procedures require a single forward pass costing O(DN^2 + D^2), rather than the O(N^3) cost commonly associated with explicit marginal-likelihood evaluations.GPU parallelization of attention and the modest network size yield near constant wall-clock time in practice [29].Source: §3 · p. 6

Inference

  • Inference generates candidate kernel expressions, optimizes hyperparameters for the top three, and combines forecast posteriors using network probabilities or another evidence proxy.The expression is extended until a stop token or maximum length is reached; additional candidates are generated by stochastic sampling weighted by their probabilities.Source: §3 · p. 7

4 Experimental Results

KITT is evaluated on synthetic kernel recovery and eight UCI regression tasks, where it generalizes to higher-dimensional structure and performs competitively with slower kernel search. Its predictions are also fast and can favor kernels with strong test performance.

Ground Truth Recovery

  • The encoder reliably recovers kernel structure on higher-dimensional tasks unseen during training, with reasonable accuracy requiring about 200 data points or more.Accuracy continues to improve as dataset size increases.Source: §4 · pp. 7–8

UCI Regression

  • KITT consistently outperforms AHGP and remains competitive with the much slower tree search on UCI regression tasks.The comparison uses negative predictive log likelihood values across the tasks.Source: Figure 4 · p. 8, §4 · p. 8
  • The UCI evaluation covers eight tasks with input dimensions from 4 to 14, using a 90/10 split and subsampling datasets above 2,000 points.Kernel captions have at most three terms; the classifier instead sums the three highest-scoring kernels.Source: §4 · p. 8
  • AHGP performs worse than the RBF baseline, which is a subset of the Spectral Mixture Product model.The authors suggest AHGP may have learned the primary component’s lengthscale.Source: §4 · p. 8
  • On Yacht, KITT assigns high probability to Linear × RBF, Linear × Matern32, and Linear × Matern52, the three strongest test performers among 34 kernels.Figure 5 compares predicted probabilities with realized test performance.Source: §4 · p. 8, Figure 5 · p. 9

Computational overhead

  • A kernel recommendation requires one encoder pass and a small number of decoder passes, each taking around 0.02 seconds.Prediction time is robust to increasing dataset size and dimensionality.Source: §4 · p. 8

5 Related work

Related work addresses GP model selection through kernel engineering, hyperparameter learning, deep representations, and alternative kernel constructions. KITT instead predicts compositional kernel labels from datasets using end-to-end training.

Amortised Hyperparameter Learning (AHGP)

  • AHGP uses self-attention to amortize GP hyperparameter learning with a fixed-component spectral mixture kernel, whereas KITT focuses on kernel design.AHGP estimates hyperparameters otherwise learned by maximizing marginal likelihood.Source: §5 · p. 9

Kernel Engineering

  • Automatic Statistician-style kernel engineering searches a finite space of composite structures built from primitive kernels, while KITT learns to regress kernel labels from datasets.Earlier handcrafted kernels often target one- or two-dimensional data.Source: §5 · p. 9

Deep Kernels

  • Deep-kernel methods transform inputs with a jointly trained neural network and GP, but can overfit when training millions of network parameters alongside GP hyperparameters.They apply a simple base kernel to learned representations rather than learning a sophisticated kernel for raw inputs.Source: §5 · p. 9

Novel Kernels

  • Spectral mixture kernels parameterize kernels through spectral density and provide closed-form replacements for stationary kernels.Source: §5 · p. 9

6 Discussion

The discussion presents KITT as a transformer trained on synthetic data that recommends suitable kernels for unseen real-world regression datasets. The authors report fast inference and pattern discovery across input dimensions and dataset sizes, while identifying one-dimensional outputs as the study’s scope.

Discussion

  • Trained solely on synthetic data, KITT can select suitable kernels for previously unseen real-world datasets.The experiments focus on high-dimensional real-valued regression and report ground-truth structure recovery on high-dimensional tasks.Source: §6 · p. 10
  • KITT recommends kernels in a fraction of a second and shows pattern discovery across a broad range of input dimensions and dataset sizes.The discussion contrasts its inference speed with greedy search and differentiable kernel networks.Source: §6 · p. 10
  • The work focuses on one-dimensional outputs; the authors propose that similar models could be developed for multi-output regression, classification, latent-variable modeling, and time-series prediction.Source: §6 · p. 10
  • The authors frame KITT as combining representation learning with inference over compatible kernels, and note its complementarity with AHGP’s hyperparameter optimization.They describe this approach as offering expressivity and ensemble uncertainty.Source: §6 · p. 10

Supplementary Material for Kernel Identification

The supplementary material reviews self-attention and multi-head self-attention. It describes attention as producing context-sensitive representations from sequences of vectors.

A Background: Self-Attention

  • Dot-product attention uses queries and keys to form weights that linearly map input values into an output representation.The values’ dimensions may differ from the query and key dimensions; row-wise softmax scaling is used to avoid exploding inputs that hamper training.Source: A Background: Self-Attention · p. 13
  • Self-attention derives queries, keys, and values from one input sequence, then computes a weighted sum of vectors for each index.The weights depend on the input vectors.Source: A Background: Self-Attention · p. 13
  • Multi-head self-attention repeats the operation with separate parameter sets and combines head outputs, with a final projection returning them to dimension D.Source: A Background: Self-Attention · p. 13

B.1 On encoding sets of datasets

The encoder treats each dataset as a set of input-output vectors, preserving invariance to datapoint order and equivariance to dimension order; the decoder makes kernel outputs invariant to both.

B.1 On encoding sets of datasets

  • Each dataset is represented as a set of N input-output vectors, with the encoder designed to preserve invariance to row order.Each vector contains the input coordinates and its output value.Source: §B.1 · p. 14
  • The sequence encoder achieves row-order invariance by reshaping the dataset, applying equivariant feed-forward and attention layers, then pooling across data points.The mean-pooling operation removes dependence on the ordering of the sequence.Source: §B.1 · pp. 15–16
  • The sequence and dimension encoders are equivariant to input-dimension permutations, so permuting dimensions permutes the encoded dimensions correspondingly.The dimension encoder is a stack of permutation-equivariant attention blocks.Source: §B.1 · p. 16
  • Because the decoder treats the encodings equally and uses no positional encoding, KITT's kernel recommendations are invariant to permutations of both datapoints and input dimensions.The encoder is invariant to datapoint order and equivariant to dimension order before decoding.Source: §B.1 · p. 17

C KITT Architecture

KITT encodes a dataset and iteratively decodes a kernel expression, reusing cached dataset encodings across decoder passes. The architecture omits positional encodings where the data or kernel algebra supplies the relevant structure.

C KITT Architecture

  • Figure 6's end-to-end architecture maps a dataset to a distribution over kernels using six attention blocks in each encoder and decoder.The decoder runs repeatedly with the same cached encodings, updating its prompt with the last proposed kernel; masked captions parallelize training predictions.Source: C KITT Architecture · p. 17
  • KITT's Set Attention Blocks use dropout at rate 0.1 and one residual connection with layer normalization, differing from the compared design's two such operations.The passage also notes a difference in implementation from the referenced Set Attention Blocks.Source: C KITT Architecture · p. 17
  • The transformer processes datapoint positions through their input values directly, so the dataset encoder does not require explicit positional encodings.Each datapoint's x values are part of the input to the transformer.Source: C KITT Architecture · pp. 17–18
  • Kernel prompts omit positional encodings because kernel addition is commutative; prior kernels are ordered by decreasing variance to retain some ordering information.The order of previously predicted kernels does not affect the next prediction.Source: C KITT Architecture · p. 18

D Further Experimental Results

The further experiments report an advantage for averaging the top three kernels over selecting only the top one, with both outperforming random selection. The priors section adjusts product-kernel lengthscales to avoid a systematic short-lengthscale cue.

E Priors

  • Product kernels use one learned variance hyperparameter, while each additive term has its own; product-component lengthscale priors are scaled to target the implied lengthscale.Without this correction, product kernels' shorter implied lengthscales could become a characteristic cue for the network.Source: E Priors · p. 18

D Further Experimental Results

  • Averaging the top three kernels offers an advantage over selecting the top one, and both approaches outperform random kernel selection.Table 4 reports predictive performance using RMSE for three model-averaging strategies.Source: D Further Experimental Results · p. 18, Table 4 · p. 19

E Priors

  • The lengthscale correction is approximate because propagation differs for kernels with non-Gaussian spectral densities, though the reported impact is relatively small.The scaling compensates for shrinkage in the implied product-kernel lengthscales.Source: E Priors · p. 19

F.1 Ground-Truth Experiments

Ground-truth experiments assess whether KITT classifies synthetic samples to their known primitive kernels across test sizes and input dimensions, with accuracy reported as the correct-classification fraction.

F.1 Ground-Truth Experiments

  • For each test-size and dimensionality combination, the experiments draw 300 GP samples from known primitive kernels and measure test-time classification accuracy.Accuracy is the percentage of samples classified as their ground-truth primitive kernel.Source: §F.1 · p. 19
  • Table 5 reports kernel-selection times on an Intel i7-8700 with a GTX 1070, and tree-search timings use at most 2,000 datapoints.The table reports means and one standard deviation over 10 repetitions.Source: Table 5 · p. 20
  • Means and standard errors are computed over three runs per combination; top-3 accuracy is reported for test size and top-3/top-5 accuracy for test dimensions.The test-dimension experiment fixes test inputs at N = 1600.Source: §F.1 · p. 20

F.2 ML-II Training

ML-II experiments initialize hyperparameters using the highest marginal likelihood among 1,000 prior draws. The broader-impact discussion argues that a pre-trained kernel-selection model could reduce experiment-time search and help users compare automated proposals with hand-crafted kernels.

F.2 ML-II Training

  • For each experiment, initial hyperparameter values were selected as the highest-marginal-likelihood choice among 1,000 random draws from the priors.Source: §F.2 · p. 20

G Broader Impact

  • A pre-trained solution could remove the burden of expensive kernel selection at experiment time and let users compare automated proposals with hand-crafted kernels.The authors suggest this approach may enjoy wide-spread adoption and popularity.Source: G Broader Impact · p. 20
  • Figure 8 contrasts triple-product kernel samples before and after induced lengthscale shrinkage, while Figure 9 shows the product prior and a scaled prior preserving the component kernels’ lengthscale range.Figure 9’s left panel compares the component lengthscale prior with the implied prior for a product of three RBF kernels; its right panel adds samples from the scaled prior.Source: Figure 8 · p. 21, Figure 9 · p. 21
Loading 2106.08185v2…