RoMA: Robust Model Adaptation for Offline Model-based Optimization

1 Introduction . Designing new objects with the desired property is a fundamental problem in a wide range of real-world domains , including chemistry [ 13 ] , biology [ 5 , 41 ] , robotics [ 3 , 28 ] , and aircraft design [ 20 ] . Unfortunately , many of them often require evaluating an expensive objective function , e.g. , synthesizing and measuring the fluorescence of the protein [ 5 ] . Model-based optimization ( MBO ) is a powerful framework to circumvent this issue , based on an approximation of the expensive evaluation of objective functions by a cheap proxy model . Specifically , MBO optimizes a solution with respect to the proxy model instead of the costly objective function . For obtaining a better proxy model , predominant works have typically assumed active ( or online ) learning scenarios , i.e. , they allow for updating the model interactively with new queries of choice to approximate the ground-truth objective function better around optima [ 22 , 37 , 39 , 40 , 43 , 45 , 46 ] . Recently , researchers have also proposed offline versions of MBO algorithms [ 5 , 8 , 11 , 25 , 52 ] which build the proxy model from a given static dataset , i.e. , input-output pairs collected from the objective function prior to building the proxy model . Such frameworks are beneficial for various real-world scenarios where the new interactive access to the objective functions may accompany serious danger or expensive cost for evaluation , e.g. , drug discovery or aircraft design . The common challenge to be addressed for offline MBO is to construct a proxy model that yields an accurate approximation of the 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . true objective function outside the training dataset , as the optimization procedure mostly leads the solution to be out of the dataset and causes severe overestimation due to a domain shift [ 12 , 48 ] . Intriguingly , Fu and Levine [ 11 ] , Trabucco et al . [ 52 , 53 ] found the gradient-based optimization of solutions to be effective for offline MBO . Such a procedure is attractive since it can utilize gradient information of the proxy model , e.g. , a deep neural network ( DNN ) , to generate fine-grained solutions . However , due to the highly non-smooth nature of DNNs , such a gradient-based optimization may generate adversarial examples [ 51 ] whose score deviates significantly between the DNN-based proxy model and the true objective function . Existing works have attempted to regularize this behavior using conservative training [ 52 ] or normalized maximum likelihood [ 11 ] . In this paper , we investigate a more direct way to regularize such a non-smooth nature of DNN . Contribution . We introduce a new framework , coined robust model adaptation ( RoMA ) , for offline MBO to optimize the solution via gradient-based updates from the DNN proxy model . Our main idea is to utilize a local smoothness prior for alleviating the brittleness of DNN at adversarial examples . To be specific , RoMA consists of two steps : ( a ) a pre-training strategy to robustly train the proxy model and ( b ) a novel adaptive adjustment procedure of the proxy model to have robust estimates for a specific set of candidate solutions . We provide an overall illustration of our framework in Figure 1 . ( a ) We first introduce a pre-training procedure of the proxy model to be locally smooth at inputs and eventually becomes robust to the gradient-based input change . The novel component of the pretraining procedure is an additional regularization to force a flat loss landscape with respect to the model parameters , rather than only employing regularization based on input-level perturbation into the training objective . Such consideration of the flatness to the model parameters is effective in improving the robustness of the proxy model to the gradient-based updates . ( b ) To strengthen the accurate update of solution candidates from the trained proxy model , we also present a novel adaptive adjustment strategy of the proxy model to enhance the local smoothness at inputs outside the training data . Specifically , we propose to adjust the model at given solution candidates by adding a small perturbation to the model weights.After that , an update of solution candidates occurs from the adjusted model rather than the original DNN proxy model . We discover that the proposed adjustment strategy aptly mitigates the fragility of the proxy model at gradient-basead updates to given inputs and improves the performance . We also emphasize that the adaptive strategy is worthwhile in offline MBO as the expressive power of the proxy model is limited under the fixed number of data and parameters ; it is difficult to sufficiently regularize all inputs with the domain shift from the dataset without any adjustments . Algorithm 1 Robust model adaptation ( RoMA ) 1 : for k = 0 to K do ⊲ Step 1 . Pre-training with weight perturbations . 2 : Find θ̃ ← argmax θ̃∈B ( θ ) E ( x , y ) ∼D , δ∼N ( 0 , σ ) [ ( f ( x+ δ ; θ̃ ) − y ) 2 ] via PGD . 3 : Reparameterize θ̃ using θ such that θ̃ℓ = θℓ + φℓ for some fixed φℓ and ℓ = 1 , . . . , L. 4 : Minimize E ( x , y ) ∼D , δ∼N ( 0 , σ ) [ ( f ( x+ δ ; θ̃ ) − y ) 2 ] over θ using gradient-based update . 5 : end for 6 : Set θ−1 ← θ and choose x ( 0 ) from the training dataset D. 7 : for t = 0 to T − 1 do ⊲ Step 2 . Optimizing the solution with iterative model adaptation . 8 : Find θt ← argminθ̃∈B ( θ ) [ ||∇xf ( x ; θ̃ ) ||2 ∣∣∣ x=x ( t ) + α ( f ( x ( t ) ; θ̃ ) − f ( x ( t ) ; θt−1 ) ) 2 ] via PGD . 9 : Compute x ( t+1 ) from x ( t ) using gradient-based update . 10 : end for 11 : Output the final solution x ( T ) . We extensively verify the effectiveness of our RoMA under Design-Bench [ 53 ] , an offline MBO benchmark on diverse domains such as chemistry , biology , and robotics . In particular , our method achieves the best result in 4 out of 6 tasks when compared to strong baselines [ 5 , 8 , 11 , 25 , 52 ] . For the remaining tasks , RoMA at least achieves the runner-up results . When averaging over all tasks , RoMA achieves the normalized 100th percentile score1 of 1.705 , while the second-best baseline [ 11 ] achieves 1.687 . Intriguingly , our framework dramatically improves the 50th percentile scores , demonstrating how it generates diverse solutions with superior quality . 2 RoMA : Robust model adaptation . 2.1 Overview of RoMA . To describe our framework , we define our problem as follows . Consider a real-valued vector x ∈ RL , e.g. , protein sequence , and a real-valued objective function f∗ ( x ) ∈ R , e.g. , fluorescence of a protein sequence or the morphology of the robot . We are interested in finding an optimal input x∗ ∈ RL that maximizes the objective function , i.e. , x∗ = argmaxx∈RL f ∗ ( x ) . Rather than having access to the objective function , we are given a static dataset D = { ( xi , yi ) } d i=1 of input-output pairs ( xi , yi ) queried from the objective function , i.e. , yi = f ∗ ( xi ) . Since the true objective function is unknown , our RoMA framework approximates it using a surrogate model f ( x ; θ ) parameterized as a deep neural network ( DNN ) with weights θ . Then one can use gradient-based updates to find a fine-grained input maximizing the surrogate model . However , without any regularization , DNNs are prone to overfitting and naïve gradient ascent algorithms can easily find adversarial examples whose score deviates significantly between the surrogate model and the true objective function . Our main contribution lies in resolving this issue by regularizing the smoothness of the surrogate model , i.e. , decreasing the sensitivity of the model with respect to input perturbations . Such a local smoothness prior for regularizing DNNs has been successfully used to enhance the generalization of DNNs in various scenarios , e.g. , adversarial robustness , generative models , and uncertainty estimation [ 1 , 7 , 9 , 14 , 29 , 31 , 35 , 38 , 47 , 51 , 56 , 57 ] . Especially , we note the striking resemblance between our framework to the adversarial training methods [ 7 , 51 , 56 ] ; they seek to prevent the existence of adversarial examples generated from gradient-based optimization of inputs . In this respect , one may expect the smoothness prior to be beneficial for our offline MBO framework . To be specific , our RoMA is a two-stage framework , described as follows . Step 1 . Train the proxy model on the dataset D to approximate the true objective function with Gaussian smoothing of inputs under worst-case weight perturbations . Step 2 . For time-steps t = 0 , . . . , T − 1 , update the solution x ( t ) after adapting the output of the proxy model to be locally smooth at the solution x ( t ) . We provide the detailed description of our framework in Algorithm 1 . 1The percentile score is computed from 128 solutions proposed by the offline MBO algorithms . In the rest of this section , we explain each component of RoMA in detail . In Section 2.2 , we describe how the RoMA pre-trains the surrogate model with Gaussian smoothing ( Step 1. ) . In Section 2.3 , we explain the procedure of updating the solution while simultaneously adapting the proxy model to yield input-level smoothness at the intermediate solution ( Step 2. ) . Finally , in Section 2.4 , we describe some other algorithmic components in RoMA , such as an approach to handling discrete inputs or a procedure to select hyperparameters without the true objective function . 2.2 Pre-training with weight perturbations . We first describe our pre-training strategy for the DNN-based proxy model f ( x ; θ ) to approximate the true objective function f∗ ( x ) on training dataset D. We focus on regularizing the training with Gaussian smoothing of inputs under worst-case weight perturbations [ 55 ] . To this end , we optimize an L-layered DNN with weight θ as the proxy model to minimize the approximation loss as follows : L ( θ ) : = max θ̃∈B ( θ ) [ E ( x , y ) ∼D , δ∼N ( 0 , σ ) [ ( f ( x+ δ ; θ̃ ) − y ) 2 ] ] , ( 1 ) where δ is a random variable sampled from the Gaussian distribution N ( 0 , σ2 ) with mean 0 and standard deviation σ . Furthermore , B ( θ ) is the set of perturbed parameters of the DNN that are near the parameter θ , defined as follows : B ( θ ) : = { θ̃ : ‖θℓ − θ̃ℓ‖F ≤ ε · ‖θℓ‖F ∀ ℓ ∈ { 1 , · · · , L } } . Here , ε > 0 indicates the maximum magnitude of the perturbation , θℓ is the parameter of the ℓ-th layer and ‖·‖F denotes the Frobenious norm of a matrix . To solve the inner maximization problem in ( 1 ) , we utilize the projected gradient descent ( PGD ) [ 32 ] , similar to Wu et al . [ 55 ] . We remark that our pre-training uses Gaussian noise for input perturbation , while Wu et al . [ 55 ] uses a bi-level worst-case input perturbation . We found that using the Gaussian noise leads to a much more stable training for our problem . We provide more details of this optimization procedure in the Appendix A . Such an input-level Gaussian smoothing , i.e. , augmenting inputs with additive noise sampled from a Gaussian distribution , has shown to effectively regularize the smoothness of DNNs in various domains [ 27 , 30 ] and moreover improve the adversarial robustness [ 7 ] . Furthermore , training the model under the worst-case perturbation of weights induces a flat loss landscape with respect to weights , which has shown to be beneficial for generalization of the training objective outside the training dataset [ 6 , 10 , 21 , 26 , 36 ] . In particular , Stutz et al . [ 49 ] and Wu et al . [ 55 ] empirically observed a strong correlation between generalization capability of the model robustness at adversarial examples and the flat loss landscape of the model parameters .

Summary: RoMA is a framework that leverages the local smoothness prior to robustly approximate black-box objective functions. RoMA consists of two steps: 1, train the proxy models (neural networks) to approximate black-box objective functions by perturbing original inputs with Gaussian inputs and minimizing worst-case weight perturbations; 2, iteratively adapting the proxy model to enforce the local smoothness of loss landscape around the input x(t). This paper devises two separate loss functions for optimizing neural networks in these two steps. As a result, this method serves to avoid adversarial optimized inputs. In the experiment section, this method is compared to the state-of-the-art methods on the Design-bench dataset. Results show that this method achieves state-of-the-art performance across a variety of tasks.