[論文レビュー] 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）を提案する。無限幅の深層ニューラルネットワークとニューラルタングェントカーネル（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.

研究の動機と目的

• オフラインモデルベース最適化における分布外問題に対処すること。具体的には、プロキシ関数に対する勾配ベース最適化がしばしば低品質な設計を生成する問題を解決すること。
• 双方向知識蒸留を通じて、高得点デザインが静的データセットのスコアを正確に予測でき、逆に静的データセットのスコアが高得点デザインのスコアを正確に予測できるようにすることで、設計品質を向上させること。
• 有限幅モデルにおける計算不能な損失問題を克服するために、閉形式損失を有する無限幅DNNを採用し、ニューラルタングェントカーネル（NTK）を用いて損失関数の閉形式表現を導出すること。
• タンパク質、DNA、材料設計などの実世界の設計タスクにおいて、目的関数の評価が高コストであることを踏まえ、双方向学習の有効性を実証すること。

提案手法

• BDIは、低得点（静的）デザインから高得点デザインへの前向きマッピングと、高得点デザインから静的データセットへの後向きマッピングを導入する。
• 後向きマッピングにより、高得点デザインが静的データセットのスコアを再構築できるようになり、より多くの情報を抽出し、トレーニング分布との整合性を高める。
• 静的データセットに含まれる最大スコアよりも大きな事前に定義されたターゲットスコアを用い、高得点デザインの最適化をガイドする。
• 正確な損失計算を可能にするために、BDIは無限幅DNNモデルを採用し、ニューラルタングェントカーネル（NTK）を用いて損失関数の閉形式表現を導出する。
• 双方向損失を同時に最小化することで、NTKに基づく閉形式勾配を用いて高得点デザインを勾配降下法で更新可能となる。
• 本手法は二段階最適化フレームワークに適用され、高得点デザインは自身のスコアと静的データセットのスコアの両方を予測できる能力に基づいて最適化される。

実験結果

リサーチクエスチョン

• RQ1 高得点デザインと静的データセットの間で双方向学習を実施することで、オフラインモデルベース最適化における一般化性能が向上し、分布外失敗が軽減されるか？
• RQ2 無限幅DNNにニューラルタングェントカーネル（NTK）を適用することで、閉形式損失が得られ、有限幅近似と比較してより正確で安定した設計更新が可能になるか？
• RQ3 タンパク質やDNA配列最適化などの多様なオフライン設計タスクにおいて、BDIは最先端の手法と比較してどの程度の性能を発揮するか？
• RQ4 高得点デザインが静的データセットのスコアを予測するよう強制する後向きマッピングは、真の目的関数との整合性をどの程度向上させるか？

主な発見

• BDIはDesign-Benchベンチマークの複数のタスクで最先端の性能を達成し、既存の手法を上回る性能を発揮した。
• NTKを用いた無限幅DNNの採用により、閉形式損失関数が得られ、有限幅モデルにおける近似損失と比較して、設計品質が顕著に向上した。
• 後向きマッピングにより、高得点デザインが静的データセットと整合的になるよう保証され、スコアの過大評価が軽減され、分布外問題が効果的に緩和された。
• 実験の結果、双方向知識蒸留により、プロキシ関数が性能を過大評価しても、生成された設計の真のスコアが高くなる傾向が確認された。
• 本手法はタンパク質設計（GFP）、DNA配列設計、ロボット設計など多様な分野にわたり、一般化性能が高く、堅牢性とスケーラビリティを示した。