[论文解读] Bidirectional Learning for Offline Infinite-width Model-based Optimization

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

Advanced Neural Network Applications被引用 7

一句话总结

本文提出双向学习用于离线无限宽度基于模型的优化（BDI），通过在高分设计与静态数据集之间训练双向映射，缓解分布外失败问题，从而提升离线基于模型的优化性能。BDI 利用无限宽度深度神经网络与神经正切核（NTK），实现精确的设计更新的闭式损失，相较于现有最先进方法，在 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.

研究动机与目标

• 为解决离线基于模型优化中的分布外问题，即在代理函数上进行基于梯度的优化常生成低质量设计。
• 通过双向知识蒸馏，确保高分设计能准确预测静态数据集的得分，反之亦然，从而提升设计质量。
• 通过采用无限宽度深度神经网络与基于神经正切核（NTK）的闭式损失，克服有限宽度模型中难以处理的损失问题。
• 在蛋白质、DNA 与材料设计等真实世界设计任务中，展示双向学习的有效性，其中目标函数评估成本高昂。

提出的方法

• BDI 引入从前置低分（静态）设计到高分设计的前向映射，以及从高分设计到静态数据集的反向映射。
• 反向映射确保高分设计能重建静态数据集的得分，从而蒸馏更多信息，提升与训练分布的一致性。
• 使用一个大于静态数据集中最大得分的预设目标得分，以引导高分设计的优化。
• 为实现精确损失计算，BDI 采用无限宽度深度神经网络模型，并利用神经正切核（NTK）推导损失函数的闭式表达式。
• 双向损失联合最小化，使高分设计可通过基于 NTK 的闭式梯度进行梯度下降更新。
• 该方法被应用于双层优化框架中，高分设计基于其预测自身得分与静态数据集得分的能力进行优化。

实验结果

研究问题

• RQ1 在离线基于模型的优化中，高分设计与静态数据集之间的双向学习是否能提升泛化能力并减少分布外失败？
• RQ2 使用无限宽度 DNN 与神经正切核（NTK）是否能实现闭式损失，从而相比有限宽度近似方法带来更精确、更稳定的更新？
• RQ3 在蛋白质与 DNA 序列优化等多样化的离线设计任务中，BDI 与最先进方法相比性能如何？
• RQ4 反向映射（强制高分设计预测静态数据集得分）在多大程度上提升了与真实目标函数的一致性？

主要发现

• BDI 在 Design-Bench 基准的多个任务中达到最先进性能，优于现有离线基于模型优化方法。
• 采用无限宽度 DNN 与 NTK 实现了闭式损失函数，显著提升了有限宽度模型使用近似损失时的设计质量。
• 反向映射通过确保高分设计与静态数据集一致，有效缓解了分布外问题，减少了得分的过度估计。
• 实验表明，双向知识蒸馏使生成设计的实测得分更高，即使代理函数高估了性能。
• 该方法在多种领域（包括蛋白质设计（GFP）、DNA 序列设计与机器人设计）中表现出良好泛化能力，展示了鲁棒性与可扩展性。