[Paper Review] Bidirectional Learning for Offline Infinite-width Model-based Optimization

Can, Yingxue Zhang+3 more|arXiv (Cornell University)|Sep 15, 2022

Advanced Neural Network Applications7 citations

TL;DR

This paper proposes Bidirectional Learning for Offline Infinite-width Model-based Optimization (BDI), a method that improves offline model-based optimization by training bidirectional mappings between high-scoring designs and a static dataset to mitigate out-of-distribution failures. By leveraging infinite-width deep neural networks and the Neural Tangent Kernel (NTK), BDI achieves a closed-form loss for accurate design updates, outperforming state-of-the-art methods on multiple tasks in Design-Bench.

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, DNA sequences, and proteins. Recent approaches train a deep neural network (DNN) 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 poor designs. To mitigate this problem, we propose BiDirectional learning for offline Infinite-width model-based optimization (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 from the static dataset into the high-scoring designs, which effectively mitigates the out-of-distribution problem. 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 at https://github.com/GGchen1997/BDI.

Motivation & Objective

• To address the out-of-distribution problem in offline model-based optimization, where gradient-based optimization on a proxy function often generates low-quality designs.
• To improve design quality by ensuring high-scoring designs can accurately predict scores of the static dataset, and vice versa, through bidirectional knowledge distillation.
• To overcome the intractable loss problem in finite-width models by adopting an infinite-width DNN with a closed-form loss via the Neural Tangent Kernel (NTK).
• To demonstrate the effectiveness of bidirectional learning in real-world design tasks such as protein, DNA, and material design, where objective function evaluations are expensive.

Proposed method

• BDI introduces a forward mapping from low-scoring (static) designs to high-scoring designs, and a backward mapping from high-scoring designs to the static dataset.
• The backward mapping ensures that high-scoring designs can reconstruct the scores of the static dataset, thereby distilling more information and improving alignment with the training distribution.
• A predefined target score, larger than the maximum score in the static dataset, is used to guide the optimization of high-scoring designs.
• To enable exact loss computation, BDI adopts an infinite-width DNN model and uses the Neural Tangent Kernel (NTK) to derive a closed-form expression for the loss function.
• The bidirectional loss is minimized jointly, allowing the high-scoring designs to be updated via gradient descent using the NTK-based closed-form gradient.
• The method is applied in a bi-level optimization framework, where the high-scoring designs are optimized based on their ability to predict both their own scores and the static dataset’s scores.

Experimental results

Research questions

• RQ1 Can bidirectional learning between high-scoring designs and a static dataset improve generalization and reduce out-of-distribution failures in offline model-based optimization?
• RQ2 Does using an infinite-width DNN with Neural Tangent Kernel (NTK) enable a closed-form loss that leads to more accurate and stable design updates compared to finite-width approximations?
• RQ3 How does BDI compare to state-of-the-art methods in terms of performance across diverse offline design tasks such as protein and DNA sequence optimization?
• RQ4 To what extent does the backward mapping, which enforces high-scoring designs to predict static dataset scores, improve alignment with the true objective function?

Key findings

• BDI achieves state-of-the-art performance on multiple tasks in the Design-Bench benchmark, outperforming existing methods in offline model-based optimization.
• The use of infinite-width DNNs with NTK enables a closed-form loss function, significantly improving design quality over finite-width models with approximate losses.
• The backward mapping effectively mitigates the out-of-distribution problem by ensuring high-scoring designs are consistent with the static dataset, reducing overestimation of scores.
• Experiments show that bidirectional knowledge distillation leads to higher ground-truth scores for generated designs, even when the proxy function overestimates performance.
• The method generalizes well across diverse domains, including protein design (GFP), DNA sequence design, and robotic design, demonstrating robustness and scalability.