Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design
Created by
Haebom
저자
Xingyu Su, Xiner Li, Masatoshi Uehara, Sunwoo Kim, Yulai Zhao, Gabriele Scalia, Ehsan Hajiramezanali, Tommaso Biancalani, Degui Zhi, Shuiwang Ji
개요
본 논문은 생체 분자 설계에서 보상 기반 생성을 위해 확산 모델을 미세 조정하는 문제를 다룹니다. 확산 모델은 복잡하고 고차원적인 데이터 분포를 모델링하는 데 매우 효과적이지만, 실제 응용 프로그램에서는 물리 기반 시뮬레이션이나 과학적 지식에 기반한 보상 함수와 같이 미분 불가능한 보상 함수에 대한 최적화를 필요로 하는 경우가 많습니다. 기존의 강화 학습(RL) 기반 방법은 불안정성, 낮은 샘플 효율성, 모드 붕괴 문제를 겪는 경우가 많습니다. 본 연구에서는 임의의 보상 함수에 대해 확산 모델을 최적화할 수 있는 반복적인 증류 기반 미세 조정 프레임워크를 제안합니다. 이 방법은 문제를 정책 증류로 바꾸어, 롤인 단계에서 오프-폴리시 데이터를 수집하고, 롤아웃 단계에서 보상 기반 소프트 최적 정책을 시뮬레이션하고, 시뮬레이션된 소프트 최적 정책과 현재 모델 정책 간의 KL 발산을 최소화하여 모델을 업데이트합니다. 오프-폴리시 공식과 KL 발산 최소화는 기존 RL 기반 방법에 비해 훈련 안정성과 샘플 효율성을 향상시킵니다. 실험 결과는 단백질, 소분자 및 조절 DNA 설계의 다양한 작업에서 제안된 방법의 효과와 우수한 보상 최적화를 보여줍니다.
시사점, 한계점
•
시사점:
◦
임의의 보상 함수에 대해 확산 모델을 효과적으로 미세 조정하는 새로운 프레임워크 제시
◦
오프-폴리시 학습과 KL 발산 최소화를 통해 기존 RL 기반 방법의 한계인 불안정성과 낮은 샘플 효율성 문제 해결
