[논문 리뷰] Bidirectional Learning for Offline Infinite-width Model-based Optimization

Can, Yingxue Zhang외 3명|arXiv (Cornell University)|2022. 09. 15.

Advanced Neural Network Applications인용 수 7

한 줄 요약

이 논문은 오프라인 무한 폭 모델 기반 최적화를 위한 이중성 학습(Bidirectional Learning)을 제안한다(BDI). 이는 고성능 설계와 정적 데이터셋 간의 이중성 매핑을 학습하여 분포 외 실패 문제를 완화함으로써 오프라인 모델 기반 최적화의 성능을 햖थ한다. 무한 폭 딥 네ural 네트워크와 신경 탄성 커널(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)을 통해 폐쇄형 손실를 갖는 무한 폭 딥 네ural 네트워크를 채택한다.
• 목적 함수 평가가 비용이 많이 드는 실제 설계 과제들, 예를 들어 단백질, DNA, 재료 설계에서 이중성 학습의 효과를 입증한다.

제안 방법

• BDI는 저성능(정적) 설계에서 고성능 설계로의 정방향 매핑과 고성능 설계에서 정적 데이터셋으로의 역방향 매핑을 도입한다.
• 역방향 매핑은 고성능 설계가 정적 데이터셋의 점수를 재구성할 수 있도록 하여 더 많은 정보를 전달하고 학습 분포와의 일치도를 향상시킨다.
• 정적 데이터셋의 최대 점수보다 큰 사전 정의된 목표 점수를 사용하여 고성능 설계의 최적화를 이끌어낸다.
• 정확한 손실 계산을 가능하게 하기 위해, BDI는 무한 폭 딥 네ural 네트워크 모델을 채택하고, 신경 탄성 커널(NTK)을 활용해 손실 함수에 대한 폐쇄형 표현을 유도한다.
• 이중성 손실는 함께 최소화되어, NTK 기반 폐쇄형 기울기를 사용해 고성능 설계를 기울기 하강법으로 업데이트할 수 있다.
• 이 방법은 이중 최적화 프레임워크에 적용되며, 고성능 설계는 자신의 점수와 정적 데이터셋의 점수를 모두 예측할 수 있는 능력에 기반해 최적화된다.

실험 결과

연구 질문

• RQ1 고성능 설계와 정적 데이터셋 간의 이중성 학습이 오프라인 모델 기반 최적화에서 일반화 능력을 향상시키고 분포 외 실패 문제를 감소시키는가?
• RQ2 신경 탄성 커널(NTK)을 활용한 무한 폭 딥 네ural 네트워크가 폐쇄형 손실를 제공함으로써, 근사 손실를 사용하는 유한 폭 모델보다 더 정확하고 안정적인 설계 업데이트를 이끌어내는가?
• RQ3 단백질 및 DNA 서열 최적화와 같은 다양한 오프라인 설계 과제에서 BDI는 최신 기술 대비 성능 면에서 어떻게 비교되는가?
• RQ4 정적 데이터셋 점수를 예측하도록 강제하는 역방향 매핑이 진정한 목적 함수와의 일치도를 얼마나 향상시키는가?

주요 결과

• BDI는 Design-Bench 벤치마크에서 다양한 작업에서 최신 기술을 초월하는 최고 성능을 달성한다.
• NTK를 활용한 무한 폭 딥 네ural 네트워크는 근사 손실를 사용하는 유한 폭 모델보다 설계 품질을 크게 향상시키는 폐쇄형 손실를 제공한다.
• 역방향 매핑은 고성능 설계가 정적 데이터셋과 일치하도록 함으로써 분포 외 문제를 효과적으로 완화하며, 점수 과대평가 문제를 줄인다.
• 실험 결과, 이중성 지식 전이가 프록시 함수가 성능을 과대평가하더라도 생성된 설계의 실제 점수를 높이는 데 기여한다.
• 이 방법은 단백질 설계(GFP), DNA 서열 설계, 로봇 설계 등 다양한 도메인에 걸쳐 잘 일반화되며, 강건성과 확장성을 입증한다.