Paper: Beyond the training set: an intuitive method for detecting distribution shift in model-based optimization (ICLR-2024)

Reviewer comment (Evaluation):
It is not clear how the threshold on the OOD scores should be determined? Is there a systematic way to do this?

Author response:
In practice, we find the OOD scores to be robust across design algorithms (e.g., the scores are similar between the Adalead run in Figure 3 and the beam search run in Suppl. Figure 6). Therefore one could choose a threshold based on the results in the paper with confidence that the results will transfer. As we highlight in Section 2.4 “Selection Using OOD Scores”, a threshold choice will be application-specific and will depend on the user’s needs and level of risk-tolerance (e.g., dependent on the batch size, goals of the experiment, etc). A user can choose to stratify their risk by choosing multiple cutoffs, or minimize risk by choosing a very low cutoff.

Reviewer comment (Theory):
The binary classifier is trained to assign designed samples a different label from the training samples even if the designed samples are similar to the training samples. What are the downsides of this?

Author response:
For designed samples that are similar (read: in-distribution) to the training samples, the OOD scores will be low and will likely fall within the distribution of OOD scores for the training data itself. This is not an issue as the model learns to assign a classification probability close to 0.5 (or lower) for these data points. The downside of this is when all designed samples are in-distribution, this method will likely have limited use as it requires there to be some degree of shift. For any practical sequence design application, however, it is inevitable that there will be a shift as the goal is to design a sample with an improved measurement relative to what is observed in the training data.

Reviewer comment (Evaluation):
It is not discussed when the proposed technique could fail.

Author response:
The OOD classifier is likely to fail when the distributions p and q are very different from each other. If the gap between these two distributions is especially large, training a classifier to distinguish points from p and q can result in achieving perfect accuracy, but it may provide a poor estimate of the density ratio between them. This emphasizes the importance of choosing a design distribution q that overlaps with p to enable effective density ratio estimation through classification. Recent work has introduced a modified density ratio method using a telescoping product [1] to tackle this specific situatation. We will update the discussion section to address this challenge.

[1] Rhodes, Benjamin, Kai Xu, and Michael U. Gutmann. "Telescoping density-ratio estimation." Advances in neural information processing systems 33 (2020): 4905-4916

Reviewer comment (Evaluation):
While the paper provides qualitative insights and findings from experiments, more quantitative metrics that measure the efficacy, false positives, and false negatives of the method would give a clearer picture.

Author response:
> “While the paper provides qualitative insights and findings from experiments, more quantitative metrics that measure the efficacy, false positives, and false negatives of the method would give a clearer picture.”

Our regret plot (Figure 3c in the main text) provides a quantitative metric for comparing our method to the Deep Ensemble baseline. We also provide extensive ablations of the metric (Top K = 10, 50, 100, 250 and transduction percentiles = 90, 95, max) in the Appendix (Figures 7, 8, 9b, and 11). Standard classification metrics are not useful for our setting because the positive and negative examples used to train the OOD classifier are not necessarily true positives or true negatives. Therefore, classification metrics (e.g., AUC) on the test set does not tell you anything informative about the method. This why we chose Regret on the observed design as our primary metric for quantitatively evaluating our method.

Reviewer comment (Experiments):
It's not clear if the approach would be as effective across diverse domains outside of the ones presented in the paper.

Author response:
> “It's not clear if the approach would be as effective across diverse domains outside of the ones presented in the paper.”

We focus on applications to protein sequence design and protein structure prediction, which has been an area of intense interest in recent years. Covariate shift induced by the model is a major roadblock to the success of applying these methods in the real-world. This is evident by Figure 3a which shows running a common design method for 15 steps results in a 0% functional rate, despite high predicted model scores.

Demonstrating feedback covariate shift in diverse domains is challenging. For example, in commonly studied simulation environments like Design-Bench (Trabucco et. al. 2022) and FLEXS (Sinai et. al. 2020), the effect of covariate shift is minimal (we are happy to add a plot to the SI showing this). This is evident by the unexpectedly high performance of naive unconstrained search algorithms, such as the “max_x f(x)” in Design-Bench and an unconstrained version of Adalead in FLEXS. Using these methods “out-of-the-box” produces close to a 0% functional rate in real-world scenarios as we show in our manuscript. This suggests that these simulation settings do not effectively mimic the feedback covariate shift issue found in real-world scenarios. To bridge the gap between simulations and real-world conditions, we conducted a real-world experiment to specifically study this effect, which has largely gone unnoticed in the literature.

Reviewer comment (Theory):
This approach overly constricts the search space of allowable design samples.

Author response:
We show evidence that the OOD score tracks shift intensity (Figure 3a), which translates to how large the search space of allowable designs is. Our method doesn’t enforce any particular constraint but lets the user decide how restrictive they want the search space to be. Our OOD scores traverse the entire range from overly restrictive to relaxed. On the restrictive side, you can use an OOD score comparable to say the 75th or 90th percentile of the training data. In this case, we see close to a 100% functional rate (the 0-5 iteration regime in the top panel of Figure 3a). On the other end, we see at certain high OOD scores the sequence space contains mostly non-functional variants (the >10 iteration regime in the top panel of Figure 3a). Then there’s a range in between which the user can select from depending on their goals.

In practice, we find the OOD scores to be robust across design algorithms (e.g., the scores are similar between the Adalead run in Figure 3 and the beam search run in Suppl. Figure 6), so risk can be controlled reliably (e.g., through stratification based on OOD score).

Reviewer comment (Theory):
This means that design samples will be forced to be similar to training samples, even if distant samples in the input space could still have correct surrogate model predictions.

Author response:
We show evidence that the OOD score tracks shift intensity (Figure 3a), which translates to how large the search space of allowable designs is. Our method doesn’t enforce any particular constraint but lets the user decide how restrictive they want the search space to be. Our OOD scores traverse the entire range from overly restrictive to relaxed. On the restrictive side, you can use an OOD score comparable to say the 75th or 90th percentile of the training data. In this case, we see close to a 100% functional rate (the 0-5 iteration regime in the top panel of Figure 3a). On the other end, we see at certain high OOD scores the sequence space contains mostly non-functional variants (the >10 iteration regime in the top panel of Figure 3a). Then there’s a range in between which the user can select from depending on their goals.

In practice, we find the OOD scores to be robust across design algorithms (e.g., the scores are similar between the Adalead run in Figure 3 and the beam search run in Suppl. Figure 6), so risk can be controlled reliably (e.g., through stratification based on OOD score).

Reviewer comment (Theory):
The statement that distribution shifts in supervised regression “typically takes the form of covariate shift” is not well supported. Assuming covariate shifts is one of the most restrictive assumptions on shift since the model predictions p(y|x) is unchanged between training and test/design distribution. Can you expand on this argument for focusing on this type of error?

Author response:
In design problems, we only observe covariate shift, which means P(x) changes. P(y|x) stays constant because the underlying distribution relating sequences x to some property y is dictated by a fixed underlying biological process. This focus on covariate shift has been previously explored in the design literature (Fannjiang et. al. 2022). We will clarify Section 2.2 to clarify distribution shifts in *sequence design* problems typically take the form of covariate shift.

Reviewer comment (Experiments):
Several approaches to refining the search space in ML-guided design are discussed in Section 3 (first paragraph). Why were these not used as baseline comparisons in the experiments?

Author response:
The first paragraph of section three focuses on offline MBO algorithms used for sequence design. We are not devising a new MBO algorithm, instead, we are devising a complementary method that can be used alongside any MBO algorithm and only applies at the “selection” step after the initial sequences are designed. The selection problem can also be viewed as post-hoc hyperparameter tuning for an offline design algorithm (Fannjiang et al 2022). We highlight at the end of the first paragraph that “while these approaches require specific surrogate models or search strategies, the OOD classifier can be flexibly paired with any surrogate model or search technique.” 

It is a well-established practice in the field (in real-world applications) to distinguish between the design stage—where millions of sequences are generated with high predicted scores using various models and search algorithms—and the selection stage—where a small subset of the most promising candidates are chosen for experimental testing. These batches can range from 10^1 to 10^5 in size. Our work enables a selection procedure that achieves low regret, better than ensemble uncertainties and greedy selection. 


With respect to baselines more generally:

We agree that adding additional baseline comparisons would strengthen the paper, although developing such baselines for our use case is not always straightforward. Existing methods to handle distribution shift that we could compare against can either be classified as scoring rules used for identifying OOD points or from improving model robustness to distribution shift. We discuss the possibilities for developing baselines from each of these categories below:
* On scoring rules: Deep Ensemble scores are considered the state of the art for handling distribution shift and is largely considered the de facto standard, due to both its simplicity and robustness. For this reason, we focus our comparison against this baseline. With that said, conformal prediction has recently been applied to sequence design as a tool to correct covariate shift (Fannjiang et. al. 2022). This approach, however, in its most effective  form requires training (n + 1) x |y| models, where n is the number of test data points and |y| is the cardinality of y. Scalability improvements can enable training of n models, but for practical high-throughput sequence design problems, this remains a difficult computational challenge. We will update the text to explain this in more detail.

* On improving model robustness: It has previously been shown that model pretraining can improve robustness to distribution shift [1]. In the protein domain, large language models have been pretrained on hundreds of millions of protein sequences [2]. These models have shown impressive performance on a variety of downstream prediction tasks. In response to the reviewer’s comment, we have completed an additional experiment showcasing that improving model robustness (e.g., by using a pretrained LLM) does not address feedback covariate shift in the design setting. We used ProtBERT [2] (a pretrained protein language model) to predict our target property using linear probing [3], a method for adapting pretrained models to downstream regression problems by fixing the pretrained weights and fitting a linear head to predict the target property. We used an identical train/test split to the CNN described in the manuscript and computed predictions on the designed data. We find identical behavior to the CNN – model scores and model error go up as a function of the distribution shift. There is a massive drop in performance from a random holdout (Spearman rho = 0.75) to the design set (Spearman rho = -0.55). We plan on updating the manuscript with a figure showcasing this result. We reiterate that our method is broadly applicable since it can be integrated with any design algorithm or predictive model, so if future models prove to be robust to design-induced distribution shift, our method can be used to further improve this result. 

[1] Hendrycks, Dan, Kimin Lee, and Mantas Mazeika. "Using pre-training can improve model robustness and uncertainty." International conference on machine learning. PMLR, 2019

[2] Elnaggar, Ahmed, et al. "Prottrans: Toward understanding the language of life through self-supervised learning." IEEE transactions on pattern analysis and machine intelligence 44.10 (2021): 7112-7127.

[3] Kumar, Ananya, et al. "Fine-tuning can distort pretrained features and underperform out-of-distribution." arXiv preprint arXiv:2202.10054 (2022)