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

Sending Samples Without Bits-Back

導語

變分上界

普通的拒絕式採樣（次最優）

修正的拒絕式採樣（次最優）

討論$

References

月章

招幕推薦系統強人

推薦系統 #2 RecSys Best Paper Awards

推薦系統 #1 Fairness and Transparency in Recommendation: The Users’ Perspective 推薦中的公平性和透明度：用戶的視角

The top 6 programming languages for IoT projects 物聯網項目的最佳 6 種編程語言

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結