기술 요약: 아마티제이션된 이산 샘플링을 위한 근사 정책 최적화 (Proximal Policy Optimization)

문제 정의

본 논문은 유한한 이산 공간 $\mathcal{X}$ 상의 확률 분포 $p(x) = R(x)/Z$로부터 샘플링하는 과제를 다룹니다. 여기서 $R(x)$는 양의 보상 함수이며, $Z$는 알려지지 않은 정규화 상수입니다. 이 문제는 통계 물리학, 베이지안 통계학, 계산 생물학 등의 분야에서 발생합니다. 저자들은 **GFlowNet(Generative Flow Networks)**에 초점을 맞추고 있으며, 이는 유향 비순환 그래프(DAG) 상에서 객체 $x$를 생성하는 과정을 순차적인 프로세스로 모델링합니다. 목표는 종단 상태(terminal states)의 분포가 $p(x)$와 일치하도록 DAG를 탐색하는 확률적 순방향 정책 $P_F$를 훈련하는 것입니다.

GFlowNet 훈련은 가치 기반 목적 함수(예: Trajectory Balance, Detailed Balance)에 의해 주도되어 왔으나, 저자들은 이러한 목적 함수들이 대규모 강화 학습(RL) 및 확산 모델(diffusion model)의 사후 훈련(post-training)에서 지배적인 역할을 하는 **근사 정책 최적화(PPO)**와 같은 정책 경사(policy gradient) 방법론이 이 맥락에서는 여전히 미개척 영역으로 남아 있다는 점에 주목합니다.

방법론

이론적 토대: Soft RL로서의 GFlowNet

본 연구는 선행 문헌 [42, 7]에서 확립된 GFlowNet과 엔트로피 정규화 RL(entropy-regularized RL) 사이의 이론적 연결 고리를 구축합니다. 할인 계수 $\gamma=1$ 및 엔트로피 계수 $\alpha=1$을 설정함으로써, 저자들은 GFlowNet 샘플링 문제가 RL 환경에서의 soft value function을 최대화하는 것과 동일함을 입증합니다.

• 보상 구조: 결정적으로, 중간 보상은 비종단 전이(non-terminal transitions)에 대해 $r(s_t, s_{t+1}) = \log P_B(s_t | s_{t+1})$로 정의되며, 여기서 $P_B$는 역방향 정책(backward policy)입니다. 종단 상태에서는 $r(s_T) = \log R(s_T)$입니다.
• Soft 정책 개선: 이 soft-RL 목적 함수에 대한 최적 정책은 타겟 샘플링 분포인 $R(x)/Z$와 일치합니다.

Ent-PPO: 새로운 정책 경사 알고리즘

본 논문은 GFlowNet의 soft-RL 공식에 적응된 PPO의 변형인 Ent-PPO(Entropic Proximal Policy Optimization)를 도출합니다. 이 유도 과정은 GFlowNet에 PPO를 단순 적용할 때 흔히 누락되는 두 가지 핵심 요소를 강조합니다:

1. 올바른 보상 공식화: 소프트-RL 최적점이 샘플링 타겟과 일치하려면 중간 보상이 반드시 위에서 정의된 $\log P_B$여야 합니다. 그렇지 않으면 목적 함수가 희소 보상(sparse rewards)을 가진 표준 RL로 축소되어, 전체 분포를 샘플링하는 대신 $\arg\max_x R(x)$를 최적화하게 됩니다(모드 붕괴 발생).
2. 해석적 KL 신뢰 영역(Analytic KL Trust Region): 휴리스틱한 엔트로피 보너스 $c_2 H(\pi)$를 추가하는 표준 PPO와 달리, Ent-PPO는 soft policy improvement 정리로부터 직접 유도된 해석적 KL 페널티를 사용합니다.
   • 대리 목적 함수(surrogate objective)는 cross-entropy 항(soft advantage 정의로부터 유래)과 엔트로피 항을 결af하여 도출됩니다.
   • 그 결과, 다음과 같은 단계별 손실 함수가 도출됩니다: $\ell_t(\theta) = \text{PPOClip}(\rho_t, \hat{A}_t) - \text{KL}(\pi_\theta(\cdot|s_t) \parallel \pi_{\text{old}}(\cdot|s_t))$.
   • KL 항은 별도의 튜닝 가능한 엔트로피 계수의 필요성을 대체하며, soft-RL 공식에 내재된 신뢰 영역 역할을 수행합니다.

분산 감소 (Variance Reduction)

저자들은 이 프레임워크 내에서 Vanilla Policy Gradient(VPG) 베이스라인을 위한 분산 감소 기법을 체계적으로 평가합니다:

• 추정기(Estimators): 가장 단순한 정책 경사인 reward-to-go, 학습된 가치 베이스라인을 포함한 reward-to-go, Generalized Advantage Estimation (GAE), 그리고 Subtrajectory Evaluation Balance (SubEB-GAE)를 비교합니다.
• 결과: GAE와 학습된 가치 베이스라인의 조합이 가장 효과적인 전략으로 식별되었으며, 이는 단순한 추정기들에 비해 유리한 편향-분산 트레이드오프를 제공하고 수렴 속도와 안정성을 크게 향상시킵니다.

주요 기여

1. GFlowNet을 위한 표준 정책 경사의 도출: 저자들은 엔트로피 정규화 RL 프레임워크 하에서 이산 샘플러를 위한 표준 정책 경사 목적 함수를 도출하였으며, 안정적인 훈련을 위한 분산 감소 기법(베이스라인 및 어드밴티지 추정)의 중요성을 실험적으로 검증하였습니다.
2. Ent-PPO 알고리즘: GFlowNet에 PPO를 성공적으로 적용한 최초의 사례인 Ent-PPO를 제안합니다. 이 유도 과정은 필요한 중간 보상($\log P_B$)과 해석적 KL 신뢰 영역을 명시적으로 식별합니다.
3. 실험적 검증: 본 논문은 합성 그리드부터 복잡한 분자 그래프 생성에 이르기까지 다양한 벤치마크를 통해 Ent-PPO가 기존의 off-policy GFlowNet 목적 함수(DB, TB, SubTB) 및 단순 PPO 변형들을 능가함을 입증합니다.

실험 결과

저자들은 다음과 같은 환경에서 방법론을 평가합니다:

• 합성/소규모: Hypergrid, TFBind8 (DNA 서열), String QM9.
• 대규모: sEH (파편 기반 분자 그래프) 및 QM9 (원자 단위 분자 그래프).

주요 결과:

• 베이스라인 대비 우월성: Ent-PPO는 소규모 문제에서 낮은 Total Variation (TV) 거리를 달eric하고, 대규모 분자 문제에서 더 나은 ELBO / EUBO를 달성하며 DB, TB, SubTB보다 일관되게 우수한 성능을 보입니다.
• 샘플 효율성: Ent-PPO는 배치당 여러 번의 업데이트 에포크($K > 1$)를 수행할 때 샘пли 효율성과 수렴 속도가 크게 향상되는 이점을 얻으며, 이는 대규모 RL에서 흔히 나타나지만 이전에는 GFlowNet에서 탐구되지 않았던 특성입니다.
• 교정의 필요성:
  • KL 페널티: KL 항을 제거하면(단순 PPO로 회귀하면), 특히 여러 업데이트 에포크를 사용할 때 훈련이 붕괴되거나 불안정해집니다. 이는 올바른 soft-RL 목적 함수를 유지하면서 행동 정책(behavior policy)으로부터 너무 멀어지는 것을 방지하기 위해 KL 항이 필수적임을 보여줍니다.
  • 클리핑(Clipping): KL 항을 유지하면서 PPO 클리핑 메커니즘을 제거하면 발산이 발생합니다. 이는 해석적 KL만으로는 중요도 비율(importance ratio)의 변동을 제한하기에 불충분함을 나타냅니다.
• 역방향 정책 학습: Ent-PPO가 off-policy 목적 함수와 같이 순방향-역방향 공동 학습을 네이티브하게 지원하지는 않지만, Trajectory Likelihood Maximization (TLM)을 사용하여 역방향 정책을 별도로 학습하면 수렴 속도가 크게 빨라진다는 것을 보여줍니다.

의의 및 주장

본 논문은 Ent-PPO가 현대적이고 확장 가능한 정책 경사 방법론과 이산 샘플링의 특정 요구 사항 사이의 간극을 메움으로써 GFlowNet 훈련의 중요한 진전을 이루었다고 주장합니다.

• 이론적 엄밀성: 본 연구는 제1원칙인 엔트로피 정규화 RL으로부터 필요한 보상 구조와 신뢰 영역을 엄격하게 도출함으로써, GFlowNet에 PPO를 적용하려 했던 이전의 시도들을 바로잡았습니다.
• 실질적 영향: 결과는 적절히 조정된 정책 경사 방법론이 기존의 가치 기반 GFlowNet 목적 함수보다 수렴 속도 및 데이터 효율성 측면에서 더 뛰어난 성능을 보일 수 있음을 시사합니다.
• 확장성: Ent-PPO가 여러 업데이트 에포크를 활용하고 복잡한 분자 생성 작업(QM9)으로 확장 가능하다는 점은, 전통적인 가치 기반 방법론이 크레딧 할당(credit assignment)이나 분배 함수(partition function) 추정에 어려움을 겪을 수 있는 더 크고 복잡한 문제에 GFlowNet을 적용할 수 있는 경로를 제시합니다.

저자들은 자신의 방법론이 현재 이산 영역에 집중하고 있지만, soft-RL 프레임워크가 향후 연속 샘플링 영역으로 정책 경사 방법론을 적응시키기 위한 유망한 방향을 제공한다고 결론짓습니다.