Why This Paper?
1편에 이어서 GFlowNet Foundations에 대한 내용을 정리하도록 한다.
GFlowNets: Learning a Flow
지금까지 나타낸 정의들을 이용하여 GFlowNet 훈련 방법을 확장해보고자 한다. 몇 가지 새로운 훈련 기법을 탐구한다.
1. 시간 스탬프 도입을 통한 순환 구조 허용
일반적으로 GFlowNet의 상태 공간은 DAG로 구성되지만 순환이 포함된 상태 공간도 고려할 수 있다. 이를 위해 확장된 상태 공간 $S' = S x \mathbb{N}$을 도입한다. 시간 스탬프 $t$를 포함하여 여러 번 방문해도 각 방문이 서로 구분되게 만든다.
2. 확률적인 보상
기존에는 보상이 결정론적이라고 가정했지만 확률적인 보상도 고려할 수 있다. 즉, 상태 $s$에서 보상 $R(s)$이 확률적으로 결정될 수도 있다.이 경우, 손실 함수를 최소화 할 때 실제로 최적화되는 것은 기대값 기반의 손실 함수가 된다.(확률적이기 때문에)
3. 오프라인 학습
이 개념은 강화학습에서 offline 학습과 비슷하다. 원래 학습된 자체 샘플링 분포 $P = \Pi(o)$가 아닌 다른 분포에서 샘플을 가져와서 학습하는 것이 가능하다.
4. 기존 데이터를 활용한 훈련
기존의 데이터셋에서 추가적인 탐색을 이용하여 훈련을 진행할 수 있다.
Conditional Flows and Free energies
Flow 네트워크의 중요한 속성 중 하나는 초기 상태의 흐름 $F(s_0)$를 이용하여 정규화 상수 $Z$를 복원할 수 있다는 점이다.(1편 정리 10 참고). 이 $Z$는 종료 보상 함수 $R$에 해당하는 terminating flow(종료 흐름)과 관련된 분할 함수를 제공한다. 그렇다면 내부 상태 $s (s_0 < s < s_f)$에 대해서는 어떻게 할까? 만약 s에서 도달할 수 있는 종료 흐름들에 대한 정규화 상수를 구할 수 있다면, 이는 상태 $s$ 를 기준으로 종료 상태의 조건부 확률을 얻는 것과 같다.
그런데 $F(s)$가 단순히 그 상태에서 도달 가능한 종료 흐름들의 합을 나타내는가?
정답은 그럴 수도 있지만 대부분 "아니오"이다.
아래의 예시를 보자.
$F(s_2)$의 값은 4인데 실제로 $s_2$에서 도달할 수 있는 종료 흐름들의 합은 6이다. $s_2$에서 도달 가능한 종료 상태는 $s_5, s_6, s_7$이지만 $(s_0 \to s_1 \to s_5)$를 통해 $s_5$로 가는 흐름이 존재한다. 그런데 이 흐름은 $s_2$를 거치지 않기 때문에 $F(s_2)$에 포합되지 않는다. 따라서 위의 질문에 대한 대답은 아니오다.
정의 27. pointed DAG $G = (S, \mathbb{A})$가 주어졌을 때, 부분 순서 관계 $\geq$를 정의하고 에너지 함수 $\mathcal{E}: S \to \mathbb{R}$을 도입하여 우리는 자유 에너지(Free Energy) $\mathcal{F}(s)$를 아래와 같이 정의한다.
$$e^{-\mathcal{F}(s)} := \sum_{s':s' \geq s} e^{-\mathcal{E}(s')}$$
즉, 상태 $s$의 자유 에너지는 해당 상태보다 크거나 같은 모든 상태들의 에너지 함수 값의 지수형 합으로 정의된다.
Conditional flow networks
앞의 정리에서는 Flow Network를 DAG와 전체 궤적 집합 $\mathcal{T}$ 위에서 정의된 함수로 구성했다. 지금부터는 이 개념을 확장하여 각 구성요소를 어떤 정보 $x$에 대해 조건부로 정의할 수 있다. 일반적으로 이 조건 변수 $x$는 Flow Network에 영향을 주는 외부 정보이거나 내부 정보일 수 있다.
정의 28. $ \mathcal{X}$를 조건 변수들의 집합이라고 하자. 우리는 $ \mathcal{X}$에 속하는 각 $x$에 대해 DAG $G_x = (S_x, A_x)$를 정의한다. 이 DAG들은 초기상태 $(s_0|x)$와 종료 상태 $(s_f|x)$를 모두 포함한다. 또한 각 $G_x$에 대해 전체 궤적들의 집합을 $\mathcal{T}_x$라고 하고 모든 $\mathcal{T}_x$ 를 합친 전체 궤적 집합을 $\mathcal{T}$로 나타낸다.
$$\mathcal{T} = \bigcup_{x \in \mathcal{X}}\mathcal{T}_x$$
조건부 흐름 네트워크(Conditional flow network)는 아래의 3가지 요소로 구성된다.
1. $\mathcal{X}$(조건 변수 집합)
2. $\{G_x, x \in \mathcal{X}\}$(각 조건 변수 $x$에 해당하는 DAG들의 집합)
3. 조건부 흐름 함수 $F$
특정 $x$에 속하지 않는 궤적 $\tau$에 대해 $F(x, \tau) = 0$이며 $F_x$는 특정 조건 $x$에 대한 궤적 $\tau$를 $F(x, \tau)$로 매핑하는 함수이다.
Reward-conditional flow networks
정의 29. 조건 변수 $\mathcal{X}$가 주어졌다고 가정하자. 그리고 DAG $G = (S, \mathbb{A})$와 flow function $F$가 있다고 하자. 여기서 종료 상태 집합 $S^f$에 대해, 조건 $x$에 따라 달라지는 보상 함수들의 집합 $\{R_x : S^f \to \mathbb{R}^+, x \in \mathcal{X}\}$를 정의할 수 있다.
이제 $\mathcal{R}$과 호환되는 reward-conditional flow network란 정의 28에서 정의한 조건부 흐름 네트워크의 특수한 형태로, 아래의 조건을 만족하는 네트워크를 의미한다.
$$\forall x \in \mathcal{X}, \;\; \forall s \in S^f, \;\; F_x(s \to s_f) = R_x(s) $$
State-conditional flow networks
정의 30. DAG $G = (S, \mathbb{A})$와 flow function $F$가 있다고 하자. 각각의 상태 $s \in S$에서 $G_s$를 $s' \geq s$인 모든 상태 $s'$를 포함하는 $G$의 서브 그래프라고 하자. 이때 state-conditional flow network $F_s$는 $\{G_s, s \in S\}$로 정의되며 $G_s$에 대해 아래와 같은 식을 만족한다.
$$F_s(s' \to s_f) = F(s' \to s_f)$$
이번에 정의한 state-conditional flow network에서는 DAG의 구조가 상태 $s$에 따라 달라진다. 특히 초기 상태 $s_0$가 $s$로 변경되지만, 최종 상태 $s_f$는 항상 유지된다.
정리 31. 아무런 DAG $G = (S, \mathbb{A})$와 flow $F$가 주어졌을 때, 우리는 state-conditional flow network를 정의 30처럼 정의할 수 있다.
정리 32. state-conditional flow network $(G, F)$가 정의 30처럼 주어졌다고 가정하자. 어떤 상태 $s$에 대해 state-conditional flow network의 초기 flow는 $s' \geq s$인 모든 $F(s' \to s_f)$를 합친 것과 같다. 이것을 식으로 정리하면 아래와 같다.
$$F_s(s_0|s) = F_s(s) = \sum_{s':s' \geq s} F(s' \to s_f) = \exp(-\mathcal{F}(s))$$
여기서 $\mathcal{F}(s)$는 free energy로 정의되며 에너지 함수 $\mathcal{E}(s') = -\log F(s' \to s_f)$와 관련된다.
결론 33. DAG $G = (S, \mathbb{A})$와 flow $F$가 주어졌다고 가정하고 정의 30에 따라 state-conditional flow network를 정의할 수 있다. 어떤 상태 s가 주어졌을 때, flow function $F_s$는 종료 직전 상태인($\{s'' \in S^f\}$) $s''$에 대해 $s'' \geq s$인 확률을 계산할 수 있으며 이것을 $P_T(\cdot\;|\;s)$으로 표기한다.
이 확률 분포하에, $G_s$에서의 궤적이 특정 종료 상태 $s'$에서 끝날 확률은 (즉, 마지막 edge가 $s' \to s_f$일 확률)은 아래와 같다.
$$P_T(s'|s) = 1_{s' \geq s}e^{-\mathcal{E}(s') + \mathcal{F}(s)}$$
여기서 $\mathcal{E}$는 에너지 함수, 즉 $s'$에 대한 $-\log F(s' \to s_f)$를 매핑하며, $\mathcal{F}$는 free energy 함수(해당 상태에서의 전체 marginalization)을 고려한 값이다.
Conditional GFlowNets
우리가 GFlowNet을 이용하여 flow network의 flow를 추정할 수 있었던 것처럼, 조건부 GFlowNet을 사용하여 특정한 목표 보상 함수를 가진 조건부 flow network를 추정할 수도 있다. 앞에서 설명된 일반적인 GFlowNet의 구조를 따르지만 정의 28번처럼 모든 값들이 조건 변수 $x \in \mathcal{X}$에 따라 달라진다는 점이 다르다.
Training Energy-Based Models with a GFlowNet
GFlowNet은 에너지 기반 모델(Energy-Based Model)의 샘플링을 수행할 수 있으며, MCMC(Markov Chain Monte Carlo)의 대안으로 활용될 수 있다. 에너지 기반 모델에서는 특정 상태 $s$의 확률을 $P_{\theta}(s) = \frac{e^{-\mathcal{E}_{\theta}(s)}}{Z}$로 정의하며, 여기서 $ \mathcal{E}_{\theta}(s) $는 에너지 함수, $Z$는 정규화 상수이다. 이러한 확률 분포에서 직접 샘플링하는 것은 어려우므로, GFlowNet을 활용하여 근사적인 샘플링을 수행한다.
GFlowNet을 학습하여 종료 상태 확률 $P_T(s)$를 학습하고 종료 상태 보상을 $R(s) = e^{-\mathcal{E}_{\theta}}(s)$로 설정하면 학습한 종료 상태 확률 $P_T(s)$가 에너지 기반 모델의 확률 $P_{\theta}(s)$를 근사하게 된다.
Estimating Entropies, Conditional Entropies and Mutual Information
정의 34. 주어진 보상 함수 $R(s)$에 대해, 엔트로피 보상 함수 $R'(s)$를 아래와 같이 정의한다
$$R'(s) = -R(s)\log R(s)$$
아래서부터는 엔트로피 추정을 위해 2가지의 GFlowNet을 학습하는 방법을 설명한다. 첫 번째는 기본적인 종료 상태 보상 함수 $R(s)$를 기반으로 학습하는 것이고 두 번째는 엔트로피 보상 함수 $R'(s)$를 기반으로 학습하는 것이다.
정리 35. 주어진 flow network $(G, F)$에 대해, 종료 flow가 보상함수 $R$과 일치한다고 가정하자.
$$\forall s \in S^f, \quad F(s \to s_f) = R(s)$$
이며, 모든 $s$에 대해 $R(s) < 1$을 만족한다.
또한, 동일한 pointed DAG를 공유하지만, 엔트로피 보상 함수 $R'$를 따르는 새로운 flow network $(G, F')$를 고려하자. $R'$는 아래와 같이 정의된다. (정의 34번 참고)
$$R'(s) = -R(s)\log R(s)$$
이 경우, 종료 상태 확률 분포 $P_T(S = s) = \frac{R(s)}{Z}$는 이렇게 정의된다.(정의 13번 참고)
이 종료 상태 확률 분포를 이용하여 종료 상태 랜덤 변수 $S$의 엔트로피 $H[S]$를 정의하면 아래처럼 된다.
$$H[S] := -\sum_s P_T(s)\log P_T(s)$$
이 엔트로피를 GFlowNet을 통해 계산하면 결론적으로 아래와 같이 표현된다.
$$H[S] = \frac{F'(s_0)}{F(s_0)} + \log F(s_0)$$
정리 36. 조건 변수 $\mathcal{X}$가 주어진 경우, GFlowNet의 조건부 흐름 네트워크(conditional flow network)를 정의할 수 있다. 이 네트워크는 특정한 목표 보상 $R_x$을 따르는 종료 흐름을 학습한다. 여기서 $R_x(s) < 1$이어야 한다.
이와 별개로, 두 번째 조건부 흐름 네트워크 $F'$를 정의할 수 있는데 이는 엔트로피 기반 보상 $R'_x$(정의 34번 참고)를 만족하는 종료 흐름을 학습한다. 이때, 조건부 엔트로피 $H[S|x]$는 아래와 같이 주어진다.
$$H[S|x] = \frac{F'(s_0|x)}{F(s_0|x)} + \log F(s_0|x)$$
그리고 랜덤 종료 상태 $S$가 조건 변수 $X$를 따를 때, 상호 정보량(Mutual Information)은 아래와 같이 주어진다.
$$MI(S; X) = H[S] - \mathbb{E}_X[H[S|X]]$$
$$ MI(S; X) = \frac{F'(s_0)}{F(s_0)} + \log F(s_0) - \mathbb{E}_X\left [ \frac{F'(s_0|X)}{F(s_0|X)} + \log F(s_0|X)\right ] $$
GFlowNets on Sets, Graphs, and to Marginalize Joint Distributions
Set GFlowNets
이제부터는 조금 더 큰 범위에서 진행한다. GFlowNet을 랜덤 집합 $S$를 생성하는 도구로 간주하며, 이를 이용해 확률, 조건부 확률, 주변 확률 등의 값을 추정할 수 있다. 이러한 집합 $S$의 원소들은 더 큰 우주(universe) 집합 $\mathcal{U}$에서 선택된다.
정의 37. 주어진 우주(universe) 집합 $\mathcal{U}$에 대해 DAG $G = (S, \mathbb{A})$를 고려하자. 여기서 상태 공간 $S$는 $\mathcal{U}$의 모든 부분집합과 추가적인 종료 상태 $S_f$를 포함하는 집합으로 정의된다.
초기 상태 $s_0$는 { }, 즉 공 집합이며 만약 $s$와 $s'$ 사이의 간선이 있다면(즉, $s \to s' \in \mathbb{A}$라면) $s'$은 $s$에서 하나의 원소 $a \in \mathcal{U}$를 추가한 집합이다. 각 간선은 현재 집합에 새로운 원소를 추가하는 것과 같다.
따라서 모든 부분 집합들은 종료 상태 $s_f$와 연결되고 이것은 모든 $s \in S$에서 $s_f$로 이동이 가능하다는 뜻이다.
GFlowNet을 이용하여 종료 상태 확률 $P_T(s)$를 정의할 수 있다.
$$P_T(s) = \frac{F(s \to s_f)}{F(s_0)} = e^{-\mathcal{E}(s) + \mathcal{F}(s_0)}$$
결론 33에서 제시된 방식과 비슷하게 주어진 상태 $s$가 상태 $s'$을 포함하는 경우의 조건부 확률은 아래와 같이 정의된다.
$$P_T(s'|s'\supseteq s) = e^{-\mathcal{E}(s') + \mathcal{F}(s)} = \frac{F(s' \to s_f)}{F(s|s)}$$
즉, 어떤 집합 $s$가 주어졌을 때, 상위 집합 $s'$이 될 확률을 의미한다.
하지만 실제 GFlowNet의 학습 과정에서는 $F(s)$가 정확히 $R(s)$와 일치하지 않을 수 있다. 그래서 확률을 추정하는 방식으로 아래의 2가지를 사용할 수 있다.
$$\hat{P}_T(s) = \frac{\hat{F}(s \to s_f)}{\hat{F}(s_0)}$$
$$\hat{P}_T(s) = \frac{R(s)}{F(s_0)}$$
정리 38. 정의 37의 표기법을 따르면 $\mathfrak{S}(s) = {s' \supseteq s}$는 집합 $s$의 모든 상위집합(supersets)을 의미한다. 주어진 set flow network에서 $\mathfrak{S}(s)$에 속하는 임의의 요소가 선택될 확률은 다음과 같이 계산된다.
$$P_T(\mathfrak{S}(s)) = \sum_{s' \supseteq s}P_T(s') = \frac{e^{-\mathcal{F}(s)}}{Z} = \frac{F(s|s)}{F(s_0)} $$
이것을 간단하게 증명해보면 우선 $P_T(s')$의 정의를 이용하면 $P_T(s') = \frac{F(s' \to s_f)}{F(s_0)}$이고 정리 32번을 이용하면 $ \sum_{s':s' \geq s} F(s' \to s_f) = F(s|s)$이다. 이것들을 위의 식에 대입하면 아래와 같이 나온다.
$$\sum_{s'\supseteq s}\frac{F(s' \to s_f)}{F(s_0)} = \frac{F(s|s)}{F(s_0)}$$
그리고 마지막으로 자유 에너지 $\mathcal{F}(s)$의 정의를 이용하면 $e^{-\mathcal{F}(s)} = F(s|s)$이므로 결론적으로 아래의 식이 나온다.
$$\frac{F(s|s)}{F(s_0)} = \frac{e^{-\mathcal{F}(s)}}{Z}$$
GFlownet on Graphs
그래프는 노드와 엣지를 포함하는 특정한 종류의 집합(set)이다. 그렇기 때문에 GFlowNet의 연산을 그래프에도 적용이 가능하며 Deleu et al. (2022)에서는 DAG(Directed Acyclic Graph) 위에서 GFlowNet을 사용하여 베이지안 네트워크(Bayesian Network)의 그래프 구조에 대한 사후 확률 분포(posterior distribution) 근사를 학습했다.
이 과정에서 flow-matching 조건에서 유도된 손실(loss)을 최소화하도록 학습을 진행하였으며, 결과적으로 얻어진 분포가 목표하는 사후 확률과 잘 일치함을 보였다.
Marginalizing over Missing Variables
GFlowNet은 확률 변수의 결합 분포(joint distribution) 학습, 조건부 확률 계산, 마진 확률 추정 등의 다양한 확률적 문제를 해결하는 데 사용할 수 있다.
특히, 확률 변수 $X = (X_1, X_2, ..., X_n)$에 대해, 각 변수 $X_i$가 특정한 값 $x_i$를 가질 때 $(i, x_i)$ 쌍을 원소로 갖는 GFlowNet을 학습할 수 있다.
GFlowNet을 활용하면 다음과 같은 작업이 가능하다.
- 마진 확률(marginal probability) 추정: 특정 부분 변수 집합이 주어졌을 때, 나머지 변수의 확률 계산
- 정규화 상수(partition function) 추정: 확률 분포의 정규화 상수 $Z$를 계산
- 확률이 높은 변수 조합 탐색: 가장 가능성이 높은 변수 설정 찾기
- 엔트로피(entropy) 계산: 확률 분포의 불확실성 측정
이를 통해 GFlowNet은 확률적 샘플링을 보다 효율적으로 수행할 수 있으며, 최적의 변수 조합을 찾거나 확률 분포를 정밀하게 모델링하는 데 활용될 수 있다.
Continuous of Hybrid Actions and States
지금까지의 GFlowNet의 전개에서는 상태나 행동에 대한 합을 이용하여 계산을 수행했다. 이는 상태나 행동이 이산적(discrete) 공간에 속한다고 가정했기 때문이다.
하지만 상태 또는 행동이 연속적(continuous)일 경우 기존의 합을 적분으로 대체할 수 있다. 이런 연속값의 도입은 GFlowNet 구조의 변화를 가져다 준다.
아래에서는 이러한 연속적 요소를 처리하는 방법에 대한 내용을 정리할 것이다.
Integrable Normalization Constants
먼저, 연속적인 상태를 처리할 수 있다고 가정하면 이산적인 부분과 연속적인 부분을 결합한 하이브리드 상태(hybrid state)도 처리할 수 있다. 이를 위해 상태 $s$를 아래와 같이 분해한다.
$$s = (s^i, s^x)$$
여기서 $s^i$는 이산적(discrete)인 상태고, $s^x$는 연속적(continuous)인 상태를 나타낸다. 이 상태들을 이용하여 전이 확률을 다음과 같이 분해할 수 있다.
$$P_F(s_{t+1}|s_t) = P(s_{t+1}^x|s_{t+1}^i, s_t)P(s_{t+1}^i|s_t)$$
상태 전이를 두 단계로 나누는 방식과 같다.
- 이산적인 선택(discrete)를 먼저 진행하여 다음 상태 $s_{t+1}^i$를 선택한다.
- 연속적인 선택(continuous)를 수행하여 주어진 $s_{t+1}^i$에 대한 $s_{t+1}^x$를 샘플링한다.
가우시안 분포 등을 사용하여 단순하게 샘플링 할 수도 있지만 보통 최적화가 어려우며 이를 대신하여 혼합 모델, 정규화 흐름, 확산 모델 등으로 확률 밀도를 근사할 수 있다.
추가적으로 균형 조건을 만족시키기 위해 flow-matching loss를 활용하거나 전이 확률 모델링이 가능하며 GFlowNet의 DAG 구조에서는 특정 규칙을 따르는 전이만 허용된다.
GFlowNets in GFlowNets
GFlowNet은 연속 상태(continuous state)도 다룰 수 있도록 하위 수준의 GFlowNet-에너지 함수(GFlowNet, energy function) 쌍을 활용하여 확장할 수 있다.
- 하위 수준 GFlowNet 활용: 연속 변수를 포함하는 엣지 플로우를 모델링하기 위해 작은 GFlowNet을 사용해 특정 전이를 관리하고, 샘플링과 조건부 확률을 추정할 수 있다.
- 이산 상태 확장: 그래프 GFlowNet에서는 특정 노드 삽입 등의 전이를 개별적으로 모델링하여 보다 정밀한 제어가 가능하다.
- 연속 상태 연구: Lahlou et al. (2023)은 GFlowNet을 연속 상태 공간으로 확장하고, 기존 손실 함수가 연속 영역에서도 효과적으로 활용될 수 있음을 검증했다.
이를 통해 GFlowNet은 이산적 상태뿐만 아니라 연속적, 하이브리드 상태에서도 활용될 가능성을 확장하고 있다.
느낀점
논문이 거의 50페이지 가까이 되기 때문에 최대한 이해하려고 하면서 정리했다. 아무래도 이론적인 부분을 다루는 내용이다 보니 계속 이해를 하기 위해 노력했고 생각보다 재밌게 읽은 것 같다.
https://dl.acm.org/doi/pdf/10.5555/3648699.3648909