Offline Model-Based Optimization by Learning to Rank

TL;DR

This paper targets offline model-based optimization (MBO), where only a fixed dataset of designs and scores is available. It shows that minimizing mean squared error (MSE) for surrogate models yields weak guidance in finding top designs due to out-of-distribution (OOD) errors, and that preserving the relative ordering of designs is the key objective. To address this, the authors introduce RaM, a ranking-based surrogate trained with learning-to-rank (LTR) losses and enhanced by data augmentation and output adaptation, together with an AUPCC-based metric that aligns with the design-optimization goal. Theoretical analysis provides a generalization bound for ranking losses, and extensive experiments on Design-Bench demonstrate RaM's superiority over twenty offline MBO methods, with strong evidence that ranking-focused methods and OOD-AUPCC correlate more with final design quality than MSE-based approaches. This work advances offline MBO by reframing the objective, offering a practical algorithm, theoretical guarantees, and broad empirical validation, with available code for reproducibility.

Abstract

Offline model-based optimization (MBO) aims to identify a design that maximizes a black-box function using only a fixed, pre-collected dataset of designs and their corresponding scores. A common approach in offline MBO is to train a regression-based surrogate model by minimizing mean squared error (MSE) and then find the best design within this surrogate model by different optimizers (e.g., gradient ascent). However, a critical challenge is the risk of out-of-distribution errors, i.e., the surrogate model may typically overestimate the scores and mislead the optimizers into suboptimal regions. Prior works have attempted to address this issue in various ways, such as using regularization techniques and ensemble learning to enhance the robustness of the model, but it still remains. In this paper, we argue that regression models trained with MSE are not well-aligned with the primary goal of offline MBO, which is to select promising designs rather than to predict their scores precisely. Notably, if a surrogate model can maintain the order of candidate designs based on their relative score relationships, it can produce the best designs even without precise predictions. To validate it, we conduct experiments to compare the relationship between the quality of the final designs and MSE, finding that the correlation is really very weak. In contrast, a metric that measures order-maintaining quality shows a significantly stronger correlation. Based on this observation, we propose learning a ranking-based model that leverages learning to rank techniques to prioritize promising designs based on their relative scores. We show that the generalization error on ranking loss can be well bounded. Empirical results across diverse tasks demonstrate the superior performance of our proposed ranking-based models than twenty existing methods.

Figures (2)

• Figure 1: Illustration of (a) OOD issue of regression-based models and (b) order-preserving ranking-based models. In (a), the regression-based method searches into suboptimal regions. Prior works focus on high OOD-MSE, while in this work, we point out that it is caused by the OOD error in preserving order. In (b), although the surrogate model also has high OOD-MSE, it can maintain the order, thus resulting in good design candidates.
• Figure 2: Scatter plots of five surrogate models (each trained using eight seeds) on the two tasks of D'Kitty and TF-Bind-8, where the $y$-axis denotes the rank of the 100th percentile score, and the $x$-axis denotes the rank of the metric in the OOD region, i.e., OOD-MSE or OOD-AUPCC. The Spearman correlation coefficients are also calculated, as shown in the title of each subfigure.

Theorems & Definitions (4)

• Theorem 1: Equivalence of Optima for Order-Preserving Surrogates
• proof
• Definition 1: AUPCC for Offline MBO
• Theorem 2: Generalization Error Bound for LTR listwise-general