[논문 리뷰] Learning GFlowNets from partial episodes for improved convergence and stability

Kanika Madan, Jarrid Rector-Brooks외 7명|arXiv (Cornell University)|2022. 09. 26.

Reinforcement Learning in Robotics인용 수 5

한 줄 요약

이 논문은 다양한 길이의 부분 동작 부분열을 통해 학습할 수 있도록 하여 경향성 편향과 분산을 균형 잡는 새로운 훈련 목표인 Subtrajectory Balance (SubTB)를 제안한다. SubTB는 수렴 속도를 향상시키고 희박하거나 장수열 환경에서 안정성을 높이며, 펩타이드 및 단백질 생성과 같은 합성 및 실세계 작업에서 기존 목표인 TB와 FM보다 보상 및 다양성 지표에서 뛰어난 성능을 발휘한다.

ABSTRACT

Generative flow networks (GFlowNets) are a family of algorithms for training a sequential sampler of discrete objects under an unnormalized target density and have been successfully used for various probabilistic modeling tasks. Existing training objectives for GFlowNets are either local to states or transitions, or propagate a reward signal over an entire sampling trajectory. We argue that these alternatives represent opposite ends of a gradient bias-variance tradeoff and propose a way to exploit this tradeoff to mitigate its harmful effects. Inspired by the TD($λ$) algorithm in reinforcement learning, we introduce subtrajectory balance or SubTB($λ$), a GFlowNet training objective that can learn from partial action subsequences of varying lengths. We show that SubTB($λ$) accelerates sampler convergence in previously studied and new environments and enables training GFlowNets in environments with longer action sequences and sparser reward landscapes than what was possible before. We also perform a comparative analysis of stochastic gradient dynamics, shedding light on the bias-variance tradeoff in GFlowNet training and the advantages of subtrajectory balance.

연구 동기 및 목표

• 지역적 목표(예: 세부 균형)는 분산은 낮지만 편향이 높고, 궤적 수준의 목표(예: 궤적 균형)는 편향은 낮지만 분산이 높은 GFlowNet 훈련에서의 편향-분산 트레이드오���을 해결하기 위함.
• 장기 행동 시퀀스 또는 희박 보상 환경에서 효과적인 훈련을 가능하게 하기 위해 불완전하거나 부분적인 에피소드에서 학습할 수 있는 유연한 훈련 목표를 개발하기 위함.
• 특히 어려운 확률 모델링 설정에서 하이퍼파ram터 선택에 대한 민감도를 낮추고 수렴 속도를 높이며 강건성을 향상시키기 위함.
• 학습 가능한 $\lambda$ 파rameter를 통해 국소 일관성과 전체 궤적에 대한 책임 할당 사이를 통합하는 프레임워크 제공

제안 방법

• 강화학습의 TD($\lambda$)에 영감을 얻어, 다양한 길이의 부분궤적에 걸쳐 책임 할당을 균형 잡는 SubTB($\lambda$) 훈련 목표를 제안.
• 서브궤적에 걸친 시간 차이 목표의 가중합을 사용하며, 가중치는 유인 추적 파rameter $\lambda$에 의해 결정되어 국소(DB 유사)와 전역(TB 유사) 학습 사이의 보간 가능.
• 궤적 균형의 일반화로써 목표를 유도하여, 향후 보상에 기반해 중간 상태에 책임을 할당함으로써 불완전한 에피소드에서의 학습을 가능하게 함.
• 서브궤적 세그먼트를 사용해 기울기를 계산하는 스토하스틱 기울기 업데이트 규칙을 적용하여 전체 궤적 역전파보다 분산을 감소.
• 동적으로 서브궤적을 샘플링하고 가중치를 조정할 수 있는 미분 가능한 메커니즘을 도입하여 데이터 효율성과 수렴 속도 향상.
• 정책 파라미터가 신경망으로 파aram터화된 정책 기반 강화학습 방법을 통해 GFlowNets를 훈련시키며, 미지 상태로의 일반화 가능.

실험 결과

연구 질문

• RQ1 부분 에피소드에서 학습하는 GFlowNet 훈련 목표가 전체 궤적 또는 국소 목표에 비해 수렴 속도와 안정성 측면에서 향상되는가?
• RQ2 SubTB($\lambda$)의 하이퍼파ram터 $\lambda$는 GFlowNet 훈련에서 편향-분산 트레이드오프에 어떤 영향을 미치는가?
• RQ3 SubTB($\lambda$)는 이전 방법이 실패하는 장수열 또는 희박 보상 함수를 가진 환경에서 효과적으로 훈련을 가능하게 하는가?
• RQ4 SubTB($\lambda$)는 궤적 균형에 비해 기울기 분산을 얼마나 줄이며, 이는 훈련 동역학에 어떤 영향을 미치는가?
• RQ5 SubTB($\lambda$)는 순차적 생성 작업에서 샘플 품질(보상)과 다양성 측면에서 모두 향상시키는가?

주요 결과

• SubTB($\lambda$)는 합성 및 실세계 환경에서 TB, FM, DB 목표에 비해 더 빠른 수렴 속도와 하이퍼파ram터 선택에 대한 낮은 민감도를 보였다.
• 항미생물 펩타이드 생성 작업에서 SubTB($\lambda$)는 평균 상위 100개 보상 0.96 ± 0.02를 기록하여 TB(0.90 ± 0.03)와 FM/DB(0.78 ± 0.05)를 크게 앞섰다.
• 237개 길이의 형광 단백질 생성 작업에서 SubTB($\lambda$)는 상위 100개 보상 1.18 ± 0.10과 다양성 204.44 ± 0.45를 기록하여 TB(0.76 ± 0.19 및 204.31 ± 0.44)와 다른 기준선을 초월했다.
• SubTB($\lambda$)로 훈련된 모델은 더 빠른 초기 단계에서 더 많은 모드를 발견하고 목표 분포와 더 높은 스피어만 상관계수를 보였으며, 이는 더 빠르고 효과적인 탐색을 의미한다.
• SubTB($\lambda$) 하에서 스토하스틱 기울기의 분산은 TB에 비해 유의미하게 낮아, SubTB가 낮은 편향을 유지하면서도 기울기 분산을 감소시킨다는 가설을 검증했다.
• SubTB($\lambda$)는 이전 GFlowNet 목표가 실패하거나 수렴하지 못하던 희박 보상과 장수열 환경에서도 효과적인 훈련을 가능하게 한다.