GFlowNet은 Generative Flow Network의 준말로, 각 state로 들어가는 flow의 합과 나가는 flow의 합을 동일하게 하여 우리가 목적으로 하는 reward에 비례하는 확률에 따라 다양한 샘플들을 생성하는 생성 모델을 의미한다. GFlowNet은 샘플의 diverse한 mode를 찾는 것을 통해 다양한 샘플을 생성하는 것을 가능하게 하고, 목적으로 하는 reward를 반영한 생성을 통해 실제 신약 개발 분야 등에서 유용하게 활용될 수 있다. 이번 글에서는 GFlowNet이 무엇이고, 이를 어떻게 학습 하는지에 대해 알아보겠다.
Directed acyclic graph (DAG)
우선, GFlowNet을 쉽게 이해하기 위해서는 GFLowNet의 DAG가 어떻게 구성되어 있는지부터 이해해야
한다. 다음 그림과 함께 설명을 시작하겠다.
위 그림에서 각 동그라미는 state, 이들을 연결하는 화살표는 action 혹은 state transition이라고 부른다. State는 어떤 compositional object, 즉 어떤 작은 component들이 합쳐진 object를 의미하고, 위 그림에서는 각 분자 구조 혹은 분자 구조의 subgraph가 된다. 또한 이러한 state는 파란색으로 표현된 처음의 initial state, 흰색으로 표현된 intermediate state, 그리고 노란색으로 표현된 최종 생성 샘플에 해당하는 terminal state로 이루어진다. Action은 이 state들 간의 연결을 의미하고, 보통의 경우 component를 더하는 것을 의미한다. 위 그림에서는 어떤 subgraph 혹은 fragment를 더하는 과정이 될 것이다.
Flow in GFlowNet
아래 그림에서 각 action에 해당하는 화살표의 굵기가 다른 것을 확인할 수 있는데, 이 굵기는 각 action에 따라 이동하는 flow의 양을 의미한다. 어떤 state로부터 해당 action을 선택할 확률은 해당 action의 flow의 양에 비례한다. 예를 들어, initial state로부터 나오는 두 개의 flow는 각각 굵은 순으로 3의 flow, 1의 flow를 의미하고, 이들을 뽑을 확률은 각각 0.75, 0.25가 되는 것이다. 이러한 flow는 흔히 물이 흐르는 것에 비유하고는 한다. 물이 어떤 state로 흘러 들어가면, 해당 state에서 같은 양의 물이 나오고, terminal state가 가지는 최종적인 물의 양이 해당 데이터를 샘플링할 확률에 비례하는 것이다. 아래부터는 약간의 수식을 사용할 예정인데, flow의 양을 $F$, score를 $R$, transition을 $T$로 정의하겠다.
이러한 flow는 다음과 같은 두 개의 중요한 특성을 가진다. 첫 번째로, 각 terminal state로 들어가는 inflow의 양은 각 terminal state 의 object가 가지는 score(혹은 reward) $R(s)$에 비례한다. 여기에서 말하는 score는 score function을 말할 때의 score가 아니라, 진짜 점수, 즉 우리가 최대화하고 싶어하는 특성을 의미한다. 예를 들어, 우리가 최적화하고 싶어하는 score가 분자의 합성 가능성이라고 하자. 그럼 위 그림에서 각 terminal state의 분자가 가지는 합성 가능성 score는 위에서부터 3, 2, 1이 되는 것이다. 이는 식으로 나타내면 다음과 같다. ($s_f$는 terminal state에서 들어가는 최종적인 하나의 state를 의미한다.)
두 번째로, 각 intermediate state에 들어가는 inflow의 양과 나오는 outflow의 양은 항상 같다. 예를 들어, initial state 다음 state 중 위에 있는 state를 보자. 이 state의 inflow는 4인데, outflow도 3+1=4이다. 이러한 특성이 각 state에 대해 항상 만족되어야 하는 것이다. 이는 flow matching condition이라고 부르고, 식으로 나타내면 다음과 같다. 식을 좀 더 구체적으로 설명하면, 어떤 state s가 가지는 flow의 양이 이 state로 들어오는 inflow의 합 혹은 이 state로부터 나가는 outflow의 합과 같기 때문에, 결론적으로 inflow의 양과 outflow의 양이 같아지는 것이다.
이 두 개의 특성을 동시에 만족시키는 것을 식으로 표현하면 다음과 같이 표현할 수 있다. 이 식을 풀어 설명하면, $s'$이 terminal state일 때는 $F(s'\rightarrow s)$가 0이 되기 때문에 $s'$로 들어가는 inflow가 score와 같다는 첫 번째 조건을 표현하고, $s'$이 terminal state가 아닐 때는 $R(s')$가 0이 되기 때문에 inflow와 outflow가 같다는 두 번째 조건을 표현하는 것이다.
이 두 개의 특성을 만족시키는 것이 중요한 이유는, 이들을 만족시킴으로써 GFlowNet이 score에 비례하는 샘플을 생성할 수 있기 때문이다. 즉, 이 flow의 양을 따라서 initial state를 시작으로 state transition을 진행하다보면 각 terminal state에 도달할 확률이 flow에 비례하게 되고, 이는 곧 score에 비례하게 된다. 결론적으로, flow를 알면, score에 비례하는 샘플을 생성할 수 있는 generation policy를 정의할 수 있는 것이다.
Training of GFlowNet
그럼 이제 우리는 우리의 목적인 score에 비례하는 샘플을 생성하기 위해 flow의 양을 정확하게 추정하는 방향으로 학습을 진행해야 한다. 가장 직관적인 objective는 다음과 같이 우리가 parameterize한 $F_\theta$가 다음 식을 만족하도록 학습되는 것이다. (이 식은 위에서 본 flow condition의 우변을 그냥 왼쪽으로 옮긴 것이다.) 이러한 objective는 GFlowNet의 학습의 가장 기본이 되는 형태로, flow matching objective라고도 부른다.
이외에도 detailed balance, trajectory balance 등 다양한 objective들이 존재하는데 이는 차차 다루어보도록 하겠다.
Experiments
GFlowNet은 다음과 같이 synthetic한 hyper-grid에서 reward가 높은 state를 찾는 task나 practical한 molecule generation과 같은 task에서 모두 좋은 성능을 보였다. 특히, 강화학습 방법들과 비교했을 때 어떤 reward를 만족하는 diverse한 solution들을 생성하는 데에 좋은 효과를 보였다.
하지만 GFlowNet도 아직까지 continuous domain에 대해서는 높은 성능을 보이지 못한다거나, 너무 search space나 trajectory 길이가 긴 task는 잘 다루지 못한다는 한계점을 가진다. 이번 글은 GFlowNet의 가장 기초적인 부분들에 대해 다루었고, 미래의 내가 추가적인 objective들과 같은 더 구체적인 내용을 다룰 것이다.
최근댓글