Sending Samples Without Bits-Back

作者：John Schulman（OpenAI）

译者：朱小虎 Xiaohu (Neil) Zhu（CSAGI / University AI）

原文链接：http://joschu.net/blog/sending-samples.html

术语：变分界（Variational Bound）; 压缩（Compression）; Monte-Carlo 估计（Monte-Carlo estimator）

导语

我将给出信息论中的一个有趣的小问题及其基于拒绝式采样（Rejection Sampling）的解 —— 一个压缩算法。这个问题受到变分界的编码解释。该压缩算法不同于常见的 bits-back 编码并且给出了一个对变分界目标函数的更加直接的解释。不过它的计算很低效，但我认为它作为原理性证明非常有意思。

问题描述 Alice 和 Bob 初始时认同一个先验分布 $p(z)$ ，并且他们有共享的随机数生成器（RNG）。然后，Alice 被给予一个不同的分布 $q(z)$ 。Alice 需要多长的消息发给 Bob 使得 Bob 通过组合该消息和 RNG，能够产生一个样本 $z\sim q(z)$ ？

更准确地说，Alice 和 Bob 认同一个关于 RNG 状态 $\omega$ 和消息 $m$ 的确定型函数 $f(\omega, m)$ 。Alice 计算消息 $m$ 作为 $q$ 的函数，而 $f(\omega, m$ 的分布必须等于 $q(z)$ 。

现在来分类讨论问题：

如果 $q(z) = p(z)$ ，那么消息长度为零：Alice 可以直接告诉 Bob 从 $p$ 中采样第一个样本。
如果 $q(z) = I[z = z_0]$ ，即对一个值 $z_0$ 的一个指示函数，那么 Alice 能做的最好的事情是以消息长度为 $\log p^{-1}(z_0)$ 发送 $z_0$ 。

注意，该问题稍微不同于 Alice 采样 $z\sim q$ （使用一个非共享 RNG）然后必须发送其给 Bob。发送任意的 $z$ 需要期望代码长度 $E_q[\log p_{-1}(z)]$ 。如果 Alice 采样 $z\sim q$ 并发送其给 Bob 使用来自 $p(z)$ 的代码，它会需要超过必需的比特才可以。特别地，会使用 $S[p]$ 个比特使得 $q(z)=p(z)$ 而不是零。

你可能会猜测一般的答案是 $\text{KL} [q,p]$ ，这对上面例子（1）和（2）是正确的。那个猜测是对的！我会在下面证明它（在加上一些细节后）。但首先，我会解释一下这个通信问题的动机。

变分上界

很多概率概念有一个使用编码和压缩的术语对应的解释。常用的隐变量模型（如变分自编码器，VAE）需要的一个关键想法是变分上界（VUB）。它通常被称作变分下界，但是我们这里对其符号取负因此能对应编码的长度。

VUB 是用来拟合隐变量概率模型的目标函数。给定一个模型 $p(x,z) = p(z) p(x|z)$ ，我们一般想要去最大化 $\log p(x)$ ，但是这里会出现一个难解的对 $z$ 上的求和。VUB 引入了一个样本分布 $q(z)$ ，给出了关于 log-loss 的一个上界。

$\log p^{-1} (x) \leq \underbrace{\text{KL}[q,p]}_{(*)} + \underbrace{E_{z\sim q}[\log p^{-1}(x|z)]}_{(**)}$

等式在 $q(z) = p(z|x)$ 时出现，并且训练（如 VAE）包含关于 $p$ 和 $q$ 联合最小化右式。

不等式的左边读作“Alice 需要的比特数目发送给 Bob 来传递 $x$ ，给定他们之前认同概率分布 $p(x)$ ”。右式是否能被解释为一个对 $x$ 的包含一个编码 $z$ 部分编码 $x$ 的具体压缩模式呢？

我们想要如下表述：

(*) 是Alice 必须发送给 Bob 的比特数目，所以 Bob 获得一个样本 $z\sim q$
(**) 是对 Bob 的期望代价来完全重构 $x$ ，给定他收到的 $z$

第二点，将 $E_{z\sim q}[\log p^{-1}(x|z)]$ 解释为给定 $z$ 时的 $x$ 的编码长度，显然正确。第一点却是非易见的，它其实就是上面给出的问题。

变分上界实际上是在一个著名的被称为 bits-back 的压缩模式的编码长度。然而，bits-back 编码并不非常匹配上面给出的简单的两部分编码解释。在 bits-back 编码中，Alice 采样 z 使用某些附加数据作为熵的来源，然后发送完全的样本给 Bob 以期望代价 $E_{z\sim q}[\log p^{-1}(z)]$ 。然后她以代价 $\log p ^{−1}(x∣z)$ 发送 $x$ 。最后使用 $x$ 来推断分布 $q$ ，Bob 恢复附加数据的 $E_{z\sim q}[\log p^{-1}(z)]$ 个比特，给出一个净编码长度 $E_{z\sim q}[\log p^{-1}(z)] + E_{z\sim q}[\log p^{-1}(x|z)]$ - $E_{z\sim q}[\log q^{-1}(z)] = \text{KL}[q,p] + E_q[\log p^{-1}(x|z)]$ 。

Bits-back 是非常巧妙的，但我总是想知道是否这样的两部编码的解释可以被直接实现，编码字 $z$ 是否可以使用代价 $\text{KL}[q,p]$ 通信而不需要使用 $x$ ？

普通的拒绝式采样（次最优）

我们回到问题本身，Alice 需要发送给 Bob 一个消息让他采样 $z\sim q$ 。一个自然的想法是使用拒绝式采样。拒绝式采样让人可以通过随机地过滤从一个不同的分布 $p(z)$ 的采样的样本的方式来从 $q(z)$ 中采样。Alice 使用她的 RNG 生成一个来自 $p(z)$ 的 IID 样本序列，但是作用了拒绝式采样使得第一个接受的样本是从 $q(z)$ 中的一个样本。然后她将第一个接受的样本的索引发送给 Bob。Bob 在自己这里运行同样的步骤获得第 n 给样本，这个和 Alice 的第 n 个样本一样，因为都是用了贡献的 RNG。

现在我们看一下这个协议的期望编码长度。在拒绝式采样中，我们采样 $z\sim p(z)$ 并以概率 $p(\text{accept}\ z) = \frac{1}{M} \frac{q(z)}{p(z)}$ 接受，其中 $M = \max_z \frac{q(z)}{p(z)}$ 。

接受一个样本的概率是 $E_z[p(\text{accept}\ z)] = \frac{1}{M}$ 。给定一个事件其概率为 $\epsilon$ ，直到它出现的期望样本数目为 $\frac{1}{\epsilon}$ 。因此，RS 过程的实验的期望数目就是 $M$ 。该整数的编码长度为 $\log M = \log \max_z \frac{q(z)}{p(z)} = \max_z \log\frac{q(z)}{p(z)}$ 。

因此，这个拒绝式采样过程获得了一个编码长度 $\max_z \log \frac{q(z)}{p(z)}$ 。但是我们已经声明最优的编码长度是 $\text{KL}[q,p] = E_{z\sim q} [\log \frac{q(z)}{p(z)}]$ 。所以拒绝式采样是次最优的，通过使用 $\max_q$ 最大化而不是期望 $E_q$ 。

修正的拒绝式采样（次最优）

拒绝式采样方法总是可用，但是它给出的是次最优的编码长度。像很多编码理论里面的想法一样，我们可以通过组合一堆消息在一起并使用大数定律的方式解决此问题。我们将 $n$ 个样本组合起来并以 $\log M \approx n\text{KL}[q,p]$ 做拒绝式采样。

让我们修改通信问题让 Alice 每次发给 Bob $n$ 个样本。在修改后的问题中，Alice 和 Bob 认同 $(p_1, p_2, \dots, p_n)$ ；Alice 需要发给 Bob 一个采样自 $(q_1, q_2, \dots, q_n)$ 的样本 $(z_1, z_2, \dots, z_n)$ 。为了简化这个论断，让所有的分布相同，即 $p_1 = p_2 = \dots = p_n$ ; $q_1 = q_2 = \dots = q_n$ 。这个论断容易被修改成分布不同的情形。

符号使用上，我们令 $z^n$ 表示样本的一个 $n$ -元组且令 $p^n(z^n)$ 表示 $n$ -元组上的联合分布。

我们如下定义通信协议。Alice 重复采样 $z^n \sim p^n$ 以概率 $\min \left(1, \frac{1}{M} \frac{q^n(z^n)}{p^n(z^n)}\right)$ 接受，其中 $\log M=n(\text{KL}[q,p] + \epsilon)$ 而 $\epsilon$ 是一个小的数字，当 $n \rightarrow \infty$ 时 $\rightarrow 0$ 。然后她发送给 Bob 一个整数 $k$ ，即直到接受的试验次数，最后他从 $p^n$ 中采样第 $k$ 个样本。

为了展示该协议是可行的，我们接下来证明两个命题：

每个样本 $z_i$ 的期望消息长度在 $n \rightarrow \infty$ 时趋向于 $\text{KL}[q,p]$
每个被解码的 $z_i$ 和 $q(z)$ 的之间的 total variation divergence 在 $n \rightarrow \infty$ 时趋向于零。

证明将基于 Shannon 引入的典型集合想法。我们也会解释如何修改协议来发送准确的 $q$ 而不是一个近似（以额外比特数为代价）。

考虑 $\text{log}$ 比例：

$\log \frac{q^n(z^n)}{p^n(z^n)} = \sum_{i=1}^n \log \frac{q(z_i)}{p(z_i)}$

对采样自 $q$ 的 $z$ ，每个这些项的期望是 $E_q[\log \frac{q(z)}{p(z)}]=\text{KL}[q,p]$ 。不严格地说，该和可能在 $n \text{KL}[q,p] \pm O(\sqrt{n})$ 附近。让我们形式化一下。

令 $\epsilon>0$ 为一个小的数。当 $n \rightarrow \infty$ 时， $\log \frac{q(z)}{p(z)}$ 样本均值达到其平均值, $\text{KL}[q,p]$ ，所以对 $z^n \sim q^n$ 有：

$\Pr \left(\frac{1}{n} \log \frac{q^n(z^n)}{p^n(z^n)} \leq \text{KL}[q,p] + \epsilon \right) \ge 1-\epsilon$

令 $S$ 表示满足 $\frac{1}{n}\log \frac{q^n(z^n)}{p^n(z^n)} \le \text{KL}[q,p] - \epsilon$ 的 $z^n$ 集合。 $S$ 满足 $\Pr(z\sim q \in S) \ge 1-\epsilon$ 。

Alice 通过采样 $z^n \sim p^n$ 来进行拒绝式采样然后以概率 $\Pr(\text{accept}\ z^n) = \min(1, \frac{1}{m} \frac{q^n(z^n)}{p^n(z^n)})$ 接受，其中 $\log M = n(\text{KL}[q,p] +\epsilon)$ 。

$\Pr(\text{sample } z^n \sim p^n \text{ and accept }) = \begin{cases} \frac{q^n(z^n)}{M} \qquad z \in S \\ <\frac{q^n(z^n)}{M} \qquad z \notin S \\ \end{cases}$

现在我们计算接受的概率：

$\begin{aligned} \Pr(\text{accept})&=\sum_{z^n}\Pr(\text{sample } z^n \sim p^n \text{ and accept })\\ &=\sum_{z^n \in S} \frac{q^n(z^n)}{M} + \sum_{z \notin S} \text{[positive value]}\\ &\ge \sum_{z^n \in S} \frac{q^n(z^n)}{M}\\ &\ge (1 - \epsilon) / M \end{aligned}$

消息长度是 $\log(\frac{1}{\Pr(\text{accept})})=\log M + O(\epsilon) = n \text{KL}[q,p] + O(\epsilon)$ , 证明了命题的第一部分。

对命题第二部分，让我们按照拒绝式采样协议下定义 $\tilde{q}^N$ 为在 $z^n$ 上的解码分布。

$\tilde{q}^n(z^n) \propto p^n(z^n)P(\text{accept } z^n)$

定义 $q_S$ 为从 $q^n$ 中采样的样本分布，基于条件 $\in S$ 。对 $z^n \in S$ ， $q^n$ 是按照如下形式成比例的：

$q^n(z^n) = q_S(z^n) P_{S|q} \quad \text{where} \quad P_{S|q}=\Pr(z^n \sim q^n \in S)\\ \tilde{q}^n(z^n) = q_S(z^n) P_{S|\tilde{q}} \quad\text{where}\quad P_{S|\tilde{q}} =\Pr(z^n \sim \tilde{q}^n \in S)$

更进一步， $1 \geq P_{S|\tilde{q}} \geq P_{S|q} \geq 1- \epsilon$ 。基本的计算得出 total variation divergence 为 $O(\epsilon)$ 。这就证明了第二部分。

最后，不太令人满意的是 Alice 没有完全准确地发送 $q^n$ ，而是发送了近似 $\tilde{q}$ 。我们可以简单地修正此问题，让Alice 准确地发送 $q^n$ 花去额外的比特作为代价。这里是大概过程。以概率 $P_{S|q}$ ，我们执行上面的协议。以概率 $1-P_{S|q}$ ， Alice 直接发送 $z^n$ 采样 $S$ 的补集, 以代价 $-\log p^n(z^n)$ 。总之，额外的代价是 $O(\epsilon)$ 。

讨论$

此过程计算难解，因为需要 Alice 从一个指数级大的元组集合 $(z_1, z_2, \dots, z_n)$ 中生成一个样本的序列。这与 bits-back 编码冲突，因为这可以被高效地实现。实际上，最近的一篇论文展示了如何使用 VAE 来实现 bits-back 编码，通过一种 ANS 技术（一个类似算术编码的方法）。

所以很可能存在一种像算术编码的方式解决我们的问题，给出一个高效的算法使得 $z$ 在一个小的离散集合中。如果 $z$ 是高维的，那么看起来在没有额外假设的情况下解决高效地传递问题似乎不太可能——我们能做的就是枚举从 $p$ 中采样的样本并给予索引。

最后，我对此问题若是已经很知名表示毫不怀疑—— 因为这看起来就像一个形式化有损数据传递想法方式。若是这样请告知我。

References

https://arxiv.org/pdf/1901.04866.pdf