变分推断—— 进阶

Posted 2020-12-12 kai-nutshell

贝叶斯推断
由上一篇我们已经了解到，对于未知的分布或者难以计算的问题，我们可以通过变分推断将其转换为简单的可计算的问题来求解。现在我们贝叶斯统计的角度，来看一个难以准确计算的案例。

推断问题可以理解为计算条件概率$p(y|x)$。利用贝叶斯定理，可以将计算条件概率（或者说后验概率，posterior）转换为计算联合概率（joint distribution）和先验概率（prior）：
egin{equation} p(y|x) = p(y,x) / p(x) label{1.3} end{equation}

其中$p(y,x)$为联合概率密度，$p(y,x)=p(x|y)p(y)=p(y|x)p(x)$，$p(x)$为边缘概率密度（marginal density），又称为evidence，或者x的先验。边缘概率由对联合概率的其他值做积分得到：
egin{equation} p(x) = int{p(y,x)} dy onumber end{equation}
式$( ef{1.3})$的方法就是贝叶斯推断。朴素贝叶斯模型（naïve Bayesian model）正是基于贝叶斯推断的模型。

朴素贝叶斯是文本分类的经典模型。在文本分类任务中，$y$为文本的类别，$x$为文本，文本有特征$x_1?x_n$，这些特征可能是某个词或词组，那么：
egin{align} p(y,x) &= p(y)p(x | y) onumber \\ &= p(y)p(x_1,?,x_2 |y) onumber end{align}
朴素贝叶斯中假设特征$x_1?x_n$是相互独立的，所以：
egin{equation} p(y,x) = p(y) prod_i {p(x_i |y)} onumber end{equation}
在进行推断时，$p(x)$会被忽略，因为文本的特征的分布对文本的类别没有影响，对所有文本，它都是固定的$p(x)$，也就是说$( ef{1.3})$中等号右边的分母对所有文本类型来说都相同，判断文本类型只需要分子部分。所以要计算p(y│x)，只需要计算$p(y)$和$p(x_i |y)$。$p(y)$和$p(x_i |y)$的计算如下：
egin{align} &p(y) = frac{|D_y |}{|D|} label{1.4} \\ &p(x_i | y) = frac{|D_(y,x_i )|}{|D_y|} onumber end{align} 
其中$D$为文本总数，$D_y$为某个类别的文本的数量，而$D_(y,x_i )$为属于某个类别且包含特征$x_i$的文本的数量。

 

$p(x)$的问题
上面用朴素贝叶斯进行文本分类的例子中，我们忽略了$p(x)$，从而使得贝叶斯推断可行。对于一些模型，$p(x)$是不能被忽略，而且计算$p(x)$的时间复杂度非常高，在这种情况下，贝叶斯推断$p(y|x)$不能直接求解。例如，我们假设数据$x$服从某种混合高斯分布（mixture of Gaussian）：
egin{align} &x_i |c_i,mu sim N(c_i^T mu,1) onumber \\ &c_i sim Categorical(frac{1}{K},?,frac{1}{K}) onumber \\ &mu_k sim N(0, sigma^2) onumber end{align}
其中$x_i$表示$n$个数据中的某个数据，它服从高斯分布$N(c_i^T mu,1)$，$c_i$是一个向量，向量中元素的取值为0或1，点乘$c_i^T mu$表示取均值向量$mu$中的某个均值，例如$mu_k$，而$mu_k$服从高斯分布$N(0,sigma^2)$，$sigma^2$是一个超参数（hyperparameter），由人为设定。我们可以想象对所有数据$x_i$进行聚类（clustering），并假设共有$K$类（cluster），那么某个数据$x_i$所属的类$c_i$为$0~K$中的某一类$k$，且属于第$k$类的概率为$frac{1}{K}$，而第$k$类的数据的均值是$mu_k$，方差设为1（可以假设$x_i$都是经过归一化处理的数据，值落在[0,1]区间）。对于这样的数据，我们可以用高斯混合模型（Gaussian Mixture Model，GMM）对其进行聚类——一种迭代的方法，和EM（Expectation Maximization，期望最大）方法有关联，以后会介绍EM——数据$x$是观测量，而它所属的类以及类的均值是隐变量$y={c_i,mu}$，所以聚类也是在求$p(y|x)$。
上面描述的数据$x$，它与$c_i$和$mu$的联合概率密度函数为：
egin{equation} p(x,c,mu) = p(mu) prod_{i=1}^n{p(c_i )p(x_i|c_i,μ)} label{1.5} end{equation}

其中$c={c_1,?,c_n}$。因为数据$x$是独立同分布的，$c_i$是$x_i$的局部变量，也就是说一个$x_i$对应一个$c_i$，与其他$c_i$没有关系，而$mu$是一个全局变量，与$i$无关，所以不能在乘积中被拆分为$mu_i$。计算联合密度$p(x,c,mu)$时，需要计算所有数据对应的$p(c_i)p(x_i |c_i,mu)$的乘积，因为是$mu$是一个全局变量，不能拆分，所以联合密度的计算为$( ef{1.5})$的形式。
因为$c_i$的先验是离散的（discrete），像1、2、3这样的“孤立”值，所以计算$x$的边缘概率密度$p(x)$时，要通过求和覆盖$c_i$可能的取值，而$mu$是连续的（continuous），要对其积分，所以$x$的边缘概率密度为：
egin{equation} p(x) = int{p(mu) prod_{i=1}^n{sum_{c_i}{p(c_i )p(x_i│c_i,mu)}}} dμ label{1.6} end{equation}

假设数据$x$有$K$个类，$( ef{1.6})$的时间复杂度（time complexity）就是$O(K^n)$。这是一个指数级的复杂度。随着数据量$n$的增加，计算将无法完成。在这种情况下，我们是无法用贝叶斯推断$( ef{1.3})$去计算后验$p(y|x)$的。

 

ELBO
上面我们已经从贝叶斯统计的角度了解到，条件概率$p(y|x)$难以计算的原因在于$x$的先验$p(x)$，下面我们还会看到，在以KL散度作为变分分布训练的目标函数时，$p(x)$同样会给计算带来麻烦。

当我们设计一个函数$q(y)$或者$q(y|x)$作为变分函数，然后使它接近真实的后验时，我们可以用KL散度来衡量它和真实分布$p(y|x)$的差异。，假设部分分布是$q(y)$，那么变分推断的目标是求最优的变分分布函数$q^* (y)$：
egin{equation} q^* (y) = argmin_{q(y)in L} {KL(q(y) parallel p(y|x))} label{1.7} end{equation}

其中$argmin_{q(y) in L}$表示从$L$中找出使得KL散度最小的$q(y)$，而$L$是所有可能的$q(y)$的族（family）。族的复杂度决定了优化问题（optimization）的复杂度。
其实，无论选择什么样的族，$( ef{1.7})$仍然是难以计算的，因为KL散度的计算如下：
egin{equation} KL(q(y) parallel p(y|x)) = E_q [log{?q(y)} ] - E_q [log{p(y|x)} ] label{1.8} end{equation}

而就像前面分析的，$p(y|x)$是计算的困难所在。为了使优化问题能够计算，我们需要将$( ef{1.8})$替换掉。

首先，我们对$( ef{1.8})$做一些转换：
egin{align} KL(q(y) parallel p(y|x)) &= E_q [log{?q(y)} ] - E_q [log{frac{p(y,x)}{p(x)}} ] onumber \\ &= E_q [log{?q(y)}] - E_q [log{p(y,x)} ]+ E_q [log{p(x)} ] label{1.9} end{align}
调整$( ef{1.9})$中各项的位置，得到：
egin{align} &E_q [log{p(x)} ]= KL(q(y) parallel p(y|x))- E_q [log{q(y)} ] + E_q [log{p(y,x)}] onumber \\ &log{p(x)} = KL(q(y) parallel p(y|x)) + ELBO(q) label{1.10} end{align}

其中$E_q [log{p(x)} ]=log{p(x)}$，因为从分布$q(y)$中对$y$采样，与$x$无关，$log{p(x)} E_q [1]=log{p(x)}$。$( ef{1.10})$中$ELBO$（Evidence Lower BOund）是log evidence的下界，也就是 $log{p(x)}$的下界，因为KL散度的值只能大于等于0。当$q(y)$等于$p(y|x)$时，KL散度为0，$log{p(x)}=ELBO(q)$，否则$ELBO(q)$小于$( ef{1.10})$等号左边。
实际上，我们只需要优化$ELBO$，因为：
egin{equation} KL(q(y) parallel p(z│x)) = log{p(x)}- ELBO(q) onumber end{equation}
而$log{p(x)}$对优化目标$q$来说是一个常数，不影响$q$的优化。对$ELBO$还可以以下面的方式表示：
egin{align} ELBO(q) &= E_q [log{p(y,x)} ]- E_q [log{q(y)} ] onumber \\ &=  E_q [log{p(x|y)} ] + E_q [log{p(y)} ] - E_q [log{q(y)} ] onumber \\ &= E_q [log{p(x|y)} ] - KL(q(y) parallel p(y)) label{1.11} end{align}

式$( ef{1.11})$等号右边第一项是expected log-likelihood，而第二项是隐变量$y$的变分分布与真实分布的KL散度，其中变分分布$q(y)$是优化的目标。
将原先的目标函数KL散度替换为$ELBO$之后，我们就可以通过优化，找到$q$来进行推断了（提醒一下，变分推断并不局限于贝叶斯推断）。

 

平均场
上面我们已经得到了变分推断模型的目标函数（objective function），也就是$ELBO$，现在我们来看一下如何选择变分分布$q(y)$的族（family）。

在上一篇【变分近似】中，我们了解到论文$[2]$的作者迈克尔·乔丹等人采用Frenchel共轭来确定$q(y)$。 Frenchel共轭要求预先知道$p(x|y)$的表达式，才能将p(x|y)转换为q(y)。这里我们来看。另一种构造变分推断的方法——平均场（mean-field）方法。
在平均场变分族中，隐变量是相互独立的，所以平均场的族可以统一表示为：
egin{equation} q(y) = prod{q_i (y_i)} onumber end{equation}
例如，服从混合高斯分布的数据$x$的隐变量$mu$和$c$满足：
egin{equation} q(mu,c) = prod_k{q(mu_k|m_k,s_k^2 )} prod_i{q(c_i|φ_i )} label{1.12} end{equation}

其中$mu_k$是第$k$类的数据的均值，它服从高斯分布$N(m_k,s_k^2)$，例如$N(0,sigma^2)$，而$c_i$是第$i$个数据对应的类，它服从多项分布，例如$Categorical(frac{1}{K},?,frac{1}{K})$，$varphi_i=frac{1}{K}$。此时，隐变量被替换为$m_k$、$s_k^2$和$varphi_i$等变分参数。变分推断模型的学习过程就是优化这些参数，从而最大化$ELBO$。
mean-field是非常简单的构造变分推断的方法，它假设所有变量是相互独立。这种简单的方法有一个缺点，就是它无法反映隐变量之间的关系。另外，论文$[1]$中指出，mean-field方法会低估真实分布的方差（variance），因为KL散度是不对称的，$KL(q parallel p)$不等于$KL(p parallel q)$，而$KL(q parallel p)$会对$q$不落在分布$p$的惩罚很大，而忽略$q$没有覆盖$p$的情况，因为：
当$q o 0$而$p o 1$，此时$q log{frac{q}{p}} o 0$，KL的影响趋于零；
当$q o 1$而$p o 0$，此时$q log{frac{q}{p}} o + infty$，KL的值非常大。
这在【GAN基础】中有分析，个人认为它是KL散度的问题，而不是mean-field本身的缺陷。

 

CAVI
现在我们已经知道了变分推断的目标函数$ELBO$，包括其中的变分分布$q(y)$，现在我们来理解变分参数的求解过程。论文$[1]$中采用的是CAVI（Coordinate Ascent Variational Inference）算法。
CAVI是一种迭代方法。在CAVI中，每次只更新一个隐变量，其他隐变量作为输入求这个变量，也就是$p(y_j |y_{-j},x)$，其中$y_j$是要计算的隐变量，$y_{-j}$是其他隐变量。例如在$( ef{1.12})$中，$y_j$是$m_k$，而$y_{-j}$代表$s_k^2$、$varphi_i$，以及$m_{-k}$。对式$( ef{1.11})$的第一行进行调整，使它只优化$q_j$，也就是固定其他隐变量，计算$y_j$，并使ELBO最大化：
egin{equation} ELBO(q_j )= E_j [E_{-j} [log{p(y_j,y_{-j},x)}]]- E_j [log{q_j (y_j )} ] + const. onumber end{equation}
其中$p(y_j,y_{-j},x)$为各个变量的联合分布密度，$const.$代表一个常数。
经过迭代计算，最终变分推断会收敛于一个局部最优解（local optimum），此时，对于所有隐变量，$ELBO$都是最优。
CAVI虽然能够求解变分推断，但是它的不适用于大规模的数据集，因为它要求每一次迭代时都计算一遍所有的数据，在下一篇，我们将会了解一种随机优化的方法，它可以除了大规模的数据集。

 

未完待续~

 

 

 [1] Blei, D. M., Kucukelbir, A., McAuliffe, J. D. (2018). “variational inference a review for statisticians”.

 [2] Jordan, M. I., Ghahramani, Z., Jaakkola, T., and Saul, L. (1999). “Introduction to variational methods for graphical models”.