[논문 리뷰] Bridging Model-Based Optimization and Generative Modeling via Conservative Fine-Tuning of Diffusion Models

Masatoshi Uehara, Yulai Zhao외 6명|arXiv (Cornell University)|2024. 05. 30.

Advanced Mathematical Modeling in Engineering인용 수 4

한 줄 요약

이 논문은 오프라인 설계 최적화를 향상시키기 위해 확산 모델과 보상 모델링을 융합한 보수적인 피지컬 튜닝 프레임워크인 BRAID를 제안한다. 보상 모델링에서 분포 외 영역을 벌어지게 하고 KL 발산 제약 조건을 적용하는 双중 보수적 접근을 통해, 확산 모델이 관측된 최고의 오프라인 설계를 초월하면서도 유효하거나 악성 생성물을 방지할 수 있도록 한다.

ABSTRACT

AI-driven design problems, such as DNA/protein sequence design, are commonly tackled from two angles: generative modeling, which efficiently captures the feasible design space (e.g., natural images or biological sequences), and model-based optimization, which utilizes reward models for extrapolation. To combine the strengths of both approaches, we adopt a hybrid method that fine-tunes cutting-edge diffusion models by optimizing reward models through RL. Although prior work has explored similar avenues, they primarily focus on scenarios where accurate reward models are accessible. In contrast, we concentrate on an offline setting where a reward model is unknown, and we must learn from static offline datasets, a common scenario in scientific domains. In offline scenarios, existing approaches tend to suffer from overoptimization, as they may be misled by the reward model in out-of-distribution regions. To address this, we introduce a conservative fine-tuning approach, BRAID, by optimizing a conservative reward model, which includes additional penalization outside of offline data distributions. Through empirical and theoretical analysis, we demonstrate the capability of our approach to outperform the best designs in offline data, leveraging the extrapolation capabilities of reward models while avoiding the generation of invalid designs through pre-trained diffusion models.

연구 동기 및 목표

• 보상 모델이 분포 외 영역에서 잘못된 방향을 유도하는 문제를 해결하기 위해 오프라인 모델 기반 최적화에서의 과최적화 문제에 대응한다.
• 정적 오프라인 데이터셋에서 학습된 보상 함수를 사용하여 사전 훈련된 확산 모델을 피지컬 튜닝하여 생성 모델링과 모델 기반 최적화를 조율한다.
• 관측된 데이터 분포를 초월한 최고의 설계를 초월하는 외삽 능력을 유지하면서도 유효한 설계 공간에 충실한 방법을 개발한다.
• 피지컬 튜닝 중에 불확실성 인식 페널티와 KL 정규화를 통합하여 분포 이탈에 대한 강건성을 확보한다.

제안 방법

• 오프라인 데이터를 사용하여 보수적인 보상 모델을 훈련시키며, 분포 외 입력에 대해 페널티를 증가시키는 불확실성 정량화 항목을 포함한다.
• 부드러운 엔트로피 정규화된 마르코프 결정 과정 프레임워크를 사용하여 보수적 정규화를 피지컬 튜닝 목적함에 포함시키는 이론적 근거를 제시한다.
• 강화 학습을 통해 보수적인 보상 모델을 최적화하면서도, 생성을 유효한 설계 공간 내에 제약하는 추가적인 KL 페널티를 적용하여 사전 훈련된 확산 모델을 피지컬 튜닝한다.
• 훈련 과정을 두 단계로 분해한다: (1) 보수적 정규화를 포함한 오프라인 보상 모델 훈련, (2) 보수적인 보상 모델을 사용한 확산 모델 피지컬 튜닝.
• 학습된 보상과 시각-언어 모델(예: LLaVA)의 일관성 신호를 조합하여 복합 보상 함수를 정의함으로써 이미지 생성에서 프롬프트 준거를 보장한다.
• Enformer와 Stable Diffusion을 사전 훈련된 모델로 사용하여 생물학적 시퀀스 설계(DNA/UTR/엔ハン서) 및 이미지 생성 작업에 이 방법을 적용한다.

실험 결과

연구 질문

• RQ1 보상 모델이 제한된 정적 데이터셋에서 훈련된 경우, 보수적인 피지컬 튜닝 접근이 오프라인 모델 기반 최적화에서의 과최적화 문제를 완화할 수 있는가?
• RQ2 보상 모델의 외삽 능력을 유효한 설계 생성을 훼손하지 않도록 어떻게 활용할 수 있는가?
• RQ3 보상 모델링에 불확실성 인식 페널티를 통합할 경우, 관측된 최고의 오프라인 설계를 초월한 일반화 및 성능 향상에 얼마나 기여하는가?
• RQ4 보수적인 보상 모델링과 KL 정규화의 조합이 시퀀스 및 이미지 생성 작업 전반에서 더 강건하고 충실한 생성을 이끌어내는가?
• RQ5 보상 성능 및 분포 충실도 측면에서 제안된 방법은 조건부 피지컬 튜닝 및 표준 강화 학습 피지컬 튜닝에 비해 어떻게 비교되는가?

주요 결과

• BRAID는 DNA 시퀀스 및 이미지 생성 작업 전반에서 관측된 데이터 분포를 초월하여 오프라인 데이터 중 최고의 설계를 능가하는 성능을 보이며, 관측된 데이터 분포를 초월한 일반화 능력을 입증한다.
• 이미지 생성 분야에서 BRAID가 생성한 이미지는 STRL이 프롬프트에서 벗어나는 것과 달리, LLaVA 및 LAION 예측기 기반으로 더 높은 미적 평가 점수를 확보하면서도 프롬프트 준거성과 시각적 품질을 유지한다.
• 학습 곡선 상 학습된 보상($\hat{r}$)의 변화를 보면 BRAID는 안정적이고 일관된 향상을 보이며, 반면 STRL는 빠른 보상 증가를 보이나 일반화 능력과 분포 이탈이 떨어진다.
• 실증 결과에 따르면 BRAID-boot와 BRAID-bonus는 보상 성능에서 STRL를 크게 능가함을 확인하여 보수적 정규화의 유용성을 입증한다.
• 이론적 분석은 소프트 엔트로피 정규화된 MDP 설정에서 이 방법을 프레임워크화하며, 보수적 피지컬 튜닝이 옹호된 보조 기반 모델보다 회귀를 감소시키고 뛰어난 성능을 달성할 수 있음을 보여준다.
• 절단 실험 결과 보수적인 보상 모델링과 KL 페널티가 모두 필수적임을 확인함—어느 한 쪽을 제거할 경우 성능 저하와 분포 외 생성 증가가 발생한다.