Contrastive losses as generalized models of global epistasis

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 18:35:20

Abstract: Fitness functions map large combinatorial spaces of biological sequences to properties of interest. Inferring these multimodal functions from experimental data is a central task in modern protein engineering. Global epistasis models are an effective and physically-grounded class of models for estimating fitness functions from observed data. These models assume that a sparse latent function is transformed by a monotonic nonlinearity to emit measurable fitness. Here we demonstrate that minimizing supervised contrastive loss functions, such as the Bradley-Terry loss, is a simple and flexible technique for extracting the sparse latent function implied by global epistasis. We argue by way of a fitness-epistasis uncertainty principle that the nonlinearities in global epistasis models can produce observed fitness functions that do not admit sparse representations, and thus may be inefficient to learn from observations when using a Mean Squared Error (MSE) loss (a common practice). We show that contrastive losses are able to accurately estimate a ranking function from limited data even in regimes where MSE is ineffective and validate the practical utility of this insight by demonstrating that contrastive loss functions result in consistently improved performance on empirical benchmark tasks.

Main paper: 
 introduction: 
 A fitness function maps biological sequences to relevant scalar properties of the sequences, such as binding affinity to a target molecule, or fluorescent brightness. Biological sequences span combinatorial spaces and fitness functions are typically multi-peaked, due to interactions between positions in a sequence. Learning fitness functions from limited experimental data (often a minute fraction of the possible space) can be a difficult task but allows one to predict properties of sequences. These predictions can help identify promising new sequences for experimentation (Wu et al., 2019) or to guide the search for optimal sequences (Brookes et al., 2019; Bryant et al., 2021) .Even in the case of where experimental measurements are available for every possible sequence in a sequence space, inferring a model of the fitness function can be valuable for understanding the factors that impact sequences' fitness (Ding et al., 2022) or how evolution might progress over the fitness landscape (Wu et al., 2016) .Numerous methods have been developed to estimate fitness functions from experimental data, including classical machine learning techniques (Yang et al., 2019) , deep learning approaches (Gelman et al., 2021) , and semi-supervised methods (Hsu et al., 2022) . Additionally, there are many methods that incorporate biological assumptions into the modeling process, such as parameterized biophysical models (Otwinowski, 2018) , non-parametric techniques (Zhou & McCandlish, 2020; Zhou et al., 2022) , and methods for spectral regularization of neural networks (Aghazadeh et al., 2021) . These latter approaches largely focus on accurately modeling the influence of "epistasis" on fitness functions, which refers to statistical or physical interactions between genetic elements, typically either amino-acids in a protein sequence or genes in a genome."Local" epistasis refers to interactions between a limited number of specific positions in a sequence, and is often modeled using interaction terms in a linear model of a fitness function (Otwinowski & Plotkin, 2014) . "Global" epistasis, on the other hand, refers to the presence of nonlinear relationships that affect the fitness of sequences in a nonspecific manner. A model of global epistasis typically assumes a simple latent fitness function is transformed by a monotonically increasing nonlinearity to produce observed fitness data (Sailer & Harms, 2017; Otwinowski et al., 2018; Tareen et al., 2022;  Further, we use an entropic uncertainty principle to show that global epistasis can result in observed fitness functions that cannot be represented using a sparse set of epistatic interactions. In particular, this uncertainty principle shows that a fitness function that is sufficiently concentrated in the fitness domain-meaning that a small number of sequences have fitness values with relatively large magnitudes-can not be concentrated in the Graph Fourier bases that represent fitness functions in terms of local epistatic interactions (Stadler, 1995; Weinberger, 1991; Brookes et al., 2022) . We show that global epistasis nonlinearities tend to concentrate observed fitness functions in the fitness domain, thus preventing a sparse representation in the epistatic domain. This insight has the implication that observed fitness functions that have been affected by global epistasis may be difficult to estimate with undersampled training data and a Mean Squared Error (MSE) loss. We hypothesize that estimating the latent ranking fitness function using a contrastive loss can be done more data-efficiently than estimating the observed fitness function using MSE, and conduct simulations that support this hypothesis. Additionally, we demonstrate the practical importance of these insights by showing that models trained with the Bradley-Terry loss outperform those trained with MSE loss on nearly all FLIP benchmark tasks (Dallago et al., 2021) . 
fitness functions and the graph fourier transform: 
 A fitness function f : S → R maps a space of sequences S to a scalar property of interest. In the case where S contains all combinations of elements from an alphabet of size q at L sequence positions, then the fitness function can be represented exactly in terms of increasing orders of local epistatic interactions. For binary sequences (q = 2), this representation takes the form:f (x) = β 0 + L i=1 β i x i + ij β ij x i x j + ijk β ijk x i x j x k + ...,where x i ∈ {-1, 1} represent elements in the sequence and each term in the expansion represents a (local) epistatic interaction with weight β {i} , with the expansion continuing up to L th order terms. Analogous representations can be constructed for sequences with any size alphabet q using Graph Fourier bases (Stadler, 1995; Weinberger, 1991; Brookes et al., 2022) . These representations can be compactly expressed as:EQUATIONwhere f is a length q L vector containing the fitness values of every sequence in S, Φ is a q L × q L orthogonal matrix representing the Graph Fourier basis, and β is a length q L vector containing the weights corresponding to all possible epistatic interactions. We refer to f and β as representing the fitness function in the fitness domain and the epistatic domain, respectively. Note that we may apply the inverse transformation of Eq. 1 to any complete observed fitness function, y to calculate the epistatic representation of the observed data, β y = Φ T y. Similarly, if f contains the predictions of a fitness model for every sequence in a sequence space, then β = Φ T f is the epistatic representation of the model.A fitness function is considered sparse, or concentrated, in the epistatic domain when β contains a relatively small number of elements with large magnitudes, and many elements equal to zero or with small magnitudes. In what follows, we may refer to a fitness function that is sparse in the epistatic domain as simply being a "sparse fitness function". A number of experimentally-determined fitness functions have been observed to be sparse in the epistatic domain (Poelwijk et al., 2019; Eble et al., 2020; Brookes et al., 2022) . Crucially, the sparsity of a fitness function in the epistatic domain determines how many measurements are required to estimate the fitness function using Compressed Sensing techniques that minimize a MSE loss function (Brookes et al., 2022) . Herein we consider the effect that global epistasis has on a sparse fitness function. In particular, we argue that global epistasis results in observed fitness functions that are dense in the epistatic domain, and thus require a large amount of data to accurately estimate by minimizing a MSE loss function. However, in these cases, there may be a sparse ranking function that can be efficiently extracted by minimizing a contrastive loss function.A model of global epistasis assumes that noiseless fitness measurements are generated according to the model:EQUATIONwhere f is a latent fitness function, g is a monotonically increasing nonlinear function. In most cases, f is assumed to include only first or second order epistatic terms and the nonlinearity is explicitly parameterized using, for example, spline functions (Otwinowski et al., 2018) or sums of hyperbolic tangents (Tareen et al., 2022) . The restriction that f includes only low-order terms is somewhat arbitrary, as higher-order local epistatic effects have been observed in fitness data (see, e.g., (Wu et al., 2016) ). In general we may consider f to be any fitness function that is sparse in the epistatic domain, and global epistasis then refers to the transformation of a sparse fitness function by a monotonically-increasing nonlinearity.Global epistasis models of the form of Eq. 2 have proved effective at capturing the variation observed in empirical fitness data (Kryazhimskiy et al., 2014; Sarkisyan et al., 2016; Otwinowski et al., 2018; Pokusaeva et al., 2019; Reddy & Desai, 2021; Tareen et al., 2022) , suggesting that global epistasis is a common feature of natural fitness functions. Further, it has been shown that global epistasis results from first-principles physical considerations that are common in many biological systems. In particular, Husain & Murugan (2020) show that global epistasis arises when the physical dynamics of a system is dominated by slow, collective modes of motion, which is often the case for protein dynamics. Aside from intrinsic/endogenous sources, the process of measuring fitness can also introduce nonlinear effects that are dependent on the experiment and not on the underlying fitness function. For example, fitness data is often left-censored, as many sequence have fitness that falls below the detection threshold of an assay. Finally, global diminishing-returns epistatic patterns have been observed widely in both single and multi-gene settings where the interactions are among genes rather than within a gene (Kryazhimskiy et al., 2014; Reddy & Desai, 2021; Bakerlee et al., 2022) .Together, these results indicate that global epistasis is an effect that can be expected in empiricallyobserved fitness functions. In what follows, we argue that global epistasis is detrimental to effective modeling of fitness functions using standard techniques. In particular, global epistasis manifests itself by producing observed data that is dense in the epistatic domain. In other words, when observed fitness data is produced through Eq. 2 then the epistatic representation of this fitness function (calculated through application of Eq. 1), is not sparse. Further we argue that this effect of global epistasis makes it to difficult to model such observed data by minimizing standard MSE loss functions with a fixed amount of data. Further, we argue that fitting fitness models aimed at extracting the latent fitness function from observed data is a more efficient use of observed data that results in improved predictive performance (in the ranking sense). While the models of global epistasis described thus far could be used for this purpose, they have the drawback that they assume a constrained form of both g and f , which enforces inductive biases that may affect predictive performance. Here we propose a flexible alternative to modeling global epistasis that makes no assumptions on the form of f or g. In particular, we interpret the latent fitness function f as a parsimonious ranking function for sequences, and the problem of modeling global epistasis as recovering this ranking function. A natural method to achieve this goal is to fit a model of f to the observed data by minimizing a contrastive, or ranking, loss function. These loss functions are designed to learn a ranking function and, as we will show, are able to recover a sparse fitness function that has been transformed by global epistasis to produce observed data. An advantage of this approach to modeling global epistasis is that the nonlinearity g is modeled non-parametrically, and is free to take any form, while the latent fitness function can be modeled by any parametric model, for example, convolutional neural networks (CNNs) or fine-tuned language models, which have been found to perform well in fitness prediction tasks (Dallago et al., 2021) . An accurate ranking model also enables effective optimization, as implied by the results in Chan et al. (2021) .Contrastive losses broadly refer to loss functions that compare multiple outputs of a model and encourage those outputs to be ordered according to some criteria. In our case, we desire a loss function that encourages model outputs to be ranked according to observed fitness values. An example of such a loss function is the Bradley-Terry (BT) loss (Bradley & Terry, 1952; Burges et al., 2005) , which has the form:EQUATIONwhere f θ is a model with parameters θ, x i are model inputs and y i are the corresponding labels of those inputs. This loss compares every pair of data points and encourages the model output f θ (x i ) to be greater than f θ (x j ) whenever y i > y j ; in other words, it encourages the model outputs to be ranked according to their labels. A number of distinct but similar loss functions have been proposed in the learning-to-rank literature (Chen et al., 2009) and also for metric learning (Chopra et al., 2005) . An example is the Margin ranking loss (Herbrich et al., 1999) , which replaces the logistic function in the sum of Eq. 3 with a hinge function.In our experiments, we largely focus on the BT loss of Eq. 3 as we found it typically results in superior predictive performance; however, we also compare to models trained with the Margin loss when presenting results in benchmark tasks.The BT loss was recently used by Chan et al. (2021) to order the latent space of a generative model for protein sequences such that certain regions of the latent space corresponding to sequences with higher observed fitness values. In this case, the BT loss is used in conjunction with standard generative modeling losses. In contrast, here we analyze the use of the BT loss alone in order to learn a ranking function for sequences given corresponding observed fitness values.A key feature of the contrastive loss in Eq. 3 is that it only uses information about the ranking of observed labels, rather than the numerical values of the labels. Thus, the loss is unchanged when the observed values are transformed by a monotonic nonlinearity. We will show that this feature allows this loss to recover a sparse latent fitness function from observed data that has been affected by global epistasis, and enables more data-efficient learning of fitness functions compared to a MSE loss. 
results: 
 Our results are aimed at demonstrating three properties of contrastive losses. First, we show that given complete, noiseless fitness data (i.e. noiseless fitness values associated with every sequence in the sequence space) that has been affected by global epistasis, minimizing the BT loss enables a model to nearly exactly recover the sparse latent fitness function f . Next, we consider the case of incomplete data, where the aim is to predict the relative fitness of unobserved sequences. In this regime, we find through simulation that minimizing the BT loss enables models to achieve better predictive performance then minimizing the MSE loss when the observed data has been affected with global epistasis. We argue by way of a fitness-epistasis uncertainty principle that this is due to the fact that nonlinearities tend to produce fitness functions that do not admit a sparse representation in the epistatic domain, and thus require more data to learn with MSE loss. Finally, we demonstrate the practical significance of these insights by showing that minimizing the BT loss results in improved predictive performance over MSE loss in nearly all tasks in the FLIP benchmark (Dallago et al., 2021) .We first consider the case of "complete" data, where fitness measurements are available for every sequence in the sequence space. The aim of our task in this case is to recover a sparse latent fitness function when the observed measurements have been transformed by an arbitrary monotonic nonlinearity. In particular, we sample a sparse fitness function f from the NK model (Kauffman & Weinberger, 1989 ), a popular model of fitness functions that has been shown to recapitulate the sparsity properties of some empirical fitness functions (Brookes et al., 2022) . The NK model has three parameters: L, the length of the sequences, q, the size of the alphabet for sequence elements, and K, the maximum order of (local) epistatic interactions in the fitness function. Roughly, the model randomly assigns K -1 interacting positions to each position in the sequence, resulting in a sparse set of interactions in the epistatic domain. The weights of each of the assigned interactions are then drawn from a independent unit normal distributions.We then transform the sampled fitness function f with a monotonic nonlinearity g to produce a set of complete observed data, y i = g(f (x i )) for all x i ∈ S. The goal of the task is then to recover the function f given all (x i , y i ) pairs. In order to do so, we model f using a two layer fully connected neural network and fit the parameters of this model by performing stochastic gradient descent (SGD) on the BT loss, using the Spearman correlation between model predictions and the y i values to determine convergence of the optimization. We then compare the resulting model, f , to the latent fitness function f in both the fitness and epistatic domains, using the forward and inverse transformation of Eq. 1 to convert between the two domains.Fig. 1 shows the results of one of these tests. In this case, we used an exponential function to represent the global epistasis nonlinearity. The exponential function exaggerates the effects of global epistasis in the epistatic domain and thus better illustrates the usefulness of contrastive losses, although the nonlinearities in empirical fitness functions tend to have a more sigmoidal shape (Otwinowski et al., 2018) . Fig. 1b shows that the global epistasis nonlinearity substantially alters the representations of the observed data y in both the fitness and epistatic domains, as compared to the latent fitness function f . Nonetheless, Fig. 1c demonstrates that the model fitness function f created by minimizing the BT loss is able to nearly perfectly recover the sparse latent fitness function (where recovery is defined as being equivalent up to an affine transformation). This is a somewhat surprising result, as there are many fitness functions that correctly rank the fitness of sequences, and it is not immediately clear why minimizing the BT loss produces this particular sparse latent fitness function. However, this example makes clear that fitting a model by minimizing the BT loss can be an effective strategy for recovering a sparse latent fitness function from observed data that has been affected by global epistasis. Similar results from additional examples of this task using different nonlinearities and latent fitness functions are shown in Appendix B.1. Further, in Appendix B.2 we compare the results of Fig. 1c to a baseline method in which a quantile transformation is applied to the observed data.Next, we consider the case where fitness data is incomplete. Our aim is to understand how models trained with the BT loss compare to those trained with MSE loss at predicting the relative fitness of unobserved sequence using different amounts of subsampled training data. We take a signal processing perspective on this problem, and consider how the density of a fitness function in the epistatic domain affects the ability of a model to accurately estimate the fitness function given incomplete data. In particular, we demonstrate that global epistasis tends to increase the density of fitness functions in the epistatic domain, and use an analogy to Compressive Sensing (CS) to hypothesize that more data is required to effectively estimate these fitness functions when using an MSE loss (Brookes et al., 2022) . In order to support this claim, we first examine the effects of global epistasis on the epistatic domain of fitness functions.Fig. 1b provides an example where a monotonic nonlinearity applied to a sparse fitness increases the density of the the fitness function in the epistatic domain. In particular, we see that many "spurious" local epistatic interactions must appear in order to represent the nonlinearity (e.g. interactions of order 3, when we used an NK model with K = 2). This effect can be observed for many different shapes of nonlinearities (Sailer & Harms, 2017; Baeza-Centurion et al., 2019) . We can understand this effect more generally using uncertainty principles, which roughly show that a function cannot be concentrated on a small number of values in two different representations. In particular, we consider the discrete entropic uncertainty principle proved by Dembo et al. (1991) . When applied to the transformation in Eq. 1, this uncertainty principle states:EQUATION)whereH(x) = -i x 2 i ||x|| 2 log x 2 i||x|| 2 is the entropy of the normalized squared magnitudes of a vector and m = 1/ √ q when q = 2, m = 1/(q -√ q) when q = 3 and m = 1 -1/(q -√ q) otherwise.Low entropy indicates that a vector is highly concentrated on a small set of elements. Thus, the fitnessepistasis uncertainty principle of Eq. 4 shows that fitness functions cannot be highly concentrated in both the fitness and epistatic domains. A sparse fitness function (in the epistatic domain) must therefore be "spread out" (i.e. dense) in the fitness domain, and vice-versa.The importance of this result for understanding global epistasis is that applying a nonlinearity to a dense vector will often have the effect of concentrating the vector on a smaller number of values. This can most easily be seen for convex nonlinearities like the exponential shown in the entropy in the fitness and epistatic domains decrease and increase, respectively, as more extreme nonlinearities are applied to a sparse latent fitness function.The uncertainty principle quantifies how global epistasis affects a fitness function by preventing a sparse representation in the epistatic domain. From a CS perspective, this has direct implications for modeling the fitness function from incomplete data. In particular, if we were to model the fitness function using CS techniques such as LASSO regression with the Graph Fourier basis as the representation, then it is well known that the number of noiseless data points required to perfectly estimate the function scales as O(S log N ) where S is the sparsity of the signal in a chosen representation and N is the total size of the signal in the representation (Candes et al., 2006) . Therefore, when using these techniques, fitness functions affected by global epistasis will require more data to effectively model. Notably, the techniques for which these scaling laws apply minimize a MSE loss functions as part of the estimation procedure. Although these scaling laws only strictly apply to CS modeling techniques, we hypothesize that the intuition that fitness functions with dense epistatic representations will require more data to train an accurate model with MSE loss, even when using neural network models and SGD training procedures. In the next section we present the results of simulations that support this hypothesis by showing that the entropy of the epistatic representation is negatively correlated with the predictive performance of models trained with an MSE loss on a fixed amount of fitness data. Further, these simulations show that models trained with the BT loss are robust to the dense epistatic representations produced global epistasis, and converge faster to maximum predictive performance as they are provided more fitness data compared to models trained with an MSE loss.We next present simulation results aimed at showing that global epistasis adversely effects the ability of models to effectively learn fitness functions from incomplete data when trained with MSE loss and that models trained with BT loss are more robust to the effects of global epistasis.In our first set of simulations, we tested the ability models to estimate a fitness function of L = 10 binary sequences given one quarter of the fitness measurements (256 measurements out of total of 2 10 = 1024 sequences in the sequence space). In each simulation, we (i) sampled a sparse latent fitness function from the NK model, (ii) produced an observed fitness function by applying one of three nonlinearities to the latent fitness function: exponential, g(f ) = exp(a • f ), sigmoid, g(f ) = (1 + e -a•f ) -1 , or a cubic polynomial g(f ) = x 3 + ax with settings of the parameter a that ensured the nonlinearity was monotonically increasing, (iii) sampled 256 sequence/fitness pairs uniformly at random from the observed fitness function to be used as training data, and (iv) trained models with this data by performing SGD on the MSE and BT losses. We ran this simulation for 50 replicates of each of 20 settings of the a parameter for each of the three nonlinearities. In every case, the models were fully-connected neural networks with two hidden layers and the optimization was terminated using early stopping with 20 percent of the training data used as a validation set. After training, we measured the extent to which the models estimated the fitness function by calculating Spearman correlation between the model predictions and true fitness values on all sequences in the sequence space. Spearman correlation is commonly used to benchmark fitness prediction methods (Dallago et al., 2021; Notin et al., 2022) .The results of each of these simulations are shown in Fig. 2b . We see that the predictive performance of models trained with the MSE loss degrades as the entropy of the fitness function in the epistatic domain increases, regardless of the type of nonlinearity that is applied to the latent fitness function. This is in contrast to the models trained with the BT loss, which often achieve nearly perfect estimation of the fitness function even when the entropy of the fitness function in the epistatic domain approaches its maximum possible value of L log 2. This demonstrates the key result that the MSE loss is sensitive to the density of the epistatic representation resulting from global epistasis (as implied by the analogy to CS), while the BT loss is robust to these effects. We additionally find that these results are maintained when the degree of epistatic interactions in the latent fitness function is changed (Appendix D.1) and when Gaussian noise is added to the observed fitness functions (Appendix D.2Next we tested how training set size effects the predictive performance of models trained with MSE and BT losses on a fitness function affected by global epistasis. In order to do so, we sampled a single L = 10, K = 2 fitness function from the NK model and applied the nonlinearity g(f ) = exp(10 • f ) to produce an observed fitness function. Then, for each of a range of training set sizes between 25 There are two important takeaways from this plot. First, the BT loss achieves higher Spearman correlations than the MSE loss in all data regimes. This demonstrates the general effectiveness of this loss to estimate fitness functions affected by global epistasis. Next, we see that models trained with BT loss converge to a maximum Spearman correlation faster than models trained with MSE loss do to a maximum Pearson correlation, which demonstrates that the difference in predictive performance between models trained with MSE and BT losses is not simply due to a result of the evaluation metric being more tailored to one loss than the other. This result also reinforces our claim that fitness functions affected by global epistasis require more data to learn effectively with MSE loss, as would be predicted by CS scaling laws. The BT loss on the other hand, while not performant with the Pearson metric as expected by a ranking loss, seems to overcome this barrier and can be used to effectively estimate a fitness function from a small amount of data, despite the effects of global epistasis.In the previous sections, we used noiseless simulated data to explore the interaction between global epistasis and loss functions. Now we present results demonstrating the practical utility of our insights by comparing the predictive performance of models trained with MSE and BT losses on experimentally-determined protein fitness data. We particularly focus on the FLIP benchmark (Dallago et al., 2021) , which comprises of a total of 15 fitness prediction tasks stemming from three empirical fitness datasets. These three datasets explore multiple types of proteins, definitions of protein fitness, and experimental assays. In particular, one is a combinatorially-complete dataset that contains the binding fitness of all combinations of mutations at 4 positions to the GB1 protein (Wu et al., 2016) , another contains data about the viability of Adeno-associated virus (AAV) capsids for many different sets of mutations to the wild-type capsid sequence (Bryant et al., 2021) , and another contains data about the thermostability of many distantly related proteins (Jarzab et al., 2020) .For each of the three datasets, the FLIP benchmark provides multiple train/test splits that are relevant for protein engineering scenarios. For example, in the GB1 and AAV datasets, there are training sets that contain only single and double mutations to the protein, while the associated test sets contain sequences with more than two mutations. This represents a typical situation in protein engineering where data can easily be collected for single mutations (and some double mutations) and the goal is then to design sequences that combine these mutations to produce a sequence with high fitness. In all of the FLIP tasks the evaluation metric is Spearman correlation between model predictions and fitness labels in the test set, since ranking sequences by fitness is the primary task that models are used for in data-driven protein engineering.In the FLIP benchmark paper, the authors apply a number of different modeling strategies to these splits, including Ridge regression, training a CNN, and a number of variations on fine-tuning the ESM language models for protein sequences (Rives et al., 2021) . All of these models use a MSE loss to fit the model to the data, along with any model-specific regularization losses. In our tests, we consider only the CNN model as it balances consistently high performance in the benchmark tasks with relatively straightforward training protocols, enabling fast replication with random restarts.We trained the CNN model on each split using the standard MSE loss as well as the Margin and BT contrastive losses. The mean and standard deviation of Spearman correlations between the model predictions and test set labels over 10 random restarts are shown in Fig. 3 . By default, the FLIP datasets contain portions of sequences that are never mutated in any of the data (e.g., only 4 positions are mutated in the GB1 data, but the splits contain the full GB1 sequence of length 56). We found that including these unmodified portions of the sequence often did not improve, and sometimes hurt, the predictive performance of the CNN models while requiring significantly increased computational complexity. Therefore most of our results are reported using inputs that contain only sequence positions that are mutated in at least one train or test data point. We found that including the unmodified portions of sequences improved the performance for the Low-vs-High and 3-vs-Rest GB1 well splits, as well as the 1-vs-rest AAV split and so these results are reported in Fig. 3 ; in these cases we found both models trained with MSE and contrastive losses had improved performance.The results in Fig. 3 show that using contrastive losses (and particularly the BT loss) consistently results in improved predictive performance across a variety of practically relevant fitness prediction tasks (for more detail, see Appendix E.1 for the results in table form). Further, in no case does the BT loss result in worse performance than MSE. Indeed, it is shown in Otwinowski et al. (2018) that a GB1 landscape closely associated with that in the FLIP benchmark is strongly affected by global epistasis. Further, many of the FLIP training sets are severely undersampled in the sense of CS scaling laws, which is the regime in which differences between MSE and contrastive losses are most apparent when global epistasis is present, as shown in Fig. 2 .Although Spearman correlation is commonly used to benchmark models of fitness functions, in practical protein engineering settings it also important to consider the ability of the model to classify the sequences with the highest fitness. To test this, we calculated the "top 10% recall" for models trained with the BT and MSE losses on the FLIP benchmark data, which measures the ability of the models to correctly classify the 10% of test sequences with the highest fitness in the test set Gelman et al. (2021) . These results are shown in Appendix E.2, where it can be seen that the BT loss consistently outperforms the MSE loss. 
discussion: 
 Our results leave open a few avenues for future exploration. First, it is not immediately clear in what situations we can expect to observe the nearly-perfect recovery of a latent fitness function as seen in Fig. 1 . A theoretical understanding of this result may either cement the promise of the BT loss, or provide motivation for the development of techniques that can be applied in different scenarios.Next, we have made a couple of logical steps in our interpretations of these results that are intuitive, but not fully supported by any theory. In particular, we have drawn an analogy to CS scaling laws to explain why neural networks trained with MSE loss struggle to learn a fitness function that has a dense representation in the epistatic domain. However, these scaling laws only strictly apply for a specific set of methods that use an orthogonal basis as the representation of the signal; there is no theoretical justification for using them to understanding the training of neural networks (although applying certain regularizations to neural network training can provide similar guarantees (Aghazadeh et al., 2022) ). Additionally, it is not clear from a theoretical perspective why the BT loss seems to be robust to the dense representations produced by global epistasis. A deeper understanding of these phenomena could be useful for developing improved techniques.Additionally, our simulations largely do not consider how models trained with contrastive losses may be affected by the complex measurement noise commonly seen in experimental fitness assays based on sequencing counts (Busia & Listgarten, 2023) . Although our simulations mostly do not consider the effects of noise (except for the simple Gaussian noise added in Appendix D.2), our results on the FLIP benchmark demonstrate that contrastive losses can be robust to the effects of noise in practical scenarios. Further, we show in Appendix F: that the BT loss can be robust to noise in a potentially pathological scenario. A more complete analysis of the effects of noise on contrastive losses would complement these results.from the set of positions {1, ..., L} \ j uniformly at random, without replacement. Now let S (L,q) be the set of all q L sequences of length L and alphabet size q. Given the sampled neighborhoods {V [j] } L j=1 , the NK model assigns a fitness to every sequence in S (L,q) through the following two steps: 1. Let s [j] := (s k ) k∈V [j] be the subsequence of s corresponding to the indices in the neighborhood V [j] . Assign a 'subsequence fitness', f j (s [j] ) to every possible subsequence, s [j] , by drawing a value from the normal distribution with mean equal to zero and variance equal to 1 /L. In other words, f j (s [j] ) ∼ N (0, 1 /L) for every s [j] ∈ S (K,q) , and for every j = 1, 2, ..., L.2. For every s ∈ S (L,q) , the subsequence fitness values are summed to produce the total fitness values f (s) = L j=1 f j (s [j] ).In order to provide a baseline for the results shown in Figure 1c , we attempted to recover the latent fitness function using quantile transformations rather than using the Bradley-Terry loss. In particular, we applied a quantile transformation the observed fitness data y to produce a vector f q that approximately follows either a uniform or Gaussian distribution. In both cases, we used 10 quantiles for discretization. The results of these baseline tests are shown in Figure 5 , below. By comparing to the results in Figure 1c , we can see that the quantile transformations are not as effective at recovering the latent fitness function as the Bradley-Terry loss. A potential disadvantage of the BT and Margin losses used in the main text is that neither takes into consideration the size of the true observed gap between the fitness of pairs of sequences. In particular, Eq. 3 shows that any two sequences in which y i > y j will be weighted equally in the loss, regardless of the magnitude of the difference y i -y j . This has implications for how measurement noise may affect these losses. In particular, the relative ranking between pairs of sequences with a small gap in fitness in more likely to have been swapped due to measurement noise, compared to a pair of sequences with a large gap. This suggests that contrastive losses may exhibit pathologies in the presence of certain types of noise.Here we examine a toy scenario in which the observed fitness data has the form y i = I(x i ) + ϵ, where I(x i ) is an indicator function that is equal to zero or one and ϵ ∼ N (0, σ 2 ) is Gaussian noise. One may expect contrastive loses to exhibit pathologies when trained on this data because the relative rankings between sequences that have the same value of I(x) is due only to noise.In order to construct this scenario, we sampled an L = 10, K = 2 binary fitness function f (x from the NK model, and then let I(x) = 0 when f (x) < med(f ) and 1 otherwise, where med(f ) is the median NK fitness of all sequences. We then added Gaussian noise with σ = 0.1 to produce the observed noisy data y. We consider I(x) to be the true binary label of a sequence and y to be the noisy label of the sequence. The relationship between the NK fitness, true labels and noisy labels is shown in Fig. 10a , below. Next, we randomly split the data into training and test sets containing 100 and 924 sequences and noisy labels, respectively. We used the training set to fit two neural network models, one using the MSE loss and the other using the BT loss. The models and training procedure were the same as described in Appendix A:. The predictions of these models on test sequences, compared to noisy labels, are shown in 10b. We next constructed Receiver Operating Characteristic (ROC) curves that test the ability of each model to classify sequences according to their true labels I(x) (Fig. 10c ). These curves show that while the MSE loss does slightly outperform the BT loss in this classification task, the BT loss still results in an effective classifier and does not exhibit a major pathology in this scenario.Notably, on empirical protein landscapes, as shown in Table 2 , the performance gain by contrastive losses out-weighs the potential drawback of further sensitivity to this kind of noise, and contrastive losses ultimately result in improved predictive performance.  
 
 Appendix:  


Decision: This paper introduces an approach to fitness function prediction by leveraging the Bradley-Terry (BT) loss, a ranking-based loss function, instead of the traditional Mean Squared Error (MSE) loss. The authors argue that the BT loss is better suited for modeling global epistasis, where observed fitness data is a non-linear transformation of the true underlying fitness. The experiments include both simulated datasets and real-world benchmarks like FLIP, showing that BT loss improves prediction accuracy and interpretability.

Reviewers generally appreciated the paper's well-motivated approach, solid theoretical foundations, and clear presentation. 

**Positive Points:**
* The BT loss's flexibility and alignment with biological interpretability, as noted by Reviewer iMMK.
* The paper's strong theoretical investigation into why BT loss outperforms MSE.
* The practical significance of modeling global epistasis in real-world protein engineering settings. 

**Negative Points:**
* Reviewer 5WR7 pointed out the lack of a rigorous theoretical analysis explaining the difference between BT and MSE losses under various conditions.
* Reviewers pZCs and GFMs noted that comparisons with more standard benchmarks like ProteinGym and other state-of-the-art methods were insufficient.
* Reviewer NoUX expressed concerns about the necessity of introducing a nonlinearity in the fitness model.


Despite some valid critiques, the paper's strengths outweigh its weaknesses. The introduction of BT loss as an alternative to MSE for fitness prediction is a valuable contribution, particularly in the context of global epistasis where traditional loss functions may falter. The authors have addressed the primary concerns raised by the reviewers, providing additional experimental results on relevant ProteinGym datasets and enhancing the theoretical grounding of their approach. While further comparisons with a broader range of benchmarks and models would strengthen the paper, the current work already offers significant insights and practical benefits.

Reviews: review: paper_summary: paper_summary: This study assumes a global epistasis model, where the observed experimental data represents a monotonic nonlinear transformation of underlying fitness, and explores the problem of learning these underlying fitness functions. Instead of directly learning fitness functions using Mean Squared Error (MSE) loss, the authors propose using ranking-based loss, such as Bradley-Terry (BT) loss, to learn fitness functions indirectly. The experiments in this paper demonstrate that ranking-based loss outperforms MSE loss.
strength_weakness: strength_weakness: 1. Several previous methods account for high-order interactions, including MSA-based methods like EVE[1] and MSA-Transformer[2], and language model-based methods such as Tranception[3] and esm1v[4]. The authors should include these methods for comparison.
2. ProteyGym[5] is a more comprehensive benchmark dataset widely used in the literature. Please also include the evaluation of this dataset.

I will be glad to raise my score if more solid comparisons can be provided.

[1] Genome-wide prediction of disease variant effects with a deep protein language model.  https://doi.org/10.1038/s41588-023-01465-0

[2] MSA Transformer.  https://proceedings.mlr.press/v139/rao21a.html

[3] Tranception: Protein Fitness Prediction with Autoregressive Transformers and Inference-time Retrieval.  https://proceedings.mlr.press/v162/notin22a

[4] Language models enable zero-shot prediction of the effects of mutations on protein function.  https://openreview.net/forum?id=uXc42E9ZPFs

[5] ProteinGym: Large-Scale Benchmarks for Protein Fitness Prediction and Design.  https://papers.nips.cc/paper_files/paper/2023/hash/cac723e5ff29f65e3fcbb0739ae91bee-Abstract-Datasets_and_Benchmarks.html
questions: questions: 1. Please refer to the Weaknesses section
2. AlphaMissense[1] should be included in the Related Work section, as it also incorporates a contrastive-like loss in its training.

[1] Accurate proteome-wide missense variant effect prediction with AlphaMissense.  https://doi.org/10.1126/science.adg7492
limitations: limitations: I have no concerns about the potential negative societal impact of their work.
paper_summary: paper_summary: This paper addresses the problem setting of modeling and extracting the sparse interactions found in global epistasis, and proposes that the Bradley-Terry loss -- a loss used in ranking problems -- is a good alternative to MSE for ranking prediction in such settings, as it does not require assumptions on the form of the non-linearity. Authors use the entropic uncertainty principle from physics and information theory to prove that observed fitness functions cannot always be represented by sparse locally epistatic interactions, due to the presence of global epistasis; since the sparsity of a fitness function determines how many measurements are necessary to estimate the fitness function in MSE settings, this renders the Bradley-Terry loss to be more performative under globally epistatic settings. This is also shown empirically; in experiments, authors examine the utility of using BT loss instead of MSE in (1) simulated data where all fitness values are available, (2) simulated data where fitness labels are only partially observed, and (3) real world benchmarks (i.e. FLIP).
strength_weakness: strength_weakness: * While the method does achieve SOTA on FLIP benchmarks, FLIP datasets are rather limited in how they approximate the real-world -- diversity within the GB1 dataset is still rather locally epistatic

Minor:
* The abbreviation "CS" is used a lot in the later sections of the paper ("CS scaling laws", "CS techniques such as LASSO"); would be good to expand this to reduce confusion.
questions: questions: n/a
limitations: limitations: n/a
paper_summary: paper_summary: The authors study global epistasis models which are used to understand the fitness landscapes of biological sequences. They start with the observation that global epistasis is observed in real-world systems as a monotonic non-linear transform of an underlying fitness function on a genotype/sequence. The fitness function can have varying levels of local epistasis. They then introduce the Bradley-Terry (BT) loss, a "contrastive loss" that penalizes model prediction ordering, as an alternative to MSE as the loss function when training a model to predict f when observables are of the form y = g(f(x)), where g is some unknown monotonic non-linearity and f is the fitness function. In a series of simulations and on real protein fitness data they show that they can recover f better when using the BT loss in place of MSE, and suggest this may be due to the distribution of the label on the domain of the underlying non-linearity.
strength_weakness: strength_weakness: Probably the largest weakness is that the authors don't providet a theoretical analysis that explains the difference in results between the loss functions in a concrete way, taking into account the conditions under which we could expect these results to hold. The analogy to compressed sensing is interesting but ultimately relies on the authors to conjecture about the causes. In particular there isn't a conclusive assessment of the effect of noise since many of the results rely on lossless simulation (although the authors do experiment with the FLIP benchmark). I think the paper could be improved with the addition of some rigorous analysis that links the hypothesis class of $g$ to its learnability under different loss functions. 

It's also maybe not surprising that the model trained on BT-loss would do better as measured by Spearman correlations, given that both are measures of rank order.
questions: questions: How do the resultant models (trained on MSE vs BT) compare in terms of MSE on test data?
limitations: limitations: The authors cover limitations well.
paper_summary: paper_summary: The paper proposes using contrastive loss, specifically the Bradley-Terry (BT) loss, as an alternative to Mean Squared Error (MSE) for training global epistasis models in fitness prediction tasks. Experiments conducted on both complete and incomplete data sets, as well as benchmarked on FLIP tasks, show that the BT loss improves performance in fitness prediction.
strength_weakness: strength_weakness: - There is a lack of comparison with other standard fitness prediction tasks: For fitness prediction tasks, there are several standard experiments not performed, such as fold classification, Gene Ontology (GO), Enzyme Commission (EC) prediction, fluorescence, localization prediction, and stability prediction. Referencing benchmarks such as TAPE[1], DeepLoc[2], and dataset in DeepFRI[3] would strengthen the evaluation.
    
- There is a lack of comparison with other standard fitness prediction models: The baseline models are restricted to CNNs. It would be beneficial to include protein language models (e.g., ESM[4], xTrimoPGLM[5]) and potentially structure-based models(e.g., ESM-GearNet[6]), as well as comparisons with other epistasis-aware networks.
    

[1] Rao, R., Bhattacharya, N., Thomas, N., Duan, Y., Chen, P., Canny, J., ... & Song, Y. (2019). Evaluating protein transfer learning with TAPE. *Advances in neural information processing systems*, *32*.

[2] Almagro Armenteros, José Juan, et al. "DeepLoc: prediction of protein subcellular localization using deep learning." *Bioinformatics* 33.21 (2017): 3387-3395.

[3] Gligorijević, Vladimir, et al. "Structure-based protein function prediction using graph convolutional networks." *Nature communications* 12.1 (2021): 3168.

[4] Lin, Zeming, et al. "Evolutionary-scale prediction of atomic-level protein structure with a language model." *Science* 379.6637 (2023): 1123-1130.

[5] Chen, Bo, et al. "xTrimoPGLM: unified 100B-scale pre-trained transformer for deciphering the language of protein." *arXiv preprint arXiv:2401.06199* (2024).

[6] Zhang, Zuobai, et al. "A systematic study of joint representation learning on protein sequences and structures." *arXiv preprint arXiv:2303.06275* (2023).
questions: questions: 1. Typically, contrastive loss involves training the model to distinguish between similar and dissimilar pairs of data points.  The Bradley-Terry (BT) loss described in the paper appears more like a regression loss. Is this interpretation correct?
limitations: limitations: The experimental evaluation is limited, as it only includes CNN baselines and FLIP dataset. The impact and novelty of the work could be limited without more thorough experimental results.
paper_summary: paper_summary: This paper studies how the contrastive loss improves the estimation of fitness functions over the MSE loss and contributes in the following ways: (1) With noiseless data in Section 3.1, they found that the contrastive loss can estimate not only the correct ranking of fitness but also the exact fitness values; (2) With modulating the sparseness of epistasis vectors in Section 3.3, the MSE performance deteriorates quickly, unlike the contrastive loss; (3) The contrastive loss performs well in most of the setups with the FLIP benchmark data.
strength_weakness: strength_weakness: Since I'm not familiar with the field of fitness estimation, my comments are mostly based on the perspective of machine learning researchers who work on contrastive learning and related topics.

+ **Unclear why we need the nonlinear $g$ in modeling**: The entire paper is based on the nonlinear measurement system $y = g(f(x))$ to incorporate the knowledge of "global" epistasis. However, the necessity of $g$ is unclear because the fitness function $f$ is modeled by a neural network, which is sufficiently capable of capturing the underlying nonlinearity. I'm not sure of the potential benefits of having an additional nonlinearity $g$.
  + One may argue that $g$ is necessary because this corroborates why the observed fitness values exhibit denseness---as described in the paragraph from l.245, the nonlinearity of $g$ makes fitness values dense even if the values of $f$ are sparse. While it makes sense that the nonlinearity yields the denseness, this only indicates that the nonlinearity is _sufficient_ but not necessary.
+ **The relationship between the entropy and sparseness might be elusive**: In Section 3.2, the authors explain the trade-off between the sparseness of fitness and epistasis by the fitness-epistasis entropy lower bound in Eq. (4). This explanation might have two concerns. First, this entropy lower bound is derived based on the "local" epistasis model $\\mathbf{f} = \\mathbf{\\Phi}\\boldsymbol{\\beta}$, but the authors argue the existence of the trade-off for the "global" epistasis, which is a leap in logic. Second, even if we focus on the local epistasis model, fitness $\\mathbf{f}$ must be sparse once epistasis $\\boldsymbol{\\beta}$ is a sparse vector (not in the sense of the entropy but in the sense of the number of nonzero elements). From this observation, the explanation by the entropy sounds a bit unnatural.
+ (Slightly minor) **Does the evaluation by the correlation coefficients make sense?**: The experimental results are evaluated by the Spearman/Pearson correlation coefficients. While it may look fine at first glance, this evaluation protocol is too advantageous for the contrastive loss that we cannot be sure whether the sparseness/denseness of epistasis is a key factor of the MSE deterioration or not, because the contrastive loss enhances correct ranking and kinda directly optimizes correlation.
questions: questions: + The difference in the experimental setups of Sections 3.1 and 3.3 is not clear. They should be different to some extent because the former section considers a noise-less setup, and the latter considers an "incomplete setup. In my understanding, their experimental setups are almost identical, including the data-generating processes, whereas the nonlinearity $g$ is modulated in Section 3.3 to see the effect of the sparseness on the estimation quality. And by the "incomplete" observation, the authors probably indicate different sparseness degrees. If this is the case, I feel the wording "incomplete" is a bit misleading.

----
Minor comments:
- It is standard to concatenate two peoples' names by ndash (--) such as "Bradley--Terry".
- At l.82, space might be missing in between "Graph Fourier bases" and the reference numbers.
- From l.137, it begins with "Here we propose," but since this section is background, it is better to separate the authors' approach into another section.
- In Figure 2b, why not having the mean and confidence intervals? The current scatter plot is a bit hard to grasp the trend.
- Section 3.2 seems to be providing background only. Why not having this in the background section?
- At l.362, "[...] practical protein engineering settings it **is** also important to consider [...]"
limitations: limitations: The authors have adequately addressed the limitations of this work in Section 4. I do not see any significant negative societal concerns in this work.

Comments: title: Thank you
Thank you
Follow up
Official Comment by Reviewer NoUX
comment: Thank you for your response. My concerns have been partially addressed, and I slightly raised my score. Good luck.
The authors have made a good faith effort to address most of my comments. I think the extended analysis and theory is also worth including in the final version. I have adjusted my score.
We sincerely thank the reviewer for their consideration. Is there anything more we could do to further improve the view of the referee of our paper?
We sincerely thank the reviewer for considering our arguments.

>> Regarding the second weakness, my concern was specifically about the experiments involving the use of BT loss to fine-tune pretrained models. Do you believe that BT loss could be effectively transferred across different models? If so, further testing of the loss function's application across various models could enhance the robustness of the findings.

We agree with the referee for fitness prediction tasks this is possible (still requires supervised fine-tuning). We are happy to add a fine-tuned model to our camera-ready version of the paper.
We'd like to follow up to see if the reviewer had a chance to reflect on our argument about the difference in zero-shot and supervised prediction tasks, as well as the fact that many of the tasks in ProteinGym would be inappropriate for our method (we added two datasets from proteinGym that are appropriate). Please let us know if we can further clarify.
Thank you for your thorough responses.

In my initial review, I misunderstood the concept of "fitness prediction". With the authors' clarifications, I've revised my understanding. I now recognize that the paper aims to offer a "flexible, interpretable, and general approach to modeling global epistasis", as stated by the authors.

Regarding the second weakness, my concern was specifically about the experiments involving the use of BT loss to fine-tune pretrained models. Do you believe that BT loss could be effectively transferred across different models? If so, further testing of the loss function's application across various models could enhance the robustness of the findings.

The authors also provided additional theoretical proofs and conducted further experiments using two datasets from ProteinGym.

In light of these improvements, I am raising my score from 3 to 5.
I appreciate the authors for dedicatedly addressing the review comments.

In my initial review comments, I had several misconceptions on the relationship between GFT and nonlinear global epistasis model. Thanks to the authors’ clarifications, I corrected my understanding. Now the trade-off between the sparsity of $f$ and $\\beta$ presented in Section 3.2 makes sense to me.

The application of the nonlinearity also makes somewhat sense to me now. It looks like an interesting application, which is based on the combination of multiple technical devices including GFT, BT model, and nonlinear modeling.

Although there are several parts where the presentation can be made better (in my minor comments, I mentioned that the background, method, and experiment sections are oftentimes mixed up), I raise my evaluation from 4 -> 5 given the interesting combination.