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: This paper proposes robust model adaptation, which first do a pretraining with adversarial weight perturbation. Then the final result is obtained by iteratively doing model adaptation on the learned model, with a local smoothness penalty objective. The proposed results shows improvements over existing methods.