Bidirectional Learning for Offline Infinite-width Model-based Optimization

Can Chen, Yingxue Zhang, Jie Fu, Xue Liu, Mark Coates

Keywords: offline model-based optimization, infinite-width neural network, neural tangent kernel, bi-level optimization, out-of-distribution

TL;DR: We propose bidirectional learning between the high-scoring designs and the static dataset to effectively mitigate the out-of-distribution problem and adopt an infinite-width DNN model to yield a closed-form loss function.

Abstract: In offline model-based optimization, we strive to maximize a black-box objective function by only leveraging a static dataset of designs and their scores. This problem setting arises in numerous fields including the design of materials, robots, DNAs, proteins, etc. Recent approaches train a deep neural network (DNN) model on the static dataset to act as a proxy function, and then perform gradient ascent on the existing designs to obtain potentially high-scoring designs. This methodology frequently suffers from the out-of-distribution problem where the proxy function often returns adversarial designs. To mitigate this problem, we propose $\textit{\textbf{B}i\textbf{D}irectional learning for offline \textbf{I}nfinite-width model-based optimization}~(\textbf{BDI})$. BDI consists of two mappings: the forward mapping leverages the static dataset to predict the scores of the high-scoring designs, and the backward mapping leverages the high-scoring designs to predict the scores of the static dataset. The backward mapping, neglected in previous work, can distill more information of the static dataset into the high-scoring designs, which effectively mitigates the out-of-distribution problem. Yet, for a finite-width DNN model, the loss function of the backward mapping is intractable and only has an approximate form, which leads to a significant deterioration of the design quality. We thus adopt an infinite-width DNN model and propose to employ the corresponding neural tangent kernel to yield a closed-form loss for more accurate design updates. Experiments on various tasks verify the effectiveness of BDI. The code is available [here](https://github.com/GGchen1997/BDI).

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2209.07507/code)