GFlowNet Foundation 笔记（三）

系列文章
GFlowNet Foundation 笔记（一）
GFlowNet Foundation 笔记（二）

条件流与自由能

Def 24. 已知自由能 $\mathcal{F}(s)$
$$e^{-\mathcal{F}(s)}=\sum _{s^{\prime }:s^{\prime }\ge s}R(s^{\prime })=\sum _{s^{\prime }:s^{\prime }\ge s}{e}^{-\mathcal{F}(s^{\prime })}$$

其中， $s^{\prime }$ 为终止状态。注意， $e^{-\mathcal{F}(s)}$ 并不等同于 $F(s)$ 。

条件化 GFlowNet

除了 F(s) 之外，还有另一个可用于估计的量。即对于 $s\le s^{\prime }$ ，如果所有通过终止边缘 $s^{\prime }\to s_f$ 的流量都转向 $s$，则通过 $s$ 的流量总和将表示自由能，如下图所示。

Def 25. 用 $s_0|x$ 和 $s_f|x$ 分别表示在条件变量 $x$ 下的源状态和汇状态。使用 $R(s|x)$ 表示在条件 $x$ 下的终止奖励函数。

Prop 17. 当 GFlowNet 训练完成时
$$F(s_0|x)=\sum _{s|x}R(s|x)=Z(x)$$

Proof.
$$\begin{array}{rl}\sum _{s|x}R(s|x)& =\sum _{s|x}F(s\to s_f|x)\\ & =\sum _{s|x}\sum _{s\in \tau }F(\tau )\\ & =\sum _{\tau |x}F(\tau )=F(s_0|x)\end{array}$$

估计自由能

Def 30. 状态条件 GFlowNet 是条件 GFlowNet 的一种特殊形式，即条件集合 $\mathcal{X}$ 为状态集合，即 $s^{\prime }|s,\forall s\in \mathcal{X},s<s^{\prime }$ 。如果 $s^{\prime }\to s^{\prime \prime }\in \mathbb{A}$ ，根据马尔可夫性质
$$P_F(s^{\prime \prime }|s^{\prime },s)=P(s^{\prime }\to s^{\prime \prime }|s^{\prime },s)=P(s^{\prime }\to s^{\prime \prime }|s^{\prime })=P_F(s^{\prime \prime }|s^{\prime })$$

状态条件 GFlowNet 的训练目标为
$$\begin{gathered} \mathcal{L}=E_{(s_0,s_1,...,s_n,s_f)\sim {\pi }_T}[\sum _{t=0}^n{E}_{0\le t^{\prime }\le t}[L(s_t,s_{t+1}|s_{t^{\prime }})]] \\ L(s^{\prime },s^{\prime \prime }|s)=(log(\frac{\delta +\hat{F}(s^{\prime }|s){\hat{P}}_F(s^{\prime \prime }|s^{\prime })}{\delta +\hat{F}(s^{\prime \prime }|s){\hat{P}}_B(s^{\prime }|s^{\prime \prime },s)}){)}^2 \end{gathered}$$

Def 31. 称 $F(s|s)$ 为条件状态自流量。 $F(s|s)$ 表示当仅允许通过 s 的轨迹时通过 s 的流，并通过所有 s0 产生所需的流 R(s0)。

Prop 19. 当 GFlowNet 训练完成时
$$\begin{array}{rl}{e}^{-\mathcal{F}(s)}& =F(s|s)\\ & =\sum _{s^{\prime }\ge s}R(s^{\prime })\\ & =\sum _{s^{\prime }\ge s}F(s^{\prime }\to s_f)\end{array}$$

Def 32. 定义条件终止概率分布为
$$\begin{array}{rl}{P}_T(s|A)& =\frac{P(s\to s_f,A)}{P(A)}\\ & =\frac{1_{s\to s_f\in A}P(s\to s_f)}{{\sum }_{s^{\prime }\to s_f\in A}P(s^{\prime }\to s_f)}\\ & =\frac{1_{s\in A}P(s\to s_f)}{{\sum }_{s^{\prime }\in A}P(s^{\prime }\to s_f)}\\ & =\frac{{\sum }_{\tau \in A,s\to s_f\in \tau }P(\tau )}{{\sum }_{\tau \in A}P(\tau )}\end{array}$$

其中， $A$ 为任意轨迹的集合。
$$P_T(s|\mathcal{T})=P_T(s)=P(s\to s_f)=\frac{R(s)}{F(s_0)}=e^{-\mathcal{E}(s)+\mathcal{F}(s_0)}$$

Prop 20. 已知 $s\le s^{\prime }$
$$\begin{array}{rl}{P}_T(s^{\prime }|s)& =\frac{F(s^{\prime }\to s_f)}{{\sum }_{s^{\prime \prime }\ge s}F(s^{\prime \prime }\to s_f)}\\ & =\frac{F(s^{\prime }\to s_f)}{F(s|s)}\\ & =\frac{R(s^{\prime })}{{\sum }_{s^{\prime \prime }\ge s}R(s^{\prime \prime })}\\ & =e^{-\mathcal{E}(s)+\mathcal{F}(s)}\end{array}$$

使用 GFlowNet 训练基于能量的模型

定义模型 $P_{\theta }(s)=e^{-{\mathcal{E}}_{\theta }(s)}/Z$ ，其中 $s$ 为终止状态。根据 ${\hat{P}}_T$ 抽取的 GFlowNet 样本可用于获得上述模型在观测数据 $x$ 下的负对数似然的随机梯度估计量
$$\begin{array}{rl}\frac{\delta -log{P}_{\theta }(x)}{\delta \theta }& =\frac{\delta {\mathcal{E}}_{\theta }(x)}{\delta \theta }+\frac{\delta logZ}{\delta \theta }\\ & =\frac{\delta {\mathcal{E}}_{\theta }(x)}{\delta \theta }+\frac{\delta log{\sum }_s{e}^{-{\mathcal{E}}_{\theta }(s)}}{\delta \theta }\\ & =\frac{\delta {\mathcal{E}}_{\theta }(x)}{\delta \theta }+\frac{1}{{\sum }_s{e}^{-{\mathcal{E}}_{\theta }(s)}}\sum _s{e}^{-{\mathcal{E}}_{\theta }(s)}(-\frac{\delta {\mathcal{E}}_{\theta }(s)}{\delta \theta })\\ & =\frac{\delta {\mathcal{E}}_{\theta }(x)}{\delta \theta }-\sum _s{P}_{\theta }(s)\frac{\delta {\mathcal{E}}_{\theta }(s)}{\delta \theta }\end{array}$$

其中， $s\sim {\hat{P}}_T(s)$ 。
引入隐变量 $h$ 后， $P_{\theta }(x,h)=e^{-{\mathcal{E}}_{\theta }(x,h)}/{\sum }_{x,h}{e}^{-{\mathcal{E}}_{\theta }(x,h)}$ ，边缘负对数似然的梯度变为
$$\begin{array}{rl}\frac{\delta -log{P}_{\theta }(x)}{\delta \theta }& =\frac{\delta -log{\sum }_h{P}_{\theta }(x,h)}{\delta \theta }\\ & =-\frac{1}{{\sum }_h{P}_{\theta }(x,h)}\sum _h\frac{\delta P_{\theta }(x,h)}{\delta \theta }\\ & =-\frac{1}{P_{\theta }(x)}\sum _h\frac{\delta }{\delta \theta }(\frac{e^{-{\mathcal{E}}_{\theta }(x,h)}}{{\sum }_{x,h}{e}^{-{\mathcal{E}}_{\theta }(x,h)}})\\ & =-\frac{1}{P_{\theta }(x)}\sum _h(-P_{\theta }(x,h)\frac{\delta {\mathcal{E}}_{\theta }(x,h)}{\delta \theta }+P_{\theta }(x,h)\sum _{s,h}{P}_{\theta }(s,h)\frac{\delta {\mathcal{E}}_{\theta }(s,h)}{\delta \theta })\\ & =\sum _h{P}_{\theta }(h|x)(\frac{\delta {\mathcal{E}}_{\theta }(x,h)}{\delta \theta }-\sum _{s,h}{P}_{\theta }(s,h)\frac{\delta {\mathcal{E}}_{\theta }(s,h)}{\delta \theta })\end{array}$$

使用 GFlowNet 进行主动学习

训练分为外循环更新，即学习真实的能量函数（奖励函数），和内循环更新，即使用学习到的能量函数作为驱动目标训练 GFlowNet。

估计熵、条件熵和互信息

Def 33. 定义 熵奖励函数( entropic reward function ) $R^{\prime }$
$$R^{\prime }(s)=-R(s)logR(s)$$

新训练一个 GFlowNet ，原来的目标为 $R$ ，新训练的目标为 $R^{\prime }$ 。

Prop 21. 终止状态随机变量 $S$ 的熵 $H[S]$ 为
$$\begin{array}{rl}H[S]& =-\sum _s{P}_T(s)log{P}_T(s)\\ & =-\sum _s\frac{R(s)}{F(s_0)}(logR(s)-logF(s_0))\\ & =\frac{-{\sum }_{s}R(s)logR(s)+logF(s_0){\sum }_{s}R(s)}{F(s_0)}\\ & =\frac{{\sum }_s{R}^{\prime }(s)}{F(s_0)}+logF(s_0)\\ & =\frac{F^{\prime }(s_0)}{F(s_0)}+logF(s_0)\end{array}$$

其中， $F^{\prime }$ 为新训练 GFlowNet 的流量度量。

Prop 22. 条件熵 $H[S|x]$ 为
$$H[S|x]=\frac{F^{\prime }(s_0|x)}{F(s_0|x)}+logF(s_0|x)$$

当 $x$ 是轨迹中的一个事件时，将考虑经过该事件的轨迹集合。当 $x=s$ 时，只考虑经过 $s$ 的轨迹
$$H[S|s]=\frac{F^{\prime }(s_0|s)}{F(s_0|s)}+logF(s_0|s)=\frac{F^{\prime }(s|s)}{F(s|s)}+logF(s|s)$$

Corollary 4. 终止状态随机变量 $S$ 和条件随机变量 $X$ 的互信息为
$$MI(S;X)=H[S]-E_X[H(S|X)]=\frac{F^{\prime }(s_0)}{F(s_0)}+logF(s_0)-E_X[\frac{F^{\prime }(s_0|X)}{F(s_0|X)}+logF(s_0|X)]$$