Sending Samples Without Bits-Back

作者:John Schulman(OpenAI)

譯者:朱小虎 Xiaohu (Neil) Zhu(CSAGI / University AI)

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

術(shù)語(yǔ):變分界(Variational Bound); 壓縮(Compression); Monte-Carlo 估計(jì)(Monte-Carlo estimator)

導(dǎo)語(yǔ)

我將給出信息論中的一個(gè)有趣的小問(wèn)題及其基于拒絕式采樣(Rejection Sampling)的解 —— 一個(gè)壓縮算法。這個(gè)問(wèn)題受到變分界的編碼解釋。該壓縮算法不同于常見的 bits-back 編碼并且給出了一個(gè)對(duì)變分界目標(biāo)函數(shù)的更加直接的解釋。不過(guò)它的計(jì)算很低效,但我認(rèn)為它作為原理性證明非常有意思。

問(wèn)題描述 Alice 和 Bob 初始時(shí)認(rèn)同一個(gè)先驗(yàn)分布 p(z),并且他們有共享的隨機(jī)數(shù)生成器(RNG)。然后,Alice 被給予一個(gè)不同的分布 q(z)。Alice 需要多長(zhǎng)的消息發(fā)給 Bob 使得 Bob 通過(guò)組合該消息和 RNG,能夠產(chǎn)生一個(gè)樣本 z\sim q(z)

更準(zhǔn)確地說(shuō),Alice 和 Bob 認(rèn)同一個(gè)關(guān)于 RNG 狀態(tài) \omega 和消息 m 的確定型函數(shù) f(\omega, m)。Alice 計(jì)算消息 m 作為 q 的函數(shù),而 f(\omega, m 的分布必須等于 q(z)

現(xiàn)在來(lái)分類討論問(wèn)題:

  1. 如果 q(z) = p(z),那么消息長(zhǎng)度為零:Alice 可以直接告訴 Bob 從 p 中采樣第一個(gè)樣本。
  2. 如果 q(z) = I[z = z_0],即對(duì)一個(gè)值 z_0 的一個(gè)指示函數(shù),那么 Alice 能做的最好的事情是以消息長(zhǎng)度為 \log p^{-1}(z_0) 發(fā)送 z_0

注意,該問(wèn)題稍微不同于 Alice 采樣 z\sim q(使用一個(gè)非共享 RNG)然后必須發(fā)送其給 Bob。發(fā)送任意的 z 需要期望代碼長(zhǎng)度 E_q[\log p_{-1}(z)]。如果 Alice 采樣 z\sim q 并發(fā)送其給 Bob 使用來(lái)自 p(z) 的代碼,它會(huì)需要超過(guò)必需的比特才可以。特別地,會(huì)使用 S[p] 個(gè)比特使得 q(z)=p(z) 而不是零。

你可能會(huì)猜測(cè)一般的答案是 \text{KL} [q,p],這對(duì)上面例子(1)和(2)是正確的。那個(gè)猜測(cè)是對(duì)的!我會(huì)在下面證明它(在加上一些細(xì)節(jié)后)。但首先,我會(huì)解釋一下這個(gè)通信問(wèn)題的動(dòng)機(jī)。

變分上界

很多概率概念有一個(gè)使用編碼和壓縮的術(shù)語(yǔ)對(duì)應(yīng)的解釋。常用的隱變量模型(如變分自編碼器,VAE)需要的一個(gè)關(guān)鍵想法是變分上界(VUB)。它通常被稱作變分界,但是我們這里對(duì)其符號(hào)取負(fù)因此能對(duì)應(yīng)編碼的長(zhǎng)度。

VUB 是用來(lái)擬合隱變量概率模型的目標(biāo)函數(shù)。給定一個(gè)模型 p(x,z) = p(z) p(x|z),我們一般想要去最大化 \log p(x),但是這里會(huì)出現(xiàn)一個(gè)難解的對(duì) z 上的求和。VUB 引入了一個(gè)樣本分布 q(z),給出了關(guān)于 log-loss 的一個(gè)上界。

\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) 時(shí)出現(xiàn),并且訓(xùn)練(如 VAE)包含關(guān)于 pq 聯(lián)合最小化右式。

不等式的左邊讀作“Alice 需要的比特?cái)?shù)目發(fā)送給 Bob 來(lái)傳遞 x,給定他們之前認(rèn)同概率分布 p(x)”。右式是否能被解釋為一個(gè)對(duì) x 的包含一個(gè)編碼 z 部分編碼 x 的具體壓縮模式呢?

我們想要如下表述:

  • (*) 是Alice 必須發(fā)送給 Bob 的比特?cái)?shù)目,所以 Bob 獲得一個(gè)樣本 z\sim q
  • (**) 是對(duì) Bob 的期望代價(jià)來(lái)完全重構(gòu) x,給定他收到的 z

第二點(diǎn),將 E_{z\sim q}[\log p^{-1}(x|z)] 解釋為給定 z 時(shí)的 x 的編碼長(zhǎng)度,顯然正確。第一點(diǎn)卻是非易見的,它其實(shí)就是上面給出的問(wèn)題。

變分上界實(shí)際上是在一個(gè)著名的被稱為 bits-back 的壓縮模式的編碼長(zhǎng)度。然而,bits-back 編碼并不非常匹配上面給出的簡(jiǎn)單的兩部分編碼解釋。在 bits-back 編碼中,Alice 采樣 z 使用某些附加數(shù)據(jù)作為熵的來(lái)源,然后發(fā)送完全的樣本給 Bob 以期望代價(jià) E_{z\sim q}[\log p^{-1}(z)]。然后她以代價(jià) \log p ^{?1}(x∣z) 發(fā)送 x。最后使用 x 來(lái)推斷分布 q,Bob 恢復(fù)附加數(shù)據(jù)的 E_{z\sim q}[\log p^{-1}(z)] 個(gè)比特,給出一個(gè)凈編碼長(zhǎng)度 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 是非常巧妙的,但我總是想知道是否這樣的兩部編碼的解釋可以被直接實(shí)現(xiàn),編碼字 z 是否可以使用代價(jià) \text{KL}[q,p] 通信而不需要使用 x

普通的拒絕式采樣(次最優(yōu))

我們回到問(wèn)題本身,Alice 需要發(fā)送給 Bob 一個(gè)消息讓他采樣 z\sim q。一個(gè)自然的想法是使用 拒絕式采樣。拒絕式采樣讓人可以通過(guò)隨機(jī)地過(guò)濾從一個(gè)不同的分布 p(z) 的采樣的樣本的方式來(lái)從 q(z) 中采樣。Alice 使用她的 RNG 生成一個(gè)來(lái)自 p(z) 的 IID 樣本序列,但是作用了拒絕式采樣使得第一個(gè)接受的樣本是從 q(z) 中的一個(gè)樣本。然后她將第一個(gè)接受的樣本的索引發(fā)送給 Bob。Bob 在自己這里運(yùn)行同樣的步驟獲得第 n 給樣本,這個(gè)和 Alice 的第 n 個(gè)樣本一樣,因?yàn)槎际怯昧素暙I(xiàn)的 RNG。

現(xiàn)在我們看一下這個(gè)協(xié)議的期望編碼長(zhǎng)度。在拒絕式采樣中,我們采樣 z\sim p(z) 并以概率 p(\text{accept}\ z) = \frac{1}{M} \frac{q(z)}{p(z)} 接受,其中 M = \max_z \frac{q(z)}{p(z)}

接受一個(gè)樣本的概率是 E_z[p(\text{accept}\ z)] = \frac{1}{M}。給定一個(gè)事件其概率為 \epsilon,直到它出現(xiàn)的期望樣本數(shù)目為 \frac{1}{\epsilon}。因此,RS 過(guò)程的實(shí)驗(yàn)的期望數(shù)目就是 M。該整數(shù)的編碼長(zhǎng)度為 \log M = \log \max_z \frac{q(z)}{p(z)} = \max_z \log\frac{q(z)}{p(z)}。

因此,這個(gè)拒絕式采樣過(guò)程獲得了一個(gè)編碼長(zhǎng)度 \max_z \log \frac{q(z)}{p(z)}。但是我們已經(jīng)聲明最優(yōu)的編碼長(zhǎng)度是 \text{KL}[q,p] = E_{z\sim q} [\log \frac{q(z)}{p(z)}]。所以拒絕式采樣是次最優(yōu)的,通過(guò)使用 \max_q 最大化而不是期望 E_q。

修正的拒絕式采樣(次最優(yōu))

拒絕式采樣方法總是可用,但是它給出的是次最優(yōu)的編碼長(zhǎng)度。像很多編碼理論里面的想法一樣,我們可以通過(guò)組合一堆消息在一起并使用大數(shù)定律的方式解決此問(wèn)題。我們將 n 個(gè)樣本組合起來(lái)并以\log M \approx n\text{KL}[q,p] 做拒絕式采樣。

讓我們修改通信問(wèn)題讓 Alice 每次發(fā)給 Bob n 個(gè)樣本。在修改后的問(wèn)題中,Alice 和 Bob 認(rèn)同 (p_1, p_2, \dots, p_n);Alice 需要發(fā)給 Bob 一個(gè)采樣自 (q_1, q_2, \dots, q_n) 的樣本 (z_1, z_2, \dots, z_n)。為了簡(jiǎn)化這個(gè)論斷,讓所有的分布相同,即 p_1 = p_2 = \dots = p_n; q_1 = q_2 = \dots = q_n。這個(gè)論斷容易被修改成分布不同的情形。

符號(hào)使用上,我們令 z^n 表示樣本的一個(gè) n-元組且令 p^n(z^n) 表示 n-元組上的聯(lián)合分布。

我們?nèi)缦露x通信協(xié)議。Alice 重復(fù)采樣 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 是一個(gè)小的數(shù)字,當(dāng) n \rightarrow \infty 時(shí) \rightarrow 0。然后她發(fā)送給 Bob 一個(gè)整數(shù) k,即直到接受的試驗(yàn)次數(shù),最后他從 p^n 中采樣第 k 個(gè)樣本。

為了展示該協(xié)議是可行的,我們接下來(lái)證明兩個(gè)命題:

  1. 每個(gè)樣本 z_i 的期望消息長(zhǎng)度在 n \rightarrow \infty 時(shí)趨向于 \text{KL}[q,p]
  2. 每個(gè)被解碼的 z_iq(z) 的之間的 total variation divergence 在 n \rightarrow \infty 時(shí)趨向于零。

證明將基于 Shannon 引入的典型集合想法。我們也會(huì)解釋如何修改協(xié)議來(lái)發(fā)送準(zhǔn)確的 q 而不是一個(gè)近似(以額外比特?cái)?shù)為代價(jià))。

考慮 \text{log} 比例:

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

對(duì)采樣自 qz,每個(gè)這些項(xiàng)的期望是 E_q[\log \frac{q(z)}{p(z)}]=\text{KL}[q,p]。不嚴(yán)格地說(shuō),該和可能在 n \text{KL}[q,p] \pm O(\sqrt{n}) 附近。讓我們形式化一下。

\epsilon>0 為一個(gè)小的數(shù)。當(dāng) n \rightarrow \infty 時(shí),\log \frac{q(z)}{p(z)} 樣本均值達(dá)到其平均值, \text{KL}[q,p],所以對(duì) 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] - \epsilonz^n 集合 。S 滿足 \Pr(z\sim q \in S) \ge 1-\epsilon。

Alice 通過(guò)采樣 z^n \sim p^n 來(lái)進(jì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}

現(xiàn)在我們計(jì)算接受的概率:

\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}

消息長(zhǎng)度是 \log(\frac{1}{\Pr(\text{accept})})=\log M + O(\epsilon) = n \text{KL}[q,p] + O(\epsilon), 證明了命題的第一部分。

對(duì)命題第二部分,讓我們按照拒絕式采樣協(xié)議下定義 \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。對(duì) 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)

更進(jìn)一步,1 \geq P_{S|\tilde{q}} \geq P_{S|q} \geq 1- \epsilon?;镜挠?jì)算得出 total variation divergence 為 O(\epsilon)。這就證明了第二部分。

最后,不太令人滿意的是 Alice 沒(méi)有完全準(zhǔn)確地發(fā)送 q^n,而是發(fā)送了近似 \tilde{q}。我們可以簡(jiǎn)單地修正此問(wèn)題,讓Alice 準(zhǔn)確地發(fā)送 q^n 花去額外的比特作為代價(jià)。這里是大概過(guò)程。以概率 P_{S|q},我們執(zhí)行上面的協(xié)議。以概率 1-P_{S|q}, Alice 直接發(fā)送 z^n 采樣 S 的補(bǔ)集, 以代價(jià) -\log p^n(z^n)??傊~外的代價(jià)是 O(\epsilon)。

討論$

此過(guò)程計(jì)算難解,因?yàn)樾枰?Alice 從一個(gè)指數(shù)級(jí)大的元組集合 (z_1, z_2, \dots, z_n) 中生成一個(gè)樣本的序列。這與 bits-back 編碼沖突,因?yàn)檫@可以被高效地實(shí)現(xiàn)。實(shí)際上,最近的一篇論文 展示了如何使用 VAE 來(lái)實(shí)現(xiàn) bits-back 編碼,通過(guò)一種 ANS 技術(shù) (一個(gè)類似算術(shù)編碼的方法)。

所以很可能存在一種像算術(shù)編碼的方式解決我們的問(wèn)題,給出一個(gè)高效的算法使得 z 在一個(gè)小的離散集合中。如果 z 是高維的,那么看起來(lái)在沒(méi)有額外假設(shè)的情況下解決高效地傳遞問(wèn)題似乎不太可能——我們能做的就是枚舉從 p 中采樣的樣本并給予索引。

最后,我對(duì)此問(wèn)題若是已經(jīng)很知名表示毫不懷疑—— 因?yàn)檫@看起來(lái)就像一個(gè)形式化有損數(shù)據(jù)傳遞想法方式。若是這樣請(qǐng)告知我。

References

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

最后編輯于
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時(shí)請(qǐng)結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點(diǎn),簡(jiǎn)書系信息發(fā)布平臺(tái),僅提供信息存儲(chǔ)服務(wù)。

友情鏈接更多精彩內(nèi)容