Bayesian Flow Networks 学习笔记

· · 科技·工程

BFN 学习笔记

Introduction

建模可看作传输信息

Alice 能接触到数据（data），她希望把数据传输给 Bob。在每一步中，Alice 都会向 Bob 发送一条消息（message），揭示有关数据的一些信息，而 Bob 希望以尽可能少的比特接收数据。在每一轮收到消息前，Bob 会对消息进行猜测，他的猜测越好，传输它所需的比特就越少。收到消息后，Bob 使用他刚刚获得的信息来改进他对下一条消息的猜测。损失（loss）函数是所有消息所需的总位数。

Autoregressive：

适用于离散信息，如语言，其 loss 是负对数自然，第二次生成与上一次有关。对应算术编码。

缺点：不适用于图像等连续且无前后依赖关系的信息，有多少变量就要有多少次更新

Diffusion：

适用于连续信息，如图像，loss 是 KL 散度。对应 bits-back coding 。

解耦生成步骤数与变量个数，但在离散信息上表现不佳

Bayesian Flow Networks

网络根据数据分布的参数运行，即使数据离散也保证生成过程连续、可导

B 有一个先验 input distribution：连续数据为标准正态分布， 离散数据为各类别概率均等的类别分布(categorical distribution)。每一轮 B 把 input distribution 的一些参数（如平均值）喂给神经网络，神经网络输出 output distribution 的参数。 A 通过对原始数据加噪来构建一个 "sender distribution"，从中采样并把采样数据发给 Bob。B 将 output distribution 与 A 加的噪声分布的卷积作为 receiver distribution。B 用采样数据，根据贝叶斯推断更新 input distribution。只要 input distribution 对数据中的所有变量独立建模，贝叶斯更新就是封闭的。loss 为 sender distribution 和 receiver distribution 的 KL 散度

input distribution 独立地接收每个变量的信息，因此无法利用上下文信息（例如图像中的相邻像素或文本中的相关词）。output distribution 由神经网络产生，该神经网络共同处理输入分布中的所有参数，使其能够访问所有可用的上下文。所以 BFN 结合了贝叶斯推理和深度学习的优势：前者提供了一种数学上最优且可精细控制的方式来收集和汇总有关单个变量的信息，而后者则擅长整合许多相互关联的变量的信息。

相对于其他模型，BFN 体现了一个分布到一个分布的函数，其优点是，即使数据离散，输入的参数也连续。

数学框架

input distribution

记 D 维数据 x 所服从的分布参数为 D 维向量 \theta，则输入分布定义为

p_I(x|\theta)=\prod_{d=1}^Dp_I(x^{(d)}|\theta^{(d)})

sender distribution

$$ p_S(y|x;\alpha)=\prod_{d=1}^Dp_S(y^{(d)}|x^{(d)},\alpha) $$ #### output distribution 设 BFN 神经网络为 $\Psi$ ， $t$ 为时间，$\theta$ 为 input distribution 的参数 $$ p_O(x|\theta,t)=\prod_{d=1}^Dp_O(x^{(d)}|\Psi^{(d)}(\theta,t)) $$ $\Psi^{(d)}$ 与 $\theta$ 所有维度有关，所以 output distribution 可利用上下文信息。 #### receiver distribution $$ p_R(y|\theta,t,\alpha)=\mathbb E_{p_O(x'|\theta,t)}[p_{S}(x'|\alpha)] $$ 可理解为，知道 sender distribution 的形式，不知道具体的 $x$。是 sender distribution 在 output distribution 下的期望 #### Bayesian update distribution 根据贝叶斯推断定义贝叶斯更新函数 $h \theta'\leftarrow h(\theta,y,\alpha)\\ p_U(\theta'|\theta,x,\alpha)=\mathbb E_{p_S(y|x,\alpha)}[\delta(\theta'-h(\theta,y,\alpha))] $$ \mathbb E_{p_U(\theta_1|\theta_0,x,\alpha_1)}\mathbb E_{p_U(\theta_2|\theta_1,x,\alpha_2)}\dots \mathbb E_{p_U(\theta_{n-1}|\theta_{n-2},x,\alpha_{n-1})}p_U(\theta_n|\theta_{n-1},x,\alpha_n)=p_U(\theta_n|\theta_{0},x,\alpha=\sum_{i=1}^n\alpha_i) $$ #### 加噪设置 $\beta(t) \beta(t)=\int_0^t\alpha(t')dt'\qquad \alpha(t')>0\\ \beta(0)=0\\ \frac{d\beta(t)}{dt}=\alpha(t)\\ \alpha=\sum_{i=1}^n\alpha_i=\sum_{i=1}^n\int_{t_{i-1}}^{t_i}\alpha(t')dt'=\beta(t)

Bayesian Flow Distribution

贝叶斯流分布是当前时刻的贝叶斯更新分布的边缘分布，它考虑了从起始至今每一步更新的先验 \theta_1,\theta_2\dots,\theta_n

给定先验参数 \theta_0和 p_U,\beta(t)

p_F(\theta|x,t)=\mathbb E_{p_U(\theta_1|\theta_0,x,\alpha_1)}\mathbb E_{p_U(\theta_2|\theta_1,x,\alpha_2)}\dots \mathbb E_{p_U(\theta_{n-1}|\theta_{n-2},x,\alpha_{n-1})}p_U(\theta|\theta_{n-1},x,\alpha_n)=p_U(\theta|\theta_0,x,\beta(t))

损失函数

假设 A一共向 B 传输了 n 条消息(噪声数据)：y_1,y_2,\dots,y_n，最后再将原始数据 x 传过去

n 步离散时间损失：传输 n 条消息的损失

发送者分布和接收者分布之间的 KL 散度

L^n(x)=\mathbb E_{p(\theta_1,\dots,\theta_{n-1})}\sum_{i=1}^nD_{KL}(p_S(\cdot|x;\alpha_i)||p_R(\cdot|\theta_{i-1},t_{i-1},\alpha_i))\\ p(\theta_1,\dots,\theta_{n-1})=\prod_{i=1}^{n-1}p_U(\theta_i|\theta_{i-1},x,\alpha_i)\\ L^n(x)=n\mathbb E_{i\sim U(1,n)}\mathbb E_{p_U(\theta_1|\theta_0,x,\alpha_1)}\dots \mathbb E_{p_U(\theta_{n-1}|\theta_{n-2},x,\alpha_{n-1})}D_{KL}(p_S(\cdot|x;\alpha_i)||p_R(\cdot|\theta_{i-1},t_{i-1},\alpha_i))\\ =n\mathbb E_{i\sim U(1,n),p_F(\theta|x,t_{i-1})}D_{KL}(p_S(\cdot|x;\alpha_i)||p_R(\cdot|\theta_{i-1},t_{i-1},\alpha_i))

可以通过蒙特卡洛采样从离散的整数均匀分布 U(1,n) 与贝叶斯流分布 p_F(\theta|x,t_{i-1}) 中采样计算得到。

重构损失
L^r=-\mathbb E_{p_F(\theta|x,t=1)}\ln p_O(x|\theta,t=1)
损失函数
L(x)=L^n(x)+L^r(x)\\

与 VAE 同理

连续时间的损失函数

令 n\to \infty

\epsilon:=\frac1n\\ \alpha(t,\epsilon):=\beta(t)-\beta(t-\epsilon)\\ L^{\infty}(x):=\lim_{n\to \infty}L^n(x)\\ L^{\infty}(x)=\lim_{\epsilon\to0}\frac1{\epsilon}\mathbb E_{t\sim U(\epsilon,1),p_F(\theta|x,t-\epsilon)}D_{KL}(p_S(\cdot|x;\alpha(t,\epsilon))||p_R(\cdot|\theta,t-\epsilon,\alpha(t,\epsilon))\\ =\mathbb E_{t\sim U(0,1),p_F(\theta|x,t)}\alpha(t)\frac{||g(x)-E[P(\theta,t)]||}{2C} ### 连续型数据 #### input distribution $$ \theta:=(\mu,\rho)\\ p_I(x|\theta)=\mathcal N(x|\mu,\frac1\rho I)\\ \theta_0:=(0,1) $$ $p_I$ 不直接用于预测，而是将该分布的参数作为 BFN 的输入 $p_I$ 的参数会根据观测样本而进行后验更新 #### 贝叶斯更新函数 若先验 $p_I(x|\theta)\sim\mathcal N(\mu_{i-1},\rho_{i-1}^{-1})$，似然 $p_S(y|x;\alpha )\sim\mathcal N(x,\alpha^{-1}I)$，先验通过样本 $y$ 和贝叶斯推断来更新后验，后验 $p(x|y)\sim \mathcal N(\mu_i,\rho_i^{-1})$。 $$ \theta_{i-1}=(\mu_{i-1},\rho_{i-1})\\ h(\theta,y,\alpha)=(\mu_{i},\rho_i)\\ \rho_i=\rho_{i-1}+\alpha\\ \mu_i=\frac{\mu_{i-1}\rho_{i-1}+y\alpha}{\rho_{i}}（看作y的线性变换）\\ \mu_i\sim\mathcal N(\frac{\mu_{i-1}\rho_{i-1}+\alpha y}{\rho_{i}},\frac{\alpha}{\rho_{i}^2}I)\\ $$ $\mu_i$ 是参数 $\theta_i$ 的唯一随机变量，因此 $$ p_U(\theta_i|\theta_{i-1},x,\alpha)=\mathcal N(\frac{\mu_{i-1}\rho_{i-1}+\alpha x}{\rho_{i}},\frac{\alpha}{\rho_{i}^2}I) $$ #### 精度可加性证明 $$ \mu_{i-1}\sim\mathcal N(\frac{\mu_{i-2}\rho_{i-2}+\alpha_a y}{\rho_{i-1}},\frac{\alpha_a}{\rho_{i-1}^2}I)\\\mu_i'=\frac{\mu_{i-1}\rho_{i-1}+\alpha_by}{\rho_{i}}\\\mu_{i}'\sim\mathcal N(\frac{\mu_{i-2}\rho_{i-2}+(\alpha_a+\alpha_b) y}{\rho_{i}},\frac{\alpha_a}{\rho_{i}^2}I)\\\mu_i\sim\mathcal N(\frac{\mu_{i-1}\rho_{i-1}+\alpha_b y}{\rho_{i}},\frac{\alpha_b}{\rho_{i}^2}I)\\\mu_i=\mu_i'+\epsilon,\epsilon\sim\mathcal N(0,\frac{\alpha_b}{\rho_{i}^2}I)\\\mu_i\sim\mathcal N(\frac{\mu_{i-2}\rho_{i-2}+(\alpha_a+\alpha_b) y}{\rho_{i}},\frac{\alpha_b+\alpha_a}{\rho_{i}^2}I)\\p_U(\theta_i|\theta_{i-2},x,\alpha_{a}+\alpha_b)=\mathcal N(\mu_i|\frac{\mu_{i-2}\rho_{i-2}+(\alpha_a+\alpha_b) x}{\rho_{i}},\frac{\alpha_b+\alpha_a}{\rho_{i}^2}I)\\\mathbb E_{p_U(\theta_{i-1}|\theta_{i-2},x,\alpha_{a})}p_U(\theta_i|\theta_{i-1},x,\alpha_b)=p_U(\theta_i|\theta_{i-2},x,\alpha_{a}+\alpha_b)\\ $$ #### 加噪设置：线性熵减 精度 $\rho_t=1+\beta(t)$，$\sigma_t^2=\frac{1}{1+\beta(t)}

由线性熵减推出 \beta(t)：

H(t)=\frac D2\ln(\frac{2\pi e}{1+\beta(t)})\\ H(t)=(1-t)H(0)+tH(1)\\ (1+\beta(t))^{-1}=(1+\beta(1))^{-t}\\ \beta(t)=\sigma_1^{-2t}-1\\ \alpha(t)=-2(\ln\sigma_1)(\sigma_1^{-2t})

Bayesian Flow Distribution

p_F(\theta|x,t)=p_U(\theta|\theta_0,x,\beta(t))=\mathcal N(\frac{\beta(t) x}{1+\beta(t)},\frac{\beta(t)}{(1+\beta(t))^2}I)\\ \gamma(t):=\frac{\beta(t)}{1+\beta(t)}=1-\sigma^{2t}_1\\ p_F(\theta|x,t)=N(\gamma(t)x,\gamma(t)(1-\gamma(t))I)

Output Distribution

根据采样的重参数化技巧，input 中的 \mu 满足

\mu\sim N(\gamma(t)x,\gamma(t)(1-\gamma(t))I)\\ \mu=\gamma(t)x+\sqrt{\gamma(t)(1-\gamma(t))}\epsilon,\epsilon\sim\mathcal N(0,1)\\ \Rightarrow x=\frac{\mu}{\gamma(t)}-\sqrt{\frac{(1-\gamma(t))}{\gamma(t)}}\epsilon\\

BFN输入 (\mu,t)，输出 \hat\epsilon(\theta,t)， 它作为使用贝叶斯流分布进行后验更新时所使用的噪声向量的估计。变换得到 \hat x(\theta,t)=\frac{\mu}{\gamma(t)}-\sqrt{\frac{(1-\gamma(t))}{\gamma(t)}}\hat\epsilon(\theta,t)，Output Distribution 为

p_O(x|\theta;t)=\delta(x-\hat x(\theta,t))

为式子有意义，在 t\to 0 时令 p_O(x|\theta;t)=0

sender distribution

p_S(y|x;\alpha)=\mathcal N(y|x,\alpha^{-1}I)

receiver distribution

p_R(y|\theta,t,\alpha)=\mathbb E_{\delta(x'-\hat x(\theta,t))}[\mathcal N(y|x',\alpha^{-1}I)]=\mathcal N(y|\hat x(\theta,t),\alpha^{-1}I)

损失函数

重构损失不用于训练。（？）

\begin{aligned}L^r(x)&=\mathbb E_{p_F(\theta|x,t=1)}D_{KL}(\mathcal N(x,\sigma^2I)||N(\hat x(\theta,1),\sigma^2I)))\\&=\mathbb E_{p_F(\theta|x,t=1)}\frac{1}{2\sigma^2}||x-\hat x(\theta,1)||^2\\\alpha_i&=\beta(t_i)-\beta(t_{i-1})\\&=\sigma_1^{\frac{-2i}{n}}-\sigma_1^{\frac{-2(i-1)}{n}}\\&=\sigma_1^{\frac{-2i}{n}}(1-\sigma_1^{\frac{2}{n}})\\L^n(x)&=n\mathbb E_{i\sim U(1,n),p_F(\theta|x,t_{i-1})}D_{KL}(p_S(\cdot|x;\alpha_i)||p_R(\cdot|\theta_{i-1},t_{i-1},\alpha_i))\\&=n\mathbb E_{i\sim U(1,n),p_F(\theta|x,t_{i-1})}D_{KL}(\mathcal N(x,\alpha_i^{-1}I)||(\mathcal N(\hat x(\theta_{i-1},t_{i-1}),\alpha_i^{-1}I))\\&=n\mathbb E_{i\sim U(1,n),p_F(\theta|x,t_{i-1})}\frac{\alpha_i}{2}||x-\hat x(\theta_{i-1},t_{i-1})||^2\\L^{\infty}(x)&=\mathbb E_{t\sim U(0,1),p_F(\theta|x,t)}\alpha(t)\frac{||g(x)-E[P(\theta,t)]||}{2C}\\&=-\ln\sigma_1E_{t\sim U(0,1),p_F(\theta|x,t)}\frac{||x-\hat x(\theta,t)||}{\sigma_1^{2t}}\\g(x)=x,&\quad C=1,\quad E[P(\theta,t)]=\hat x(\theta,t),\quad \alpha(t)=-2(\ln\sigma_1)(\sigma_1^{-2t})\end{aligned}

离散化的（连续）数据

离散化：把连续数据划分成多个等长的数值区间，k(x)=(k(x^{(1)}),\dots,k(x^{(d)}))\in \{1,2,\dots,K\} 表示数据各维分量经过离散化后所处区间的索引。第 k 个区间的左、中、右端点分别用 k_l,k_c,k_r 表示。用 k_c(x) 来表示 x

把数据归一化至 [-1,1] 后， k_c=\frac{2k-1}{K}-1, k_l=k_c-\frac1K, k_r=k_c+\frac1K

对 output distribution 离散化，其他与连续情形相同。

区间宽度为最后时刻的精度 \rho_1=\frac1{\sigma_1^2} 的设置提供了参考

output distribution

在连续数据时 BFN 的输出是其估计的高斯噪声向量。将输出变为两个 D 维向量：估计的高斯噪声分布的均值 \mu_\epsilon，标准差的对数 \ln\sigma_\epsilon。

\hat x(\theta,t)=\frac{\mu}{\gamma(t)}-\sqrt{\frac{(1-\gamma(t))}{\gamma(t)}}\hat\epsilon(\theta,t)\\ \mu_x=\begin{cases}0,\quad t<t_{min}\\\frac{\mu}{\gamma(t)}-\sqrt{\frac{(1-\gamma(t))}{\gamma(t)}}\mu_\epsilon,otherwise\end{cases}\\ \sigma_x=\begin{cases}0,\quad t<t_{min}\\ \sqrt{\frac{(1-\gamma(t))}{\gamma(t)}}\exp(\ln\sigma_\epsilon),otherwise\end{cases}

其中 \mu 是定值，是经过贝叶斯更新后的结果(后验更新值)

对原高斯分布的分布函数 F 做 [-1,1] 的截断，得到离散化后的分布函数

F(x)=\frac12(1+\text{erf}(\frac{x-\mu}{\sqrt2\sigma}))\\ G(x|\mu_x^{(d)},\sigma_x^{(d)})=\begin{cases}0\quad x\le -1\\1\quad x\ge 1\\F(x|\mu_x^{(d)},\sigma_x^{(d)}),otherwise\end{cases}\\ p_O^{(d)}(k|\theta;t):=G(k_r|\mu_x^{(d)},\sigma_x^{(d)})-G(k_l|\mu_x^{(d)},\sigma_x^{(d)})\\ p_O(x|\theta,t)=\prod_{d=1}^Dp_O^{(d)}(k(x^{(d)})|\theta;t)

Receiver Distribution

\begin{aligned}p_R(y|\theta,t,\alpha)&=\mathbb E_{p_O(x'|\theta,t)}\mathcal N(y|k_c(x'),\alpha^{-1}I)\\&=\prod_{d=1}^D\int_{x'}p_O^{(d)}(k(x')|\theta;t)\mathcal N(y^{(d)}|k_c(x'),\alpha^{-1}I)dx'\\&=\prod_{d=1}^D\sum_{k=1}^Kp_O^{(d)}(k|\theta;t)\mathcal N(y^{(d)}|k_c,\alpha^{-1}I)\\\end{aligned}

损失函数

\begin{aligned}L^r&=-\mathbb E_{p_F(\theta|x,t=1)}\ln p_O(x|\theta,t=1)\\&=-\mathbb E_{p_F(\theta|x,t=1)}\sum_{d=1}^D\ln p_O^{(d)}(k(x^{(d)})|\theta;1)\\L^n(x)&=n\mathbb E_{i\sim U(1,n),p_F(\theta|x,t_{i-1})}D_{KL}(p_S(\cdot|x;\alpha_i)||p_R(\cdot|\theta_{i-1},t_{i-1},\alpha_i))\\&=n\mathbb E_{i\sim U(1,n),p_F(\theta|x,t_{i-1}),\mathcal N(y|x,\alpha_i^{-1}I)}\ln\mathcal N(y|x,\alpha_i^{-1}I)-\sum_{d=1}^D\ln(\sum_{k=1}^Kp_O^{(d)}(k|\theta;t_{i-1})\mathcal N(y^{(d)}|k_c,\alpha_{i}^{-1}I))\\L^{\infty}(x)&=-\ln\sigma_1E_{t\sim U(0,1),p_F(\theta|x,t)}\frac{||x-\hat k(\theta,t)||}{\sigma_1^{2t}}\\\end{aligned}

计算 L^n 时显式地从发送者分布中采样噪声样本 y

离散型数据

input distribution

$$ p_I(x|\theta)=\prod_{d=1}^D\theta^{(d)}_k $$ 先验分布 $\theta_0=\frac1K$ 为 $K\times D$ 维向量 #### output distribution $$ p_O^{(d)}(k|\theta;t)=(softmax(\Psi^{(d)}(\theta,t)))_k\\softmax(x)=\frac{e^{x_i}}{\sum_ie^{x_i}}\\p_O(x|\theta;t)=\prod_{d=1}^Dp_O^{(d)}(k|\theta;t) $$ 特别地，对于只有两类的数据，网络在每个维度上只输出一个值 $$ p_O^{(d)}(1|\theta;t)=Sigmoid(\Psi^{(d)}(\theta,t))\\ p_O^{(d)}(2|\theta;t)=1-Sigmoid(\Psi^{(d)}(\theta,t))\\ Sigmoid(x)=\frac1{1+e^{-x}} $$ #### sender distribution $$ p(k^{(d)}|x^{(d)};\omega):=\frac{1-\omega}{K}+\omega\delta_{k^{(d)}x^{(d)}}\\\delta_{ij}=\begin{cases}0,\quad i=j\\1,\quad i\not=j\end{cases}\\ $$ $\alpha(x^{(d)},\omega)=(p(1|x^{(d)};\omega),\dots,p(K|x^{(d)};\omega))$ 为一个 $K$ 类的类别分布 对 $\alpha(x,\omega)$ 进行 $m$ 次独立采样，$c=(c_1,\dots,c_K)\in\{1,\dots,m\}^K$ 表示 $K$ 个类别分别出现了多少次，$c$ 服从多项分布： $$ p(c|x,\omega)=\frac{m!}{c_1!\dots c_K!}\prod_{k=1}^K(\frac{1-\omega}{K}+\omega\delta_{kx})^{c_k} $$ 根据大数定律 $\lim_{m\to\infty}\frac{c_k}{m}=p(k|x,\omega)$，则对任意有限的 $\omega$， $m$ 充分大时可从 $c$ 推断出 $x$。$\omega$ 较小时，$p(k|x,w)$趋于均匀分布， $m$ 要更大才能推断出 $x$。定义精度 $\alpha=m\omega^2$，则 $p(c|x,w)$ 定义了一个从均匀分布开始（$\alpha\to0,\omega\to 0$），连续变化的 sender distribution。根据中心极限定理 （这里怀疑原作者概率论没学好） $$ p:=p(k|x,w)\\\frac{c-mp}{\sqrt{mp(1-p)}}\sim\mathcal{N}(0,I)\\\lim_{m\to\infty}p(c_k|x,w)=\mathcal N(mp,mp(1-p))=\frac{1}{\sqrt{2\pi mp(1-p)}}\exp(-\frac{(c_k-mp)^2}{2mp(1-p)})\\\xi:=1+\frac{\omega K}{1-\omega}\\sender\,sample\quad y_k:=(c_k-\frac mK)ln\xi\\ $$ $(c_k-\frac m K)$ **衡量了当前类别被采样到的次数与平均值相比如何**，$y_k$ 是类别 $k$ 对应的样本。$y_k>0$ 且越大时越能反映出真实数据的信息。$m\omega^2$ 增大时，数据集中程度增加，$y_k$ 反映出更多真实数据 $x$ 的信息。 $$ \frac{dc_k}{dy_k}=\frac{1}{\ln\xi}\\p(y_k|x,w)=|\frac{dc_k}{dy_k}|p(c_k|x,w)=\frac{1}{\ln\xi\sqrt{2\pi mp(1-p)}}\exp(-\frac{(\frac{y_k}{\ln\xi}+\frac mK-mp)^2}{2mp(1-p)})\\=\frac{1}{\frac1\omega\ln\xi\sqrt{2\pi \alpha p(1-p)}}\exp(-\frac{\big[\frac{y_k}{\frac1\omega\ln\xi}+\frac{\alpha}{\omega}(\frac1K-p)\big]^2 }{2\alpha p(1-p)})\\=\frac{1}{\frac1\omega\ln\xi\sqrt{2\pi \alpha p(1-p)}}\exp(-\frac{\big[\frac{y_k}{\frac1\omega\ln\xi}+\alpha(\frac1 K-\delta_{k^{(d)}x^{(d)}})\big]^2 }{2\alpha p(1-p)})\\\ln\xi=\ln(1+\frac{\omega K}{1-\omega})=\frac{\omega K}{1-\omega}+O(\omega^2)\\\lim_{\omega\to 0}\frac1\omega \ln\xi=K\\\lim_{\omega\to 0}p(k|x,\omega)=\frac{1}{K}\\p_s(y_k|x,\omega):=\lim_{\omega\to 0}p(y_k|x,w)=\frac{1}{\sqrt{2\pi \alpha (K-1)}}\exp(-\frac{\big[y_k-\alpha(K\delta_{k^{(d)}x^{(d)}}-1)\big]^2 }{2\alpha (K-1)})\\=\mathcal N(\alpha(K\delta_{k^{(d)}x^{(d)}}-1),\alpha(K-1))\\p_s(y|x,\omega)=\mathcal N(y|\alpha(Ke_X-1),\alpha(K-1)I)\\ $$ #### receiver distribution $$ p_R^{(d)}(y^{(d)}|\theta,t,\alpha)=\sum_{k=1}^K(softmax(\Psi^{(d)}(\theta,t)))_k\mathcal N(\alpha(Ke_k-1),\alpha(K-1)I)\\p_R(y|\theta,t,\alpha)=\prod_{d=1}^Dp_R^{(d)}(y^{(d)}|\theta,t,\alpha) $$ 每一维都是混合高斯分布 #### 精度可加性 贝叶斯更新函数（即后验概率） $h(\theta_{i-1},y)_k=p(k|c;\omega)=\frac{(\theta_{i-1})_kp(c|x=k)}{\sum_{k'=1}^K(\theta_{i-1})_{k'}p(c|x=k')}=\frac{e^{y_k}(\theta_{i-1})_k}{\sum_{k'=1}^Ke^{y_k'}(\theta_{i-1})_{k‘}}$，直接将$p(c|x,\omega),y_k $ 代入计算得到 $$ \begin{aligned}h(y_b,h(y_a,\theta_{i-2}))&=\frac{\exp(y_b)\frac{\exp(y_a)\theta_{i-2}}{\sum_{k'=1}^K\exp((y_a)_{k'})(\theta_{i-2})_{k'}}}{\sum_{k=1}^K\exp((y_b)_k)\frac{\exp((y_a)_k)(\theta_{i-2})_k}{\sum_{k'=1}^K\exp((y_a)_{k'})(\theta_{i-2})_{k'}}}\\&=\frac{\exp(y_b)\exp(y_a)\theta_{i-2}}{\sum_{k=1}^K\exp((y_b)_k)\exp((y_a)_k)(\theta_{i-2})_k}\\&=\frac{\exp(y_b+y_a)\theta_{i-2}}{\sum_{k=1}^K\exp((y_b+y_a)_k)(\theta_{i-2})_k}\\&=h(y_a+y_b,\theta_{i-2})\\y_a+y_b&\sim\mathcal N((\alpha_a+\alpha_b)(Ke_X-1),(\alpha_a+\alpha_b)(K-1)I)\\\mathbb E_{p_U(\theta_{i-1}|\theta_{i-2},x,\alpha_{a})}p_U(\theta_i|\theta_{i-1},x,\alpha_b)&=p_U(\theta_i|\theta_{i-2},x,\alpha_{a}+\alpha_b)\\\end{aligned} $$ #### 精度设置 $$ \beta(t)=t^2\beta(1)\\\alpha(t)=2t\beta(1) $$ #### Bayesian Flow Distribution $$ p_F(\theta|x,t)=\mathbb E_{\mathcal N(y|\beta(t)(Ke_X-1),\beta(t)(K-1)I)}\delta(\theta-\frac{e^{y}\theta_{0}}{\sum_{k'=1}^Ke^{y_k'}(\theta_{0})_{k'}})\\\theta_0=\frac 1K\Rightarrow p_F(\theta|x,t)=\mathbb E_{\mathcal N(y|\beta(t)(Ke_X-1),\beta(t)(K-1)I)}\delta(\theta-softmax(y))\\ $$ 只要先采样 $y\sim N(\beta(t)(Ke_X-1),\beta(t)(K-1)I)$，再令 $\theta=softmax(y)$ 就能得到后验更新 随着 $\alpha$ 增大，$y\to x,\theta\to e_X$，神经网络得到的关于 $x$ 的信息越来越多 #### 损失函数 $$ \begin{aligned}L^r&=-\mathbb E_{p_F(\theta|x,t=1)}\ln p_O(x|\theta,t=1)\\\alpha_i&=\beta(t_i)-\beta(t_{i-1})\\&=\beta(1)((\frac{i}{n})^2-(\frac{i-1}{n})^2)\\&=\beta(1)\frac{2i-1}{n^2}\\L^n(x)&=n\mathbb E_{i\sim U(1,n),p_F(\theta|x,t_{i-1})}D_{KL}(p_S(\cdot|x;\alpha_i)||p_R(\cdot|\theta_{i-1},t_{i-1},\alpha_i))\\&=n\mathbb E_{i\sim U(1,n),p_F(\theta|x,t_{i-1})}D_{KL}\big(\mathcal N(\alpha_i(Ke_X-1),\alpha_i(K-1)I)||\prod_{d=1}^D\sum_{k=1}^K(softmax(\Psi^{(d)}(\theta,t_{i-1})))_k\mathcal N(\alpha_i(Ke_k-1),\alpha_i(K-1)I\big)\\&=n\mathbb E_{i\sim U(1,n),p_F(\theta|x,t_{i-1}),\mathcal N(\alpha_i(Ke_X-1),\alpha_i(K-1)I)}\big[\ln \mathcal N(\alpha_i(Ke_X-1),\alpha_i(K-1)I)-\sum_{d=1}^D\ln(\sum_{k=1}^K(softmax(\Psi^{(d)}(\theta,t_{i-1})))_k\mathcal N(\alpha_i(Ke_k-1),\alpha_i(K-1)I))\big]\\L^{\infty}(x)&=\mathbb E_{t\sim U(0,1),p_F(\theta|x,t)}\alpha(t)\frac{||g(x)-E[P(\theta,t)]||}{2C}\\v:&=\frac{y}{\alpha}+1\\y^{(d)}&\sim p_S(\cdot|x^{(d)};\alpha_i)\Rightarrow v^{(d)}\sim\mathcal N(Ke_{x^{(d)}},\frac{K-1}{\alpha}I)\\y^{(d)}&\sim p_R(\cdot|\theta_{i-1},t_{i-1},\alpha_i)\Rightarrow v^{(d)}\sim \sum_{k=1}^K(softmax(\Psi^{(d)}(\theta,t_{i-1})))_k\mathcal N(Ke_{k},\frac{K-1}{\alpha}I)=K\sum_{k=1}^K(softmax(\Psi^{(d)}(\theta,t_{i-1})))_k\delta(\cdot-e_{k})*\mathcal N(0,\frac{K-1}{\alpha}I)\\P^{(d)}(\theta,t)&=K\sum_{k=1}^K(softmax(\Psi^{(d)}(\theta,t_{i-1})))_k\delta(\cdot-e_{k})\\g(x)=Ke_X,&\quad C=K-1,\quad E[P^{(d)}(\theta,t)]=K\sum_{k=1}^K(softmax(\Psi^{(d)}(\theta,t_{i-1})))_ke_{k},\quad \alpha(t)=2t\beta(1)\\L^{\infty}(x)&=K^2\beta(1)\mathbb E_{t\sim U(0,1),p_F(\theta|x,t)}t\frac{||e_X-\sum_{k=1}^K(softmax(\Psi^{(d)}(\theta,t_{i-1})))_ke_{k}||}{K-1}\\\end{aligned} $$ ### 参考文献 * [Bayesian Flow Networks (一)：生成即压缩！结合贝叶斯统计与深度学习的生成模型 —— 贝叶斯流网络 - 知乎 (zhihu.com)](https://zhuanlan.zhihu.com/p/659111986) * [2308.07037 (arxiv.org)](https://arxiv.org/pdf/2308.07037)