[논문 리뷰] Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation

Emmanuel Bengio, Moksh Jain외 3명|arXiv (Cornell University)|2021. 06. 08.

Machine Learning in Materials Science인용 수 49

한 줄 요약

GFlowNet은 양의 보상에 비례하는 확률로 후보를 샘플링하기 위한 흐름 기반 생성 정책을 학습하여, 비반복적이고 다양한 고보상 출력(예: 분자)을 가능하게 한다.

ABSTRACT

This paper is about the problem of learning a stochastic policy for generating an object (like a molecular graph) from a sequence of actions, such that the probability of generating an object is proportional to a given positive reward for that object. Whereas standard return maximization tends to converge to a single return-maximizing sequence, there are cases where we would like to sample a diverse set of high-return solutions. These arise, for example, in black-box function optimization when few rounds are possible, each with large batches of queries, where the batches should be diverse, e.g., in the design of new molecules. One can also see this as a problem of approximately converting an energy function to a generative distribution. While MCMC methods can achieve that, they are expensive and generally only perform local exploration. Instead, training a generative policy amortizes the cost of search during training and yields to fast generation. Using insights from Temporal Difference learning, we propose GFlowNet, based on a view of the generative process as a flow network, making it possible to handle the tricky case where different trajectories can yield the same final state, e.g., there are many ways to sequentially add atoms to generate some molecular graph. We cast the set of trajectories as a flow and convert the flow consistency equations into a learning objective, akin to the casting of the Bellman equations into Temporal Difference methods. We prove that any global minimum of the proposed objectives yields a policy which samples from the desired distribution, and demonstrate the improved performance and diversity of GFlowNet on a simple domain where there are many modes to the reward function, and on a molecule synthesis task.

연구 동기 및 목표

• 높은 보상을 가진 후보들의 분포를 하나의 최상 후보가 아닌, 샘플링하도록 동기를 부여한다.
• 생성 과정을 흐름-네트워크 관점으로 도입하여 같은 최종 객체로 이어지는 여러 액션 시퀀스를 처리한다.
• 흐름 일관성을 강제하고 보상에 비례하는 정책을 얻는 학습 목표를 제공한다.
• 합성 도메인과 분자 생성 작업에서 확장성과 다양성 향상을 입증한다.
• 훈련 목표의 오프폴리시/오프라인 수렴 특성을 보여준다.

제안 방법

• 생성 과정을 단일 원천과 후보 객체에 대응하는 종단 싱크를 갖는 흐름 네트워크로 모델링한다.
• 액션 확률을 pi(a|s) = F(s,a)/F(s)로 정의하고, 여기서 F는 간선의 흐름을 나타내며 F는 흐름-일관성 방정식을 만족한다.
• 로그 흐름에 대한 흐름 매칭 목표를 사용하여 학습을 안정화한다: 들어오는 흐름의 로그와 나가는 흐름의 로그의 차이를 제곱해서 최소화한다 (Eq. 12).
• 흐름 목표의 전역 최소가 터미널 상태에서 pi(x)가 R(x)에 비례하는 정책을 산출한다는 것을 보인다.
• 이 방법은 오프폴리시임을 보여준다: 훈련은 넓은 지원의 정책으로부터의 트래젝토리를 사용하면서도 원하는 흐름과 분포로 수렴한다(명제 3).
• 합성 DAG/MDP 도메인과 도킹 기반 보상 프록시를 사용하는 대규모 분자 생성을 적용한다.

실험 결과

연구 질문

• RQ1 여러 액션 시퀀스가 동일한 최종 상태로 이어지는 경우에도 흐름-네트워크 기반 정책이 주어진 보상 R(x)에 비례하는 확률로 샘플을 생성할 수 있는가?
• RQ2 흐름 매칭 목표가 모든 고보상 모드를 회복하고 다양성 및 샘플 효율성 측면에서 전통적 MCMC/RL 기반과 비교하여 우수한가?
• RQ3 분자 생성과 같은 큰 조합 공간에 대해 방법이 확장 가능하고 오프라인/오프폴리시 훈련에 강건한가?
• RQ4 합성 및 분자 도메인에서 실제로 생성된 분포가 정규화된 보상 분포와 어떻게 비교되는가?

주요 결과

• GFlowNet은 C가 비단일적(동일한 경로로 가는 경로가 많아도)인 경우에도 R(x)에 비례하는 확률로 터미널 상태를 샘플링한다.
• 로그 흐름 매칭 목표는 원하는 분포에 해당하는 전역 최솟값을 산출하고 오프폴리시/오프라인 학습을 가능하게 한다.
• 하이퍼그리드 도메인에서 GFlowNet은 MCMC보다 더 적은 샘플로 pi(x) ∝ R(x)로 수렴하며 모드 분리에 강인하다.
• 대규모 분자 생성에서 GFlowNet은 MARS, PPO, JT-VAE+BO 기반보다 더 높은 보상과 더 다양한 분자를 더 빨리 생성한다.
• GFlowNet은 많은 모드를 찾고 커버하며(예: R>8인 모드가 1500개 이상), 상위 샘플 간 평균 쌍간 유사도를 기준선보다 낮게 달성하여 다양성이 더 높음을 시사한다.
• 이 방법은 상당한 액션 공간을 가진 분자 그래프에 확장 가능하며 같은 계산 예산 내에 더 많은 고점수 분자를 독특하게 생성한다.