Kermut: Composite kernel regression for protein variant effects

Authors: peter mørch groth
mads herbert kerrn
lars  olsen
jesper  salomon
novonesis  wouter boomsma
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 15:05:57

Abstract: Reliable prediction of protein variant effects is crucial for both protein optimization and for advancing biological understanding. For practical use in protein engineering, it is important that we can also provide reliable uncertainty estimates for our predictions, and while prediction accuracy has seen much progress in recent years, uncertainty metrics are rarely reported. We here provide a Gaussian process regression model, Kermut, with a novel composite kernel for modeling mutation similarity, which obtains state-of-the-art performance for supervised protein variant effect prediction while also offering estimates of uncertainty through its posterior. An analysis of the quality of the uncertainty estimates demonstrates that our model provides meaningful levels of overall calibration, but that instance-specific uncertainty calibration remains more challenging.

Main paper: 
 introduction: 
 Accurately predicting protein variant effects is crucial for both advancing biological understanding and for engineering and optimizing proteins towards specific traits. Recently, much progress has been made in the field as a result of advances in machine learning-driven modeling [1] [2] [3] , data availability [4, 5] , and relevant benchmarks [6, 7] .While prediction accuracy has received considerable attention, the ability to quantify the uncertainties of predictions has been less intensely explored. This is of immediate practical consequence. One of the main purposes of protein variant effect prediction is as an aid for protein engineering and design, to propose promising candidates for subsequent experimental characterization. For this purpose, it is essential that we can quantify, on an instance-to-instance basis, how trustworthy our predictions are. Specifically, in a Bayesian optimization setting, most choices of acquisition function actively rely on predicted uncertainties to guide the optimization, and well-calibrated uncertainties have been shown to correlate with optimization performance [8] .Our goal with this paper is to start a discussion on the quality of the estimated uncertainties of supervised protein variant effect prediction. Gaussian Processes (GP) are a standard choice for uncertainty quantification due to the closed form expression of the posterior. We therefore first ask the question whether state-of-the-art performance can be obtained within the GP framework. We propose a composite kernel that comfortably achieves this goal, and subsequently investigate the quality of the uncertainty estimates from such a model. Our results show that while standard approaches like reliability diagrams give the impression of good levels of calibration, the quantification of per-instance uncertainties is more challenging. We make our model available as a baseline and encourage the community to place greater emphasis on uncertainty quantification in this important domain. Our contributions can be summarized as follow:• We introduce Kermut, a Gaussian process with a novel composite kernel for modeling mutation similarity, leveraging signals from pretrained sequence and structure models;• We evaluate this model on the comprehensive ProteinGym substitution benchmark and show that it is able to reach state-of-the-art performance in supervised protein variant effect prediction, outperforming recently proposed deep learning methods in this domain;• We provide a thorough calibration analysis and show that while Kermut provides wellcalibrated uncertainties overall, the calibratedness of instance-specific uncertainties remains challenging;• We demonstrate that our model can be trained and evaluated orders of magnitude faster and with better out-of-the-box calibration than competing methods.2 Related work 
protein property prediction: 
 Predicting protein function and properties using machine-learning based approaches continues to be an innovative and important area of research.Recently, unsupervised approaches have gained significant momentum where models trained in a self-supervised fashion have shown impressive results for zero-shot estimates of protein fitness and variant effects relative to a reference protein [3, [9] [10] [11] .Supervised learning is a crucial method of utilizing experimental data to predict protein fitness. This is particularly valuable when the trait of interest correlates poorly with the evolutionary signals that unsupervised models capture during training or if multiple traits are considered. Supervised protein fitness prediction using machine learning has been explored in detail in [12] , where a comprehensive overview can be found. A common strategy is to employ transfer learning via embeddings extracted from self-supervised models [13, 14] , an approach which increasingly relies on large pretrained language models such as ProtTrans [15] , ESM-2 [16] , and SaProt [17] . In [18] , the authors propose to augment a one-hot encoding of the aligned amino acid sequence by concatenating it with a zero-shot score for improved predictions. This was further expanded upon with ProteinNPT [19] , where sequences embedded with the MSA Transformer [20] and zero-shot scores were fed to a transformer architecture for state-of-the-art supervised variant effect prediction with generative capabilities.Considerable progress has been made in defining meaningful and comprehensive benchmarks to reliably measure and compare model performance in both unsupervised and supervised protein fitness prediction settings. The FLIP benchmark [6] introduced three supervised predictions tasks ranging from local to global fitness prediction, where each task in turn was divided into clearly defined splits. The supervised benchmarks often view fitness prediction through a particular lens. Where FLIP targeted problems of interest to protein engineering; TAPE [21] evaluated transfer learning abilities; PEER [22] focused on sequence understanding; ATOM3D [23] considered a structure-based approach; FLOP [24] targeted wild type proteins; and ProteinGym focused exclusively on variant effect prediction [11] . The ProteinGym benchmark was recently expanded to encompass more than 200 standardized datasets in both zero-shot and supervised settings, including substitutions, insertions, deletions, and curated clinical datasets [11, 7] .Kernel methods have seen much use for protein modeling and protein property prediction. Sequencebased string kernels operating directly on the protein amino acid sequences are one such example, where, e.g., matching k-mers at different ks quantify covariance. This has been used with support vector machines to predict protein homology [25, 26] . Another example is sub-sequence string kernels, which in [27] is used in a Gaussian process for a Bayesian optimization procedure. In [28] , string kernels were combined with predicted physicochemical properties to improve accuracy in the prediction of MHC-binding peptides and in protein fold classification. In [29] , a kernel leveraging the tertiary structure for a protein family represented as a residue-residue contact map was used to predict various protein properties such as enzymatic activity and binding affinity. In [30] , Gaussian process regression (GPR) was used to successfully identify promising enzyme sequences which were subsequently synthesized showing increased activity. In [31] , the authors provide a comprehensive study of kernels on biological sequences which includes a thorough review of the literature as well as both theoretical, simulated, and in-silico results.Most similar to our work is mGPfusion [32] , in which a weighted decomposition kernel was defined which operated on the local tertiary protein structure in conjunction with a number of substitution matrices. Simulated stability data for all possible single mutations were obtained via Rosetta [33] , which was then fused with experimental data for accurate ∆∆G predictions of single-and multimutant variants via GPR, thus incorporating both sequence, structure, and a biophysical simulator. In contrast to our approach, the mGPfusion-model does not leverage pretrained models, but instead relies on substitution matrices for its evolutionary signal. A more recent example of kernel-based methods yielding highly competitive results is xGPR [34] , in which Gaussian processes with custom kernels show high performance when trained on protein language model embeddings, similarly to the sequence kernel in our work (see Section 3.3) . Where xGPR introduces a set of novel random feature-approximated kernels with linear-scaling, Kermut instead uses the squared exponential kernel for sequence modeling while additionally modeling local structural environments. The models in xGPR were shown to provide both high accuracy and well-calibrated uncertainty estimation on the FLIP and TAPE benchmarks.Uncertainty quantification (UQ) for protein property prediction continues to be a promising area of research with immediate practical consequences. In [35] , residual networks were used to model both epistemic and aleatoric uncertainty for peptide selection. In [36] , GPR on MLP-residuals from biLSTM embeddings was used to successfully guide in-silico experimental design of kinase binders and fluorescent proteins. The authors of [37] augmented a Bayesian neural network by placing biophysical priors over the mean function by directly using Rosetta energy scores, whereby the model would revert to the biophysical prior when the epistemic uncertainty was large. This was used to predict fluorescence, binding, and solubility for drug-like molecules. In [38] , state-of-the-art performance on protein-protein interactions was achieved by using a spectral-normalized neural Gaussian process [39] with an uncertainty-aware transformer-based architecture working on ESM-2 embeddings.In [40] , a framework for evaluating the epistemic uncertainty of deep learning models using confidence interval-based metrics was introduced, while [41] conducted a thorough analysis of uncertainty quantification methods for molecular property prediction. Here, the importance of supplementing confidence-based calibration with error-based calibration as introduced in [42] was highlighted, whereby the predicted uncertainties are connected directly to the expected error for a more nuanced calibration analysis. We evaluate our model using confidence-based calibration as well as error-based calibration following the guidelines in [41] . In [43] , the authors conducted a systematic comparison of UQ methods on molecular property regression tasks, while [44] investigated calibratedness of regression models for material property prediction. In [45] , the above approaches were expanded to protein property prediction tasks where the FLIP [6] benchmark was examined, while [46] benchmarked a number of UQ methods for molecular representation models. In [47] , the authors developed an active learning approach for partial charge prediction of metal-organic frameworks via Monte Carlo dropout [48] while achieving decent calibration. In [49] , a systematic analysis of protein regression models was conducted where well-calibrated uncertainties were observed for a range of input representations.Much work has been done to solve the inverse-folding problem, where the most probable amino acid sequence to fold into a given protein backbone structure is predicted [50] [51] [52] [53] [54] [55] [56] . Inverse-folding models are trained on large datasets of protein structures and model the physicochemical and evolutionary constraints of sites in a protein conditioned on their structural contexts. These will form the basis of the structural featurization in our work. Local structural environments have previously been used for protein modeling. In [57] , a 3D CNN was used to predict amino acid preferences giving rise to novel substitution matrices. In [58] and [59] , surface-level fingerprinting given structural environments was used to model protein-protein interaction sites and for de novo design of protein interactions. In [60] , chemical microenvironments were used to identify potentially beneficial mutations, while [61] used a similar approach, where they investigated the volume of the local environments and observed that the first contact shell delivered the primary signal thus emphasizing the importance of locality. In [62] , a composition Hellinger distance metric based on the chemical composition of local residue environments was developed and used for a range of structure-related experiments. Recently, local structural environments were used to model mutational preferences for protein engineering tasks [63] , however not in a Gaussian process framework as we propose here. 
preliminaries: 
 We want to predict the outcome of an assay measured on a protein, represented by its amino acid sequence x of length L. We will assume that we have a dataset of N such sequences available, and that these are of equal length and structure, such that we can meaningfully refer to the effect at specific positions (sites) in the protein. In protein engineering, we typically consider modifications relative to an initial wild type sequence, x WT . We will assume that the 3D structure for the initial sequence, s, is available (either experimentally determined or provided by a structure predictor like AlphaFold [64] ). Lastly, for variant x with mutations at sites M ⊆ {1, ..., L}, let x m denote the variant which has the same mutation as x at site m for m ∈ M and otherwise is equal to x WT .To predict protein variant effects, we rely on Gaussian process regression, which we shall now briefly introduce. For a comprehensive overview, see [65] , which this section is based on.Let X and Y be two random variables on the measurable spaces X and R, respectively, and let X = x 1 , ..., x N and y = y 1 , ..., y N be realizations of these random variables. We assume that y i = g(x i ) + ϵ, where g represents some unknown function and ϵ ∼ N (0, σ 2 ϵ ) accounts for random noise. Our objective is to model the distributions capturing our belief about g.Gaussian processes are stochastic processes providing a powerful framework for modeling distributions over functions. The Gaussian process framework allows us to not only make predictions but also to quantify the uncertainty associated with each prediction. A Gaussian process is entirely specified by its mean and covariance functions, m(x) and k(x, x ′ ). We assume that the covariance matrix, K, of the outputs {y 1 , ..., y N } can be parameterized by a function of their inputs {x 1 , ..., x N }. The parameterization is defined by the kernel, k : X × X → R yielding K such that K ij = k(x i , x j ). For k to be a valid kernel, K needs to be symmetric and positive semidefinite.Let f represent our approximation of g, f (x) ≈ g(x). Given a training set D = (X, y) and a number of test points X * , the function f * predicts the values of y * at X * . Using rules of normal distributions, we derive the posterior distribution p(f * |X * , D), providing both a prediction of y * at X * and a confidence measure, often expressed as ±2σ, where σ is the posterior standard deviation at a test point x * .The kernel function often contains hyperparameters, η. These can be optimized by maximizing the marginal likelihood, p(y|X, η), which is known as type II maximum likelihood.It has been shown that local structural dependencies are useful for determining mutation preferences [57, 60, 61, 63] . We therefore hypothesize that constructing a composite kernel with components incorporating information about the local structural environments of residues will be able to model protein variant effects. To this end, we define a structure kernel, k struct , which models mutation similarity given the local environments of mutated sites. A schematic of how the structure kernel models covariances can be seen in Figure 1 . In the following we shall define k 1 struct , a structure kernel that operating on single-mutant variants. Subsequently, we shall extend it to multi-mutant variants resulting in the structure kernel, k struct .We hypothesize that for a given site in a protein, the distribution over amino acids given by a structureconditioned inverse folding model will reflect the effect of a mutation at that site. We consider such an amino acid distribution a representation of the local environment for that site as it reflects the various physicochemical and evolutionary constraints that the site is subject to. We thus presume that two sites with similar local environments will behave similarly if mutated. For instance, mutations at buried sites in the hydrophobic core of the protein will generally correlate more with each other than with surface-level mutations.For single mutant variants we quantify site similarity using the Hellinger kernel kH (x, x ′ ) = exp (-γ 1 d H (f IF (x), f IF (x ′ ))), with γ 1 > 0 [66], where d H is the Hellinger distance (see Appendix B.1). The function f IF : X 1 → [0, 1] 20 takes a single-mutant sequence, x, as input and returns a probability distribution over the 20 naturally occurring amino acids at the mutated site in x given by the inverse folding model. The Hellinger kernel will assign maximum covariance when two sites are identical. This however means that k H is incapable of distinguishing between different mutations at the same site since d H (x, x ′ ) = 1, when x and x ′ are mutated at the same site.To increase flexibility and to allow intra-site comparisons, we introduce a kernel operating on the specific mutation likelihoods. We hypothesize that two variants with mutations on sites that are close in terms of the Hellinger distance will correlate further if the log-probabilities of the specific amino acids on the mutated sites are similar (i.e., the probability of the amino acid that we mutate to is similar at the two sites). We incorporate this by definingk p (x, x ′ ) = k exp (f IF1 (x), f IF1 (x ′ )) = exp(-γ 2 ||f IF1 (x) -f IF1 (x ′ )||), where f IF1 : X 1 → [0, 1] takes a single-mutant sequence, x, as input and returns the log-probability (given by an inverse folding model) of the observed mutation, and where k exp is the exponential kernel.Finally, we hypothesize that the effect of two mutations correlate further if the sites are close in physical space. Hence, we multiply the kernel with an exponential kernel on the Euclidean distance between sites: k d (x, x ′ ) = exp (-γ 3 d e (s i , s j )). Thereby, the closer two sites are physically, the more similar -and thus comparable -their local environments will be.Taking the product of these kernel components, we get the following kernel for single-mutant variants, which assigns high covariance when two single mutant variants have mutations that have similar environments, are physically close, and have similar mutation likelihoods:EQUATIONwhere the kernel has been scaled by a non-negative scalar, λ > 0.In [63] , the authors showed that a simple linear model operating on one-hot encoded mutations is sufficient to accurately predict mutation effects given sufficient data. Thus, we generalize the kernel to multiple mutations by summing over all pairs of sites differing at x and x ′ : The structure kernel models mutations linearly and cannot capture epistatic effects. We propose to add epistatic signals through a sequence kernel. Drawing on the rich literature for modeling protein sequences with Gaussian processes on embeddings [34, 45, 49, 67] , we use a squared exponential kernel which operates on sequence embeddings from a pretrained model. We use the 650M parameter ESM-2 protein language model [16] , and perform mean-pooling across the length dimension, yielding z = f 1 (x), where f 1 produces mean-pooled embeddings, z ∈ R 1280 , of sequence x. We thus model the covariance between these representations asEQUATIONEQUATIONWe choose to add and weigh the structure and sequence kernels resulting in our final kernel formulation, whereby the model can leverage either structure or sequence similarities, depending on the presence and strength of each signal as determined through hyperparameter optimization:EQUATIONAdditional details on both sequence and structure kernels, including a proof of the validity of the structure kernel, implementation details, computational complexity details, and an automatic model selection procedure can be found in Appendices B and C.Kermut can be used with a constant mean function, m(x) = α, where α is a hyperparameter optimized through the marginal likelihood. However, we posit that additional performance can be gained by using an altered mean function which operates on zero-shot fitness estimates, which are often available at relatively low cost: m(x) = αf 0 (x) + β, where f 0 is a zero-shot method evaluated on input sequence x. This is similar to the approach employed in [37] , where Rosetta scores are used as a biophysical prior. We use ESM-2 [16] , which yields the log-likelihood ratio between the variant and wild type residue as in [10] . For details, see Appendix B.While Kermut is based on relatively simple principles, its components are non-trivial in that they exploit learned features from pretrained models: ESM-2 provides protein sequence level embeddings and zero-shot scores, while ProteinMPNN provides a featurization of the local structural environments. We stress that these pretrained components can be readily replaced by other pretrained models. Models that generate (1) protein sequence embeddings, (2) structure-conditioned amino acid distributions, and (3) zero-shot scores are plentiful and such models will progress further in future years. In our work, we have not sought to find the optimal combination of these. Our work should instead be seen  
results: 
 We evaluate Kermut on the 217 substitution DMS assays from the ProteinGym benchmark [7] .The overall benchmark results are an aggregate of three different cross-validation schemes: In the "random" scheme, variants are assigned to one of five folds randomly. In the "modulo" scheme, every fifth position along the protein backbone are assigned to the same fold, and in the "contiguous" scheme, the protein is split into five equal-sized segments along its length, each constituting a fold. For all three schemes, models are trained on four combined partitions and tested on the fifth for a total of five runs per assay, per scheme. The results are processed using the functionality provided in the ProteinGym repository. The average and per-scheme aggregated results can be seen in Table 1 .Our model reaches both higher Spearman correlations and lower mean squared errors than competing methods and thereby achieves state-of-the-art performance in all three schemes, with the largest gains in the challenging modulo and contiguous settings. In Table E .2, the performance per functional category is shown, where we observe a significant performance increase in binding-and stabilityrelated prediction tasks, likely explained by inclusion of the structure kernel. In addition to its high accuracy, Kermut is significantly faster compared to deep learning methods. We provide wall-clock times for running a 5-fold CV loop for a single split-scheme for four select datasets for both Kermut and ProteinNPT in Table C .1. Generating results for all three split schemes for the four datasets thus takes Kermut approximately 10 minutes while ProteinNPT takes upwards of 400 hours.The non-parametric bootstrap standard error for each model relative to Kermut can be seen in Tables E.1 and E.3. In Appendix M, we provide additional results using alternate zero-shot functions.The average and per-split performance for individual assays can be seen in Figures O.1 to O.4 while additional details on computation time can be seen in Appendix C. Lastly, visualizations of the distributions of optimized hyperparameters can be seen in Appendix N.To examine the impact of Kermut's components, we conduct an ablation study, where we ablate its two main kernels from Equation (4) -the structure and sequence kernels -as well as the structure kernel's subcomponents. We similarly investigate the importance of the zero-shot mean function.The ablation study is carried out on all split schemes on a subset of 174 datasets. The difference between the ablation results and the Kermut results can be seen in Table 2 , where larger values indicate large component importance. For the absolute values, see Appendix F, where we additionally include alternative kernel formulations for both structure and sequence kernels as well as for kernel composition.As indicated by the largest drop in performance, the single most important component is the structure kernel. While removing the sequence kernel leads to comparable decreases for all three schemes, The three first elements correspond to the three split-schemes from ProteinGym. The third and fourth correspond to training on both single and double mutants, and testing on each, respectively. For the last column, we train on single and test on double mutants, corresponding to an extrapolation setting.removing the structure kernel primarily leads to drops in the challenging contiguous and modulo schemes. This shows that the structure kernel is crucial for characterizing unseen sites in the wild type protein. While removing the site comparison and mutation probability kernels leads to small and medium drops in performance, we observe that removing both leads to an even larger performance drop, indicating a synergy between the two. In Table E .2, the ablation results per functional category are shown, where observe that the inclusion of the structure kernel is crucial for the increased performance in structure-related prediction tasks such as binding and stability.By inspecting the posterior predictive variance, we can analyze model uncertainty. To this end, we define several mutation domains of interest [49] . We designate the three split schemes from the ProteinGym benchmark as three such domains. These are examples of interpolation, where we both train and test on single mutants (1M→1M). While the main benchmark only considers single mutations, some assays include additional variants with multiple mutations. We consider a number of these and define two additional interpolation domains where we train on both single and double mutations (1M/2M) and test on singles and doubles, respectively. As a challenging sixth domain, we train on single mutations only and test on doubles (1M→2M), constituting an extrapolation domain.For details on the multi-mutant splits, see Appendix G.Figure 2 shows the distributions of mean predictive variances in the six domains. In the three single mutant domains, we observe that the uncertainties increase from scheme to scheme, reflecting the difficulties of the tasks and analogously the expected performance scores (Table 1 ). When training on both single and double mutants (1M/2M), we observe a lower uncertainty on double mutants than single mutants. For many of the multi-mutant datasets, the mutants are not uniformly sampled but often include a fixed single mutation. A possible explanation is thus that it might be more challenging to decouple the signal from a double mutation into its constituent single mutation signals. In the extrapolation setting, we observe large predictive uncertainties, as expected. One explanation of the discrepancy between the variance distributions in the multi-mutant domains might lie in the difference in target distributions between training and test sets. For reference, we include the results for the multi-mutant domains in Table G .1 in the appendix.To clarify the relationship between model uncertainty and expected performance, we proceed with a calibration analysis. First, we perform a confidence interval-based calibration analysis [41] , resulting in calibration curves which in the classification setting are known as reliability diagrams [68] . The results for each dataset are obtained via five-fold cross validation, corresponding to five separately trained models for each split scheme. We select four diverse datasets as examples (Table 3 ), reflecting both high and low predictive performance. The mean calibration curves can be seen in Figure 3a . For method details and results across all datasets, see Appendix J. The mean expected calibration error (ECE) is shown in the bottom of each plot, where a value of zero indicates perfect calibration. Overall, the uncertainties appear to be well-calibrated both qualitatively from the curves and quantitatively from the ECEs. Even the smallest dataset (fourth row, N = 165) achieves decent calibration, albeit with larger variances between folds.While the confidence interval-based calibration curves show that we can trust the uncertainty estimates overall, they do not indicate whether we can trust individual predictions. We therefore supplement the above analysis with an error-based calibration analysis [42] , where a well-calibrated model will have low uncertainty when the error is low. The calibration curves can be seen in Figure 3b . We compute the per CV-fold expected normalized calibration error (ENCE) and the coefficient of variation (c v ), which quantifies the variance of the predicted uncertainties. Ideally, the ENCE should be zero while the coefficient of variation should be relatively large (indicating spread-out uncertainties).While the confidence interval-based analysis suggested that the uncertainty estimates are wellcalibrated with some under-confidence, the same is not as visibly clear for the error-based calibration plots, suggesting that the expected correlation between model uncertainty and prediction error is not a given. We do however see an overall trend of increasing error with increasing uncertainty in three of the four datasets, where the curves lie close to the diagonal (as indicated by the dashed line).The second row shows poorer calibration -particularly in the modulo and contiguous schemes. The curves however remain difficult to interpret, in part due to the errorbars on both axes. To alleviate this, we compute similar metrics for all 217 datasets across the three splits, which indicate that, though varying, most calibration curves are well-behaved (see Appendix J.3).We supplement the above calibration curves with Figures K.1 to K.4 in the Appendix, which show the true values plotted against the predictions. These highlight the importance of well-calibrated uncertainties and underline their role in interpreting model predictions and their trustworthiness.Monte Carlo (MC) dropout [48] is a popular uncertainty quantification technique for deep learning models. Calibration curves for ProteinNPT with MC dropout for the same four datasets across the three split schemes can be seen in Appendix L.2 while figures showing the true values plotted against the predictions with uncertainties are shown in Appendix L. 3 . These indicate that employing MC dropout on a deep learning model like ProteinNPT seems to provide lower levels of calibration, providing overconfident uncertainties across assays and splits. Due to the generally low uncertainties and the resulting difference in scales, the calibration curves are often far from the diagonal. The trends in the calibration curves however show that the model errors often correlate with the uncertainties, suggesting that the model can be recalibrated to achieve decent calibration.Other techniques for uncertain quantification in deep learning models certainly exist, and we by no means rule out that other techniques can outperform our method (see Discussion below). We note, however, that many uncertainty quantification methods will be associated with considerable computational overhead compared to the built-in capabilities of a Gaussian process. 
discussion: 
 We have shown that a carefully constructed Gaussian process is able to reach state-of-the-art performance for supervised protein variant effect prediction while providing reasonably well-calibrated uncertainty estimates. For a majority of datasets, this is achieved orders of magnitude faster than competing methods.While the predictive performance on the substitution benchmark is an improvement over previous methods, our proposed model has its limitations. Due to the site-comparison mechanism, our model is unable to handle insertions and deletions as it only operates on a fixed structure. Additionally, as the number of mutations increases, the assumption of a fixed structure might worsen, depending on the introduced mutations, which can affect reliability as the local environments might change. An additional limitation is the GP's O(N 3 ) scaling with dataset size. While not a major obstacle in the single mutant setting, dataset sizes can quickly grow when handling multiple mutants. The last decades have however produced a substantial literature on algorithms for scaling GPs to larger datasets [72] [73] [74] , which could alleviate the issue, and we therefore believe this to be a technical rather than fundamental limitation. An additional limitation might present itself it in the multi-mutant setting, where the lack of explicit modeling of epistasis can potentially hinder extrapolation to higher-order mutants, prompting further investigation.Well-calibrated uncertainties are crucial for protein engineering; both when relying on a Bayesian optimization routine to guide experimental design using uncertainty-dependent acquisition functions and similarly to weigh the risk versus reward for experimentally synthesizing suggested variants. We therefore encourage the community to place a greater emphasis on uncertainty quantification and calibration for protein prediction models as this will have measurable impacts in real-life applications like protein engineering -perhaps more so than increased prediction accuracy. We hope that Kermut can serve as a fruitful step in this direction.The codebase is publicly available at https://github.com/petergroth/kermut under the open source MIT License.The structure kernel is comprised of three components, each increasing model flexibility. The sitecomparison kernel, k H , compares site-specific, structure-conditioned amino acid distributions. Given two such discrete probability distributions, p := f IF (x) and q := f IF (x ′ ), their distance is quantified via the Hellinger distance [75] dH (p, q) = 1 √ 2 20 i=1 ( √ p i - √ q i ) 2 ,which is used in the Hellinger kernel [66] kH (x, x ′ ) = exp (-γ 1 d H (p, q)) = exp   -γ 1 1 √ 2 20 i=1 ( √ p i - √ q i ) 2   .For the mutation probability kernel, k p , we use the log-probabilities rather than the amino acid identities to reflect that amino acids with similar probability on sites with similar distributions should have similar biochemical effects on the protein. We do not include the log-probabilities of the wild type amino acids, as the inverse folding model by definition is trained to assign high probabilities to the wild type sequence. The exact probability of a wild type amino acid depends on how many other amino acids that are likely to be at the given site. For example, we would expect a probability close to one for a functionally critical amino acid at a particular site. Conversely, for a less critical surfacelevel residue requiring, e.g., a polar uncharged amino acid, we would expect similar probabilities for the four amino acids of this type. Thus, variations of log-probabilities of the wild type amino acids should be reflected in the distribution on the sites captured by the Hellinger kernel.For the per-residue amino acid distributions in k H and k p , we use ProteinMPNN [52] . ProteinMPNN relies on a random decoding order and thus benefits from multiple samples. We decode the wild type amino acid sequence a total of ten times while conditioning on the full structure and the sequence that has been decoded thus far. We then compute a per-residue average distribution. We use the v_48_020 weights.For the distance kernel, k d , we calculate the Euclidean distance between α carbon atoms, where the unit is in Ångstrøm. All wild type structures used in k d and by ProteinMPNN are predicted via AlphaFold2 [64] and are provided by ProteinGym [7] .The sum over all pairs of mutations in equation 2 is motivated by [63] who showed that a linear model was sufficient to effectively model the mutation effects. Note that Cov(EQUATION). Hence, if the Y i 's are the effects of single mutations, we can find the covariance between two double mutant variants by calculating the pairwise sum and similarly for other number of mutations.For the sequence kernel, we use embeddings extracted from the ESM-2 protein language model [16] .We use the esm2_t36_650M_UR50D model with 650M parameters. The embeddings are mean-pooled across the sequence dimension such that each variant is represented by a 1280 dimensional vector.While it has been shown that other aggregation method can lead to large increases in performance [76] , we consider alternate methods such as training a bottleneck model out of the scope of this paper.An alternative formulation of the kernel, where we omit λ from Equation (1) and π from Equation (4), and instead provide the structure and sequence kernels with separate coefficients have shown to provide close to identical results when using a smoothed box prior (constrained between 0 and 1). While this approach is somewhat more elegant, this does not justify the re-computation of all results.For the zero-shot mean function, we download and use the pre-computed zero-shot scores from ProteinGym at https://github.com/OATML-Markslab/ProteinGym. The zero-shot value is calculated as the log-likelihood ratio between the variant and the wild type at the mutated residue as described in [10] :f 0 (x) = i∈M log p(x i ) -log p(x WT i )The values can be computed straightforwardly using the ESM suite (using the masked-marginals strategy). For multi-mutants, the sum of the ratios is taken.We will in the following argue why Kermut's structure kernel is a valid kernel. Recall that a function k : X × X → R is a kernel if and only if the matrix K, where K ij = k(x i , x j ), is symmetric positive semi-definite [65] . From literature we know a number of kernels and certain ways these can be combined to create new kernels. We will argue that Kermut is a kernel by showing how it is composed of known kernels, combined using valid methods.Note that, if we have a mapping f : X → Z and a kernel k Z , on Z, then k X defined by kX (x, x ′ ) := k Z (f (x), f (x ′ )) is a kernel on X .Let X be the space of sequences parameterized with respect to a reference sequence. Let f 1 : X → Z be a transformation of the sequences defined as in Section 3.3. k seq is the squared exponential kernel on the transformed variants. Hence, k seq is is a kernel on X .Let X 1 ⊂ X denote the subspace of single mutant variants and f 2 : X 1 → R 3 be a function mapping single mutant variants into the 3D coordinates of the α-carbon of the particular mutation. k d is the exponential kernel on this transformed space. Thus k d is a kernel on X 1 .Let f IF : X 1 → G ⊆ [0, 1] 20 be defined as in Section 3.3, where G is the space of probability distributions over the 20 amino acids. f IF (x) is the probability distribution over the mutated site of x given by an inverse folding model. k H is the Hellinger kernel on the single mutation variants transformed by f IF , hence, a valid kernel on X 1 . Likewise k p is the exponential kernel of a transformation, f IF1 : X 1 → [0, 1] as defined in Section 3.3, of the sequences, hence also a kernel.Scaling, multiplying, and adding kernels result in new kernels, making k 1 struct a valid kernel for single mutations [65] . We need to show that k struct and thereby k is valid for any number of mutations.Let f 4 : X → B be a function taking a variant x with M mutations and mapping it to a set b = {x m } m∈M of all the single mutations which constitutes x. Define the set kernel [77] k set (b, b ′ ) := x m ∈b,x ′m ∈b ′ λk 1 struct (x m , x ′m )k struct is the set kernel on the transformed input and thus also a kernel. We have thereby shown that Kermut is a kernel for variants with any number of mutations.Kermut has been developed from the ground up as described Section 3.3, which led to separate structure and sequence kernels which are added together. Alternatively, an automatic model selection scheme can be employed for data-driven kernel composition as proposed in Chapter 3 in [78] . For demonstrative purposes, we conduct such a model selection procedure. We define four base kernels (k H , k p , k d , k seq ) as well as sum and product operations. We choose a subset of 17 ProteinGym assays and fit a GP (with a zero-shot mean function) and each of the four kernels, where the best performing kernel across splits is kept. For the second round, the remaining three kernels are either added or multiplied to the existing kernel. This process is continued until either all four kernels are used or until the test performance no longer increases. The results can be seen in Table B .1, where the final kernel is the product of the four base kernels. The results on the 174 ablation datasets using this kernel (denoted as "Kermut (product)") can be seen with the main Kermut GP in Table F We build our kernel using the GPyTorch framework [79] . We assume a homoschedastic Gaussian noise model, on which we place a HalfCauchy prior [80] with scale 0.1. We fit the hyperparameters by maximizing the exact marginal likelihood with gradient descent using the AdamW optimizer [81] with learning rate 0.1 for a 150 steps, which proved to be sufficient for convergence for a number of sampled datasets.All experiments are performed on a Linux-based cluster running Ubuntu 20.04.4 LTS, with a AMD EPYC 7642 48-Core Processor with 192 threads and 1TB RAM. NVIDIA A40s were used for GPU acceleration both for fitting the Gaussian processes and for generating the protein embeddings.There are multiple factors to consider when evaluating the training time of Kermut. A limitation of using a Gaussian process framework is the cubic scaling with dataset size. For this reason, we ran the ablation study on only 174 of the 217 datasets.Training Kermut on these datasets for a single split-scheme using the aforementioned hardware (single GPU) takes approximately 1 hour and 30 minutes. Getting ablation results for all three schemes thus takes between 4-5 hours. This however assumes that 1. the embeddings for the sequence kernel have been precomputed, 2. the probability distribution for all sites in the wild type protein have been precomputed, 3. and that the zero-shot scores have been precomputed.We do not see any of these as major limitations as the same applies to ProteinNPT and similar models.Scaling the experiments to the full ProteinGym benchmark is however costly. We were able to train/evaluate Kermut on 215/217 datasets using an NVIDIA A40 GPU with 48GB VRAM. The remaining two, POLG_CXB3N_Mattenberger_2021 and POLG_DEN26_Suphatrakul_2023, datasets were too large to fit into GPU memory without resorting to reduced precision. For these, we trained and evaluated the model using CPU only which takes considerable time.The structures used for obtaining the site-wise probability distributions are predicted by Al-phaFold2 [64] and are provided in the ProteinGym repository [7] . The provided structures for A0A140D2T1_ZIKV and POLG_HCVJF do not contain the full structures however, but only a localized area where the mutations occur due to the long sequence lengths. Since our model only operates on sites with mutations, this is not an issue for neither the inverse-folding probability distributions nor the inter-residue distances.For BRCA2_HUMAN, three PDB files are provided due to the long wild type sequence length. We use ProteinMPNN to obtain the distributions at all sites in each PDB file and stitch them together in a preprocessing step. Calculating the inter-residue distances is however non-trivial and would require a careful alignment of the three structures. Instead, we drop the distance term, k d , in the kernel for the BRCA2_HUMAN_Erwood_2022_HEK293T dataset (equivalently setting it to one:k d (x, x ′ ) = 1).The sequence kernel operates on ESM-2 embeddings. The ESM-2 model has a maximum sequence length of 1022 amino acids. Protein sequences that are longer than this limit are truncated.The wall-clock times for generating the results used for the calibration curves in Figures 3a and 3b and Appendix L.2 for one split scheme can be seen in Table C .1 for Kermut and ProteinNPT. The experiments were carried out using identical hardware. The test system is however a shared compute cluster with sharded GPUs, so variance is expected between runs. Generating the full results for the figures thus takes Kermut approximately 10 minutes while ProteinNPT takes 400 hours. This shows the significant reduction in computational burden that Kermut allows for. Both ProteinNPT and Kermut assumes that sequence embeddings are available a priori. Evaluating the kernel for two variants x and x ′ , involves computation of the two components k seq and k struct :For each variant k seq requires a forward pass through ESM-2, which is based on the transformer architecture and has quadratic scaling with respect to the sequence length. Given the ESM-2 embeddings, the computational complexity is constant in sequence length and number of mutations.k struct requires a single forward pass through ProteinMPNN for the wild type protein. Given the output of ProteinPMNN, the computational complexity for evaluation of the kernel is m 1 × m 2 for two proteins with m 1 and m 2 number of mutations. The computational complexity of k 1 struct is constant in sequence length and number of mutations.All data and evaluation software is accessed via the ProteinGym [7] repository at https://github. com/OATML-Markslab/ProteinGym which is under the MIT License.The ProteinGym suite provides an aggregation procedure, whereby the predictive performance across both cross-validation schemes and functional categories can be gauged. We provide these results in Tables E.1 to E.4. We mirror the label normalization from [19] and [7] , where, for each fold of cross-validation, train and test labels are normalized given the mean and standard deviation of the training labels. The Spearman correlation coefficient and MSE are then computed in the normalized space across folds, leading to a single Spearman coefficient and MSE per assay per split (i.e., not as an average of metrics per CV-fold). The modulo and contiguous splits in ProteinGym were updated in April 2024. For completeness, we here provide results using Kermut on the updated splits for all 217 DMS assays. These can be seen in Table E . 5 , where we observe a slight decrease in performance for the contiguous and modulo splits, compared to the reference results in Table 1 . 69 of the datasets from the ProteinGym benchmark include multi-mutants. In addition to the random, modulo, and contiguous split, these also have a fold_rand_multiples split. We here show the results for Kermut in this setting. Of the 69 datasets, we select 52 which (due to the cubic scaling of fitting GPs) include fewer than 7500 variant sequences. We additionally ignore the GCN4_YEAST_Staller_2018 dataset which has a very large number of mutations. This leads to a total of 51 datasets. All results where the training domain is "1M/2M→" are from models trained using the above split. The "1M→" domain results correspond to training the model once on single mutants and evaluating it on double mutants. In addition to Kermut, we include as a baseline results from a GP using the sequence kernel operating on mean-pooled ESM-2 embeddings. This is equivalent to setting the structure kernel to 0, k struct = 0 as in Table 2 . For results on the multi-mutant GB1 landscape from FLIP [6] , see Appendix H. To further investigate Kermut's performance in a multi-mutant setting, we apply it to the GB1 fitness landscape from FLIP [6] . The results can be seen in Table H .1. Kermut's base configuration severely underperforms in the 1-vs-rest split, while reaching similar correlation coefficients in the 2-vs-rest and 3-vs-rest splits. A reason for the initial low score might be the that the accuracy of zero-shot methods at different mutation orders tend to decrease with higher mutation count. Despite this, we still see low performance compared to other models when removing the zero-shot mean function.The GB1 landscape is comprised of 149,361 mutations at exactly four highly epistatic positions in the GB1 binding domain of Protein G. This is a challenging task for Kermut, whose structural kernel directly compares sites. With only four sites to compare, Kermut fails to accurately model the fitness landscape. Additionally, as described in the main text, the only epistatic modeling in Kermut is via the mean-pooled protein language model embeddings in the sequence kernel, which proves insufficient to capture the interplay of these highly epistatic sites. Further experimentation is required to thoroughly gauge Kermut's performance across diverse multi-mutant assays where the number of mutated residues is variable and epistatis plays a central role. Given a collection of mean predictions and uncertainties, we wish to gauge how well-calibrated the uncertainties are. The posterior predictive mean and variance for each data point is interpreted as a Gaussian distribution and symmetric intervals of varying confidence are placed on each prediction [41] . In a well-calibrated model, approximately x % of predictions should lie within a x % confidence interval, e.g., 50% of observations should fall in the 50% confidence interval. The confidence intervals are discretized into K bins and the fraction of predictions falling within in bin is calculated. The calibration curve then plots the confidence intervals vs. the fractions, whereby a diagonal line corresponds to perfect calibration. Given the fractions and confidence intervals, the expected calibration error (ECE) is calculated asECE = 1 K K i=1 |acc(i) -i|,where K is the number of bins, i indicates the equally spaced confidence intervals, and acc(i) is the fraction of predictions falling within the ith confidence interval.An alternative method of gauging calibratedness is error-based calibration where the prediction error is tied directly to predictions [41, 42] . The predictions are sorted according to their predictive uncertainty and placed into K bins. For each bin, the root mean square error (RMSE) and root mean variance (RMV) is computed. In error-based calibration, a well-calibrated model as equal RMSE and RMV, i.e., a diagonal line. The x and y values in the resulting calibration plot are however not normalized from 0 to 1 as in confidence interval-based calibration. The expected normalized calibration error (ENCE) can be computed asENCE = 1 K K i=1 |RMV(i) -RMSE(i)| RMV(i) .Additionally, we compute the coefficient of variation (c v ) asc v = N n=1 (σn-µσ) 2 N -1 µ σ ,where µ σ = 1 N N n=1 σ n , and where n indexes the N data points [42] .To quantitatively describe the calibratedness of Kermut across all datasets, we compute the above calibration metrics for all split schemes and folds. We do this for Kermut and a baseline GP using the sequence kernel on ESM-2 embeddings (equivalent to the sequence kernel from Equation (3) with a constant mean). These values can be seen in Figure J .1 for each of the three main split-schemes, the 1M/2M→1M/2M split ("Multiples"), and the 1M→2M split ("Extrapolation"). For the three main splits, we see that the calibratedness measured by the ENCE correlates with performance, where the lowest values are seen in the random setting. Generally, we see that the inclusion of the structure kernel improves not only performance (see Table 2 ) but also calibration, as the ECE and ENCE values are consistently better for Kermut, with the exception of ENCE in the extrapolation domain. We however see that the sequence kernel (squared exponential) consistently provide predictive variances that themselves vary more, which is generally preferable.Overall, we can conclude that Kermut appears to be well-calibrated both qualitatively and quantitatively. While the ECE values are generally small and similar, the ENCE values suggest a more nuanced calibration landscape, where we can expect low errors when our model predicts low uncertainties, particularly in the random scheme.For each dataset, split, and fold, we perform a linear regression to the error-based calibration curves and summarize the slope and intercept. As perfect calibration corresponds to a diagonal line, we want the distribution over the slopes to be centred on one and the distribution over intercepts to be centred on zero. Boxplots for this analysis can be seen in Appendix J. We use ProteinNPT using the provided software in the paper with the default settings. We generate the MSA Transformer embeddings manually using the provided software from ProteinNPT. During evaluation, we predict using Monte Carlo dropout (with 25 samples, as described in the ProteinNPT appendix). An uncertainty estimate per test sequence is obtained by taking the standard deviation over the 25 samples as prescribed. Kermut uses a linear transformation of a variant's zero-shot score as its mean function. In the main results ESM-2 was used. We here provide additional results where different zero-shot methods are used. The experiments are carried out as the ablation results in Section 4.1, i.e., on 174/217 datasets. All zero-shot scores are pre-computed and are available via the ProteinGym suite.Using a zero-shot mean function instead of a constant mean evidently leads to increased performance. The magnitude of the improvement depends on the chosen zero-shot method, where the order roughly corresponds to that of the zero-shot scores in ProteinGym. We do however see that opting for EVE yields the largest performance increase. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [No] Justification: We have not described safeguards due to the low risk of misuse, while we however describe potential misuse in Appendix P.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: Credit is given to the benchmark data and software in Appendix D.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix:  


Decision: There was broad agreement among the reviews that this is a well-written, interesting paper with significant and novel contributions. I fully agree with this positive evaluation, and therefore I recommend acceptance of this work.

Reviews: review: paper_summary: paper_summary: This paper proposes a kernel regression model to predict protein mutational effects. The model includes kernel functions crafted for the task. In specific, it includes a kernel that measures sequence similarity based on ESM-2 features, a local structure similarity kernel based on ProteinMPNN probability, and other kernels that impose priors such as spatial correlation between mutation sites.
strength_weakness: strength_weakness: - Current formulation of Kermut does not provide transferability to different protein sequences, while previous zero-shot prediction methods are capable of predicting variant effects without prior experimental data on the same sequence.
questions: questions: - Why is the structure kernel function based on amino acid probability rather than the last-layer feature of ProteinMPNN? And why is the sequence kernel based on last-layer features rather than probability? What is the consideration behind these choices?
- Why is it reasonable to assume that mutations on the same site have similar effects (L135)? It seems that the effects depend more on the amino acid types.
limitations: limitations: See weakness section.
paper_summary: paper_summary: The authors suggest a method to predict the effect of mutations given sparse data. Their method is based on identifying the similarity of different sites on a protein by embeddings from large language models and structure. They show that their method performs state of the art mutation effect prediction. They also show that their method is only slightly overconfident.
strength_weakness: strength_weakness: Epistasis is only included through sequence embeddings; in particular, the impact of structure is purely linear.

Despite not suggesting any radically new technique, this is a practical method that cleverly and effectively uses available tools. For this reason however, I think it is reasonable to expect that the authors try hard to build as strong a method as possible. In particular, this paper is missing a more thorough investigation of model choices -- what if the Hellinger kernel is replaced with something else? what if a different kernel is used to compare embeddings? what if equation 2 is replaced with a kernel with a term for every combination the 3 kernels? what if the kernel is meta-learned on a subset ProteinGym and applied to the rest? The paper would be substantially strengthened by applying the methodology of Duvenaud, David. 2014. “Automatic Model Construction with Gaussian Processes.” University of Cambridge.
questions: questions: Could you more thoroughly investigate building an accurate kernel method on this data? For example, trying nonlinear structure kernels might help? Using something like $\tilde k(X, X')=\exp(-k(X, X'))$ or an IMQ formulation as in Amin, Alan Nawzad, Eli Nathan Weinstein, and Debora Susan Marks. 2023. “Biological Sequence Kernels with Guaranteed Flexibility.” arXiv [Stat.ML]. arXiv. http://arxiv.org/abs/2304.03775. might be a reasonable thing to try. Or more flexible combinations of kernels as described in the weaknesses section.
limitations: limitations: Partially addressed. I would like a longer discussion about epistasis mentioned in the weaknesses section.
paper_summary: paper_summary: The authors introduce a family of Gaussian process based regression models for protein variant effect prediction. The "composite" kernel introduced makes use of structural information (i.e. closeness in 3d space) as well as pre-trained sequence and/or structure models like ESM2 and ProteinMPNN (via embeddings and/or amino acid preference distributions). The model shows promising predictive performance on the ProteinGym benchmark, including w.r.t. uncertainty metrics like expected calibration error (ECE).
strength_weakness: strength_weakness: - The description and discussion of the Kermut kernel is not very easy to follow and could be considerably improved. 
  - For example, much of the discussion on lines 128-144 seems a bit besides the point, since the authors do not in the end develop a model that is linear in one-hot-features. 
  - The authors should do a much better job of foreshadowing/signposting that their construction "for single mutant variants" is but a stepping stone to a multi-mutant construction; otherwise the reader is liable to get confused about what's going on. 
  -  Many of the choices that lead to the final kernel construction are either not motivated or only briefly discussed. Why not use, for example, a product kernel $k_{\rm struct} \times k_{\rm seq}$? Some of these alternative choices should be discussed and, ideally, included in ablations. 
  - What does this mean? [Line 167] "preventing the comparison of different mutations at the same site" 
- The paper is missing a discussion of the computational complexity of computing the kernel w.r.t. the sequence length, the number of mutations from the wildtype, etc.
- The discussion of prior work is inadequate, especially w.r.t. work on sequence kernels and previous applications of GPs to protein modeling. Granted there is only so much space in the main text, but the reader deserves a more detailed discussion. (Perhaps some of the discussion can be relegated to the appendix.) To name just a few examples, there is a lot of work on sequence kernels (see e.g. [A] and references therein) and it is negligent to ignore this body of work. Also the method comparison to mGPfusion [24] is inadequate. The authors did not invent sequence kernels or pioneer their application to protein modeling (e.g. [B, C, D] for more recent work) and should be much more liberal and informative in their discussion of the literature. Granted the relevant literature can be scattered (arxiv, biorxiv, etc.), but the readers deserve (much) better. Discussion of relevant work is more than a required chore: if well done it adds significant value to the reader and the literature as a whole.

References:
- [A] "Biological Sequence Kernels with Guaranteed Flexibility," Alan Nawzad Amin, Eli Nathan Weinstein, Debora Susan Marks, 2023.
- [B] Moss, Henry, et al. "Boss: Bayesian optimization over string spaces." Advances in neural information processing systems 33 (2020): 15476-15486.
- [C] Parkinson, Jonathan, and Wei Wang. "Scalable Gaussian process regression enables accurate prediction of protein and small molecule properties with uncertainty quantitation." arXiv preprint arXiv:2302.03294 (2023).
- [D] Greenman, Kevin P., Ava P. Amini, and Kevin K. Yang. "Benchmarking uncertainty quantification for protein engineering." bioRxiv (2023): 2023-04.
questions: questions: - Have the authors considered pooling strategies other than mean-pooling?
- What are typical kernel hyperparameters learned across the ProteinGym benchmark? $\gamma_1$, $\pi$, etc. It would be great to see a histogram or similar.
- The throwaway comment [Line 316] "While this limitation can be fixed by online protein structure prediction, the computational cost would increase significantly" seems a bit misleading without further clarification, since kernels are "pairwise" objects and so it's not immediately clear how you would incorporate multiple structures into your kernel. Can you please comment?
limitations: limitations: To my mind the major limitation of this work is that reliance on the ProteinGym benchmark means that there is little signal on how well this class of models would perform in more challenging problem settings where the model is asked to extrapolate out many mutations away from the wild type. There are increasingly many public datasets that make this kind of evaluation possible in principle, and I encourage the authors to apply their method to more challenging scenarios.

For example:
- Chinery, Lewis, et al. "Baselining the Buzz. Trastuzumab-HER2 Affinity, and Beyond!." bioRxiv (2024): 2024-03.

Comments: title: Response
References used in section 2.3
References used in section 2.2
New section 2.3
Revised section 2.2
comment: Thank you for you quick response - and the raised score!
We agree that epistasis most likely plays a significant role when multiple mutations are introduced. We will make sure to include additional results in the final paper.
I thank the authors for engaging in such depth with reviewer feedback. Accordingly I have raised my score to a 7. I look forward to reading the camera ready version of the paper.
I appreciate you've included a longer discussion of epistasis. I appreciate the results in CoVES and I think you're right that including epistasis is unlikely to substantially increase values on the tested benchmarks. However, I think drawing the conclusion from CoVES that epistasis doesn't matter is not sound; rather I think that this is a result of their benchmark being "hackable". Ultimately, I similarly hope you endeavor to do more than fit single and double mutants in ProteinGym: the sorts of kernels you design could be useful in fitting more diverse data (say from an iterative design experiment) or predicting predicting epistasis as its own end -- in these cases nonlinearity could be more useful.

It seems you've done a lot of experiments. I would suggest dumping as many of those results into the appendix as is convenient for you, even if they don't result in major differences.

Your additions address my main concerns about the potential to improve the method and I'm happy to recommend the paper be accepted.
[39] Fredrik K. Gustafsson, Martin Danelljan, and Thomas B. Schön. Evaluating Scalable Bayesian Deep Learning Methods for Robust Computer Vision, April 2020.

[40] Gabriele Scalia, Colin A. Grambow, Barbara Pernici, Yi-Pei Li, and William H. Green. Evaluating Scalable Uncertainty Estimation Methods for Deep Learning-Based Molecular Property Prediction. ACS, 2020.

[41] Dan Levi, Liran Gispan, Niv Giladi, and Ethan Fetaya. Evaluating and Calibrating Uncertainty Prediction in Regression Tasks, February 2020.

[42] Lior Hirschfeld, Kyle Swanson, Kevin Yang, Regina Barzilay, and Connor W Coley. Uncertainty quantification using neural networks for molecular property prediction. Journal of Chemical Information and Modeling, 60(8):3770–3780, 2020.

[43] Kevin Tran, Willie Neiswanger, Junwoong Yoon, Qingyang Zhang, Eric Xing, and Zachary W Ulissi. Methods for comparing uncertainty quantifications for material property predictions. Machine Learning: Science and Technology, 1(2):025006, 2020.

[44] Kevin P Greenman, Ava P Amini, and Kevin K Yang. Benchmarking uncertainty quantification for protein engineering. bioRxiv, pages 2023–04, 2023.

[45] Yinghao Li, Lingkai Kong, Yuanqi Du, Yue Yu, Yuchen Zhuang, Wenhao Mu, and Chao Zhang. Muben: Benchmarking the uncertainty of molecular representation models. Transactions on Machine Learning Research, 2023.

[46] Stephan Thaler, Felix Mayr, Siby Thomas, Alessio Gagliardi, and Julija Zavadlav. Active
learning graph neural networks for partial charge prediction of metal-organic frameworks via
dropout monte carlo. npj Computational Materials, 10(1):86, 2024.

[47] Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing
model uncertainty in deep learning. In Maria Florina Balcan and Kilian Q. Weinberger,
editors, Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1050–1059, New York, New York, USA, 20–22 Jun 2016. PMLR.

[48] Richard Michael, Jacob Kæstel-Hansen, Peter Mørch Groth, Simon Bartels, Jesper Salomon, Pengfei Tian, Nikos S. Hatzakis, and Wouter Boomsma. A systematic analysis of regression models for protein engineering. PLOS Computational Biology, 20(5):e1012061, May 2024.
[5] Christian Dallago, Jody Mou, Kadina E. Johnston, Bruce J. Wittmann, Nicholas Bhattacharya, Samuel Goldman, Ali Madani, and Kevin K. Yang. FLIP: Benchmark tasks in fitness landscape inference for proteins, January 2022

[24] Christina Leslie, Eleazar Eskin, and William Stafford Noble. The spectrum kernel: A string kernel for svm protein classification. In Biocomputing 2002, pages 564–575. WORLD SCIENTIFIC, December 2001. 

[25] Christina S. Leslie, Eleazar Eskin, Adiel Cohen, Jason Weston, and William Stafford Noble.Mismatch string kernels for discriminative protein classification. Bioinformatics, 20(4):467–476, March 2004.

[26] Henry Moss, David Leslie, Daniel Beck, Javier Gonzalez, and Paul Rayson. Boss: Bayesian optimization over string spaces. Advances in neural information processing systems, 33:15476–15486, 2020.

[27] Nora C Toussaint, Christian Widmer, Oliver Kohlbacher, and Gunnar Rätsch. Exploiting
physico-chemical properties in string kernels. BMC bioinformatics, 11:1–9, 2010.

[28] Philip A. Romero, Andreas Krause, and Frances H. Arnold. Navigating the protein fitness landscape with Gaussian processes. Proceedings of the National Academy of Sciences, 110(3):E193–E201, January 2013.

[29] Jonathan C Greenhalgh, Sarah A Fahlberg, Brian F Pfleger, and Philip A Romero. Machine learning-guided acyl-acp reductase engineering for improved in vivo fatty alcohol production. Nature communications, 12(1):5825, 2021.

[30] Alan Nawzad Amin, Eli Nathan Weinstein, and Debora Susan Marks. Biological sequence kernels with guaranteed flexibility. arXiv preprint arXiv:2304.03775, 2023.

[31] Emmi Jokinen, Markus Heinonen, and Harri Lähdesmäki. mGPfusion: Predicting protein
stability changes with Gaussian process kernel learning and data fusion. Bioinformatics,
34(13):i274–i283, July 2018.

[32] Andrew Leaver-Fay, Michael Tyka, Steven M Lewis, Oliver F Lange, James Thompson, Ron Jacak, Kristian W Kaufman, P Douglas Renfrew, Colin A Smith, Will Sheffler, et al. Rosetta3: an object-oriented software suite for the simulation and design of macromolecules. In Methods in enzymology, volume 487, pages 545–574. Elsevier, 2011.

[33] Jonathan Parkinson and Wei Wang. Linear-scaling kernels for protein sequences and small molecules outperform deep learning while providing uncertainty quantitation and improved interpretability. Journal of Chemical Information and Modeling, 63(15):4589–4601, 2023.

[34] Haoyang Zeng and David K Gifford. Quantification of uncertainty in peptide-mhc binding
prediction improves high-affinity peptide selection for therapeutic design. Cell systems, 9(2):159–166, 2019.

[35] Brian Hie, Bryan D. Bryson, and Bonnie Berger. Leveraging Uncertainty in Machine Learning Accelerates Biological Discovery and Design. Cell Systems, 11(5):461–477.e9, November 2020.

[36] Hunter Nisonoff, Yixin Wang, and Jennifer Listgarten. Coherent blending of biophysics-based knowledge with bayesian neural networks for robust protein property prediction. ACS Synthetic Biology, 12(11):3242–3251, 2023.

[37] Young Su Ko, Jonathan Parkinson, Cong Liu, and Wei Wang. Tuna: An uncertainty aware transformer model for sequence-based protein-protein interaction prediction. bioRxiv, pages 2024–02, 2024.

[38] Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshminarayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. Advances in neural information processing systems, 33:7498–7512, 2020.
**2.3 Uncertainty quantification and calibration**

Uncertainty quantification (UQ) for protein property prediction continues to be a promising area of research with immediate practical consequences. In [34], residual networks were used to model both epistemic and aleatoric uncertainty for peptide selection. In [35], GPR on MLP-residuals from biLSTM embeddings was used to successfully guide in-silico experimental design of kinase binders and protein fluorescence, amongst others. The authors of [36] augmented a Bayesian neural network by placing biophysical priors over the mean function by directly using Rosetta energy scores, whereby the model would revert to the biophysical prior when the epistemic uncertainty was large. This was used to predict fluorescence, binding and solubility for drug-like molecules. In [37], state-of-the-art performance on protein-protein interactions was achieved by using a spectral-normalized neural Gaussian process [38] with an uncertainty-aware transformer-based architecture working on ESM-2 embeddings.

In [39], a framework for evaluating the epistemic uncertainty of deep learning models using confidence interval-based metrics was introduced, while [40] conducted a thorough analysis of uncertainty quantification methods for molecular property prediction. Here, they highlighted the importance of supplementing confidence-based calibration with error-based calibration as introduced in [41], whereby the predicted uncertainties are connected directly to the expected error for a more nuanced calibration analysis. We evaluate our model using confidence-based calibration as well as error-based calibration following the guidelines in [40]. In [42], the authors conducted a systematic comparison of UQ methods on molecular property regression tasks, while [43] investigated calibratedness of regression models for material property prediction. In [44], the above approaches were expanded to protein property prediction tasks where the FLIP [5] benchmark was examined, while [45] benchmarked a number of UQ methods for molecular representation models. In [46], the authors developed an active learning approach for partial charge prediction of metal-organic frameworks via Monte Carlo dropout [47] while achieving decent calibration. In [48], a systematic analysis of protein regression models was conducted where well-calibrated uncertainties were observed for a range of input representations.
*The following is the revised section 2.2 “Kernel methods for protein sequences” (formerly “Kernel methods”), followed by the newly added section 2.3 “Uncertainty quantification and calibration”, and references. We split them into separate comments due to character constraints.*

**2.2 Kernel methods for protein sequences**

Kernel methods have seen much use for protein modeling and specifically protein property prediction. Sequence-based string kernels operating directly on the protein amino acid sequences are one such example, where, e.g., matching k-mers at different ks quantify covariance. This has been used with support vector machines to predict protein homology [24,25]. Another example is sub-sequence string kernels, which in [26] is used in a Gaussian process for a Bayesian optimization procedure. In [27], string kernels were combined with predicted physico-chemical properties to further improve accuracy in the prediction of MHC-binding peptides and in protein fold classification. In [28], a kernel leveraging the tertiary structure for a protein family represented as a residue-residue contact
map, was used to predict various protein properties such as enzymatic activity and binding affinity. In [29], Gaussian process regression (GPR) was used to successfully identify promising enzyme sequences which were subsequently synthesized showing increased activity. In [30], the authors provide a comprehensive study of kernels on biological sequences which includes a thorough review of the literature as well as both theoretical, simulated, and in-silico results.

Most similar to our work is mGPfusion [31], in which a weighted decomposition kernel was defined which operated on the local tertiary protein structure in conjunction with a number of substitution matrices. Simulated stability data for all possible single mutations were obtained via Rosetta [32], which was then fused with experimental data for accurate ∆∆G predictions of single- and multi-mutant variants via Gaussian process regression, thus incorporating both sequence, structure, and a biophysical simulator. In contrast to our approach, the mGPfusion-model does not leverage pretrained models, but instead relies on substitution matrices for its evolutionary signal. A more recent example of kernel-based methods yielding highly competitive results is xGPR [33], in which GPs with custom kernels show high performance when trained on protein language model embeddings, similarly to the sequence kernel in our work (see Section 3). They use a set of novel random feature-approximated kernels with linear-scaling, where we use the squared exponential kernel. Moreover xGPR does not model structural environments, unlike our structure kernel. Their methods were shown to provide both high accuracy and well-calibrated uncertainty estimation on the FLIP and TAPE benchmarks.