[논문 리뷰] Trajectory balance: Improved credit assignment in GFlowNets

Nikolay Malkin, Moksh Jain외 3명|arXiv (Cornell University)|2022. 01. 31.

Topic Modeling인용 수 12

한 줄 요약

이 논문은 유연한 행동 시퀀스에서의 신용 할당을 향상시키기 위해 전체 경로에 대해 최적화하는 것이지 국소적 유량 제약 조건이 아닌, GFlowNets를 위한 새로운 학습 목표인 경로 균형을 제안한다. 이는 전체 경로에 대해 최적화함으로써 장기 행동 시퀀스에서의 신용 할당을 향상시킨다. 이는 궁극적으로 궤적 균형의 전역 최소화가 목표 분포에서 정확한 샘플링을 보장함을 증명하며, 실증적으로는 더 빠른 수렴, 더 높은 샘플 다양성, 장기 시퀀스 및 큰 행동 공간 설정에서의 향상된 강건성과 함께 검증된다.

ABSTRACT

Generative flow networks (GFlowNets) are a method for learning a stochastic policy for generating compositional objects, such as graphs or strings, from a given unnormalized density by sequences of actions, where many possible action sequences may lead to the same object. We find previously proposed learning objectives for GFlowNets, flow matching and detailed balance, which are analogous to temporal difference learning, to be prone to inefficient credit propagation across long action sequences. We thus propose a new learning objective for GFlowNets, trajectory balance, as a more efficient alternative to previously used objectives. We prove that any global minimizer of the trajectory balance objective can define a policy that samples exactly from the target distribution. In experiments on four distinct domains, we empirically demonstrate the benefits of the trajectory balance objective for GFlowNet convergence, diversity of generated samples, and robustness to long action sequences and large action spaces.

연구 동기 및 목표

• 유연한 행동 시퀀스에서 느린 시간적 신용 할당 문제를 겪는 기존 GFlowNet 학습 목표인 유량 매칭과 세부 균형과 같은 목표의 비효율적인 신용 전파 문제를 해결하기 위해.
• 복합적이고 장기적인 경로에서 더 효율적이고 정확한 신용 할당을 가능하게 하는 새로운 학습 목표를 개발하기 위해.
• 궤적 균형이 전역적으로 최소화될 경우, 목표 비정규화된 밀도 R(x)에서 정확히 샘플링하는 정책을 도출함을 증명하기 위해.
• 네 가지 다른 도메인에서 장기 시퀀스 및 큰 행동 공간 설정에서 수렴 속도, 샘플 다양성, 강건성 측면에서 궤적 균형의 우수성을 실증적으로 검증하기 위해.
• GFlowNets에서 세부 균형 목표에 대한 첫 실증 평가를 제공하여 궤적 균형 및 유량 매칭과의 비교 분석을 가능하게 하기 위해.

제안 방법

• 초기 상태에서 종료 상태에 이르는 전체 행동 시퀀스(경로)에 기반하여 국소적 엣지 유량이나 쌍방향 상태 전이가 아닌, 전체 경로에 대해 작용하는 새로운 학습 목표인 궤적 균형을 제안한다.
• 궤적 균형 목표를 샘플된 경로들에 대한 합으로 정의하며, 각 경로의 기여는 경로를 따라 흐름이 목표 보상 R(x)와 얼마나 다를지를 기반으로 한다.
• 신경망을 사용해 엣지 유량 F(u→v)를 매개변수화하며, 정책는 F(u→v)/∑v′F(u→v′)로 유도하여 확률적 행동 선택을 보장한다.
• 경로 파악 도함수와 재구성 기반 기법을 사용해 미분 가능 모델에 대해 기울기 하강법을 통해 궤적 균형 목표를 최적화한다.
• 궤적 균형 목표의 전역 최소화자가 존재할 경우, 목표 분포 p(x) ∝ R(x)에서 정확히 샘플링하는 정책을 도출함을 증명한다.
• 여러 기준 환경에서 표준 GFlowNet 학습 파ip라인을 사용해 궤적 균형을 유량 매칭 및 세부 균형 목표와 비교·구현한다.

실험 결과

연구 질문

• RQ1 유량 기반 목표인 유량 매칭 및 세부 균형과 비교해, 궤적 기반 목표가 GFlowNets에서의 신용 할당 효율성을 향상시킬 수 있는가?
• RQ2 궤적 균형이 장기 시점 생성 작업에서 더 빠른 수렴과 더 높은 샘플 다양성을 이끌어낼 수 있는가?
• RQ3 궤적 균형은 복합적 물체 생성에서 큰 행동 공간과 고차원 상태 공간에 대해 강건한가?
• RQ4 표본 품질과 학습 안정성 측면에서 궤적 균형은 실증적으로 유량 매칭 및 세부 균형과 어떻게 비교되는가?
• RQ5 궤적 균형은 이전에 가능하지 않았던 더 긴 시퀀스에서 GFlowNets의 학습을 가능하게 할 수 있는가?

주요 결과

• 모든 평가 도메인에서, 특히 장기 시퀀스 및 큰 행동 공간 설정에서, 궤적 균형은 유량 매칭 및 세부 균형보다 더 빠른 학습 수렴을 달성한다.
• 엔트로피 및 커버리지와 같은 지표로 측정한 결과, 기준 목표 대비 샘플 다양성이 크게 향상됨을 확인했다.
• 이전 목표들이 수렴하지 못했던 100단계에 이르는 행동 시퀀스에서도 궤적 균형은 GFlowNets의 성공적인 학습을 가능하게 했다.
• 실증 결과로 궤적 균형은 생성 샘플과 목표 분포 R(x) 사이의 발산을 감소시켜, 원하는 보상 함수와의 보다 나은 일치를 나타냄을 보여주었다.
• 세부 균형 목표의 첫 실증 검증을 통해 이론적으로 타당한 바를 확인했지만, 궤적 균형에 비해 더 느린 수렴 속도와 더 높은 샘플 분산을 드러냈다.
• 궤적 균형 목표의 전역 최소화가 정확히 목표 분포 p(x) ∝ R(x)에서 샘플링하는 것을 보장함을 증명함으로써 이론적 정확성을 입증했다.