[논문 리뷰] Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review

Masatoshi Uehara, Yulai Zhao외 2명|arXiv (Cornell University)|2024. 07. 18.

Traffic control and management인용 수 4

한 줄 요약

이 튜토리얼은 단백질 안정성 또는 분자 활성도와 같은 하류 보상 함수를 최적화하기 위해 사전 훈련된 확산 모델을 강화학습(RL)-기반 프레임워크로 통합적으로 미세조정하는 방법을 제시한다. 노이즈 제거 과정을 마르코프 결정 과정으로 공식화함으로써, PPO, 보상 가중 최대가능도 추정, 경로 일致 학습과 같은 RL 알고리즘을 사용하여 원하는 지표를 직접 최대화하면서도 샘플 품질과 일반화 능력을 유지할 수 있다.

ABSTRACT

This tutorial provides a comprehensive survey of methods for fine-tuning diffusion models to optimize downstream reward functions. While diffusion models are widely known to provide excellent generative modeling capability, practical applications in domains such as biology require generating samples that maximize some desired metric (e.g., translation efficiency in RNA, docking score in molecules, stability in protein). In these cases, the diffusion model can be optimized not only to generate realistic samples but also to explicitly maximize the measure of interest. Such methods are based on concepts from reinforcement learning (RL). We explain the application of various RL algorithms, including PPO, differentiable optimization, reward-weighted MLE, value-weighted sampling, and path consistency learning, tailored specifically for fine-tuning diffusion models. We aim to explore fundamental aspects such as the strengths and limitations of different RL-based fine-tuning algorithms across various scenarios, the benefits of RL-based fine-tuning compared to non-RL-based approaches, and the formal objectives of RL-based fine-tuning (target distributions). Additionally, we aim to examine their connections with related topics such as classifier guidance, Gflownets, flow-based diffusion models, path integral control theory, and sampling from unnormalized distributions such as MCMC. The code of this tutorial is available at https://github.com/masa-ue/RLfinetuning_Diffusion_Bioseq

연구 동기 및 목표

• 엔트로피 정규화된 MDP에서 강화학습 기반 확산 모델 미세조정을 순차적 결정 문제로 공식화하기.
• PPO, 보상 가중 최대가능도 추정, 값 가중 샘플링, 경로 일치 학습과 같은 다양한 RL 알고리즘을 비교·분석하여 확산 모델의 미세조정에 활용하기.
• 보상 피드백 가용성(온라인 대 비온라인, 알려진 대 알려지지 않은 보상) 기반으로 미세조정 시나리오를 분류하고, 알고리즘 선택에 대한 영향을 분석하기.
• 강화학습 기반 미세조정의 형식적 목표를 명확히 하며, 비온라인 학습에서의 목표 분포 및 분포 이탈 위험을 설명하기.
• 분류기 가이드라인, GflowNets, 플로우 기반 모델, 경로 적분 제어 이론과 같은 관련 개념과의 연결 고리를 설정하기.

제안 방법

• 각 노이즈 제거 단계를 누적 보상을 최대화하기 위한 결정으로 간주하는 마르코프 결정 과정(MDP)으로 확산 노이즈 제거 과정을 공식화하기.
• 프록시멀 정책 최적화(PPO), 미분 가능한 보상 역전파, 보상 가중 최대가능도 추정(MLE)과 같은 표준 RL 알고리즘을 정책 최적화에 적용하기.
• 값 가중 샘플링과 경로 일치 학습(PCL)을 사용하여 고차원 및 다모드 분포에서의 샘플 품질과 훈련 안정성을 향상시키기.
• 사전 훈련된 확산 모델을 기준 SDE(스토크래틱 미분 방정식)로 활용하여 사전 분포를 새로 학습할 필요 없이 더 빠른 추론을 가능하게 하기.
• 비온라인 RL에서의 분포 이탈 문제를 분석하여, 분포 외 샘플에 대한 과도한 최적화 위험을 규명하고 대응 전략을 제안하기.
• 관련 프레임워크와의 통합: 분류기 가이드라인(값 가중 샘플링의 특수 케이스로 간주), GflowNets(플로우 기반 모델링용), 비정규화된 분포 샘플링을 위한 MCMC

Figure 1 : Illustrative examples of RL-based fine-tuning, aimed at optimizing pre-trained diffusion models to maximize downstream reward functions.

실험 결과

연구 질문

• RQ1 확산 모델의 순차적 노이즈 제거 과정은 어떻게 자연스럽게 강화학습 문제로 프레임워크화할 수 있는가?
• RQ2 PPO, 보상 가중 MLE, PCL과 같은 다양한 RL 알고리즘의 성능 및 한계는 확산 모델의 하류 보상 최대화를 위한 미세조정에서 어떻게 다르게 나타나는가?
• RQ3 피드백 확보 시나리오(온라인, 비온라인, 알려진 대 알려지지 않은 보상)는 강화학습 기반 미세조정에서 알고리즘 설계와 성능에 어떤 영향을 미치는가?
• RQ4 강화학습 기반 미세조정의 형식적 목표는 무엇이며, 목표 분포 및 분포 이탈과의 관계는 어떻게 되는가?
• RQ5 강화학습 기반 미세조정 방법은 분류기 가이드라인, GflowNets, 경로 적분 제어 이론과 같은 기존 개념과 어떻게 연결되는가?

주요 결과

• 강화학습 기반 미세조정은 단백질 안정성, RNA 번역 효율성, 분자 QED 점수와 같은 하류 보상 함수를 직접 최적화할 수 있게 하여, 타겟 생성에서 비-RL 기반 베이스라인을 능가하는 성능을 보인다.
• 보상 가중 MLE와 PPO는 효과적인 베이스라인으로 작용하며, 특히 안정성과 정책 최적화 프레임워크 덕분에 고차원, 다모드 환경에서 뛰어난 성능을 보인다.
• 경로 일치 학습(PCL)과 값 가중 샘플링은 안정적인 훈련 동역학을 제공하고, 복잡한 보상 지도에서 모드 붕괴를 줄이는 데 기여한다.
• 사전 훈련된 확산 모델을 기준 SDE로 활용하면 훈련 시간을 크게 줄이고, MCMC 기반 방법과 달리 사전 분포 추정이 필요 없어진다.
• 비온라인 RL 환경에서는 피드백 커버리지가 제한되어 있어, 분포 외 샘플에 대한 과도한 최적화가 발생할 수 있으며, 이는 분포에 대한 강건성의 필요성을 강조한다.
• 이 프레임워크는 강화학습 기반 미세조정과 분류기 가이드라인 간의 형식적 연결 고리를 설정하였으며, 분류기 가이드라인이 특정 보상 형상화를 가진 값 가중 샘플링의 특수 케이스임을 보여준다.

Figure 4 : Offline scenario