基于文本的圖像生成模型火了，出圈的不止有擴散模型，還有開源的Stable Diffusion模型。

最近一位瑞士的軟件工程師Matthias Bühlmann無意間發(fā)現(xiàn)，Stable Diffusion不僅能用來生成圖像，還可以用來壓縮位圖圖像，甚至比JPEG和WebP的壓縮率更高。

比如一張美洲駱駝的照片，原圖為768KB，使用JPEG壓縮到5.66KB，而Stable Diffusion可以進(jìn)一步壓縮到4.98KB，而且能夠保留更多高分辨率的細(xì)節(jié)以及更少的壓縮偽影，肉眼可見地優(yōu)于其他壓縮算法。

不過這種壓縮方式也存在缺陷，即不適合壓縮人臉和文本圖像，在某些情況下，甚至?xí)?span style="color: #ff6827;">生成一些原圖并不存在內(nèi)容。

雖然重新訓(xùn)練一個自編碼器也能做到類似于Stable Diffusion的壓縮效果，但使用Stable Diffusion的一個主要優(yōu)勢在于，有人已經(jīng)投入了上百萬的資金幫你訓(xùn)練了一個，你又何必重新花錢訓(xùn)練一個壓縮模型呢？

Stable Diffusion如何壓縮圖像

擴散模型正在挑戰(zhàn)生成模型的霸主地位，對應(yīng)的開源Stable Diffusion模型也在機器學(xué)習(xí)社區(qū)掀起一場藝術(shù)革命。

Stable Diffusion由三個訓(xùn)練后的神經(jīng)網(wǎng)絡(luò)串聯(lián)得到，即一個變分自編碼器（VAE），U-Net模型和一個文本編碼器。

變分自編碼器對圖像空間中的圖像進(jìn)行編碼和解碼，從而獲得該圖像在潛空間的表征向量，以一個分辨率更低（64x64)具有更高精度(4x32bit)的向量來表示源圖像(3x8或4x8bit的512x512)。

VAE在將圖像編碼到潛空間的訓(xùn)練過程主要依賴自監(jiān)督學(xué)習(xí)，即輸入和輸出都是源圖像，因此隨著模型進(jìn)一步訓(xùn)練，不同版本的模型的潛空間表征可能會看起來不同。

使用Stable Diffusion v1.4的潛空間表征通過重新映射和解釋為4通道彩色圖像后，看起來就是下圖的中間圖像，源圖像中的主要特征仍然可見。

需要注意的是，VAE往返編碼一次并不是無損的。

比如在解碼之后，藍(lán)色帶子上的ANNA名字就沒有源圖像那么清晰了，可讀性顯著降低。

Stable Diffusion v1.4中的變分自編碼器不太擅長表示小文本以及人臉圖像，不知道在v1.5版本中是否會改善。

Stable Diffusion的主要壓縮算法就是利用圖像的這種潛空間表征，從短文本描述中生成新的圖像。

從潛空間表征的隨機噪聲開始，使用充分訓(xùn)練的U-Net迭代去除潛空間圖像的噪聲，用一種更簡單的表征輸出模型認(rèn)為它在這個噪聲中「看到」的預(yù)測，有點像我們在看云的時候，從不規(guī)則的圖形中還原出腦海里的形狀或面孔。

當(dāng)使用Stable Diffusion來生成圖像時，這個迭代去噪步驟是由第三個組件，即文本編碼器引導(dǎo)的，該編碼器為U-Net提供關(guān)于它應(yīng)該嘗試在噪聲中看到什么的信息。

不過對于壓縮任務(wù)來說，并不需要文本編碼器，所以實驗過程只創(chuàng)建了一個空字符串的編碼用于告訴U-Net在圖像重建過程中進(jìn)行非引導(dǎo)去噪。

為了使用Stable Diffusion作為圖像壓縮編解碼器，算法需要有效地壓縮由VAE產(chǎn)生的潛表征。

在實驗中可以發(fā)現(xiàn)，對潛表征進(jìn)行下采樣或者直接使用現(xiàn)有的有損圖像壓縮方法，都會大大降低重建圖像的質(zhì)量。

但作者發(fā)現(xiàn) VAE 的解碼似乎對潛表征的量化（quantization）非常有效。

通過對從浮點到8位無符號整數(shù)的潛量化進(jìn)行縮放、拖拽（clamping）和重新映射，只會產(chǎn)生很小的可見重構(gòu)錯誤。

通過量化8位的潛表征，圖像表示的數(shù)據(jù)大小現(xiàn)在是64*64*4*8bit=16kB ，遠(yuǎn)小于未壓縮源圖像的512*512*3*8bit=768kB

如果潛表征的位數(shù)小于8bit，無法產(chǎn)生比較好的效果。

如果對圖像進(jìn)一步執(zhí)行調(diào)色板（palettizing）和抖動（dithering），則量化效果就會再次提升。

使用256*4*8位向量和Floyd-Steinberg抖動的潛表征創(chuàng)建了一個調(diào)色板表示，使數(shù)據(jù)大小進(jìn)一步壓縮到64*64*8+256*4*8bit=5kB

潛空間調(diào)色板的抖動會引入噪聲，從而扭曲了解碼結(jié)果。但由于Stable Diffusion是基于潛噪聲的去除，所以可以使用U-Net去除抖動引起的噪聲。

經(jīng)過4次迭代，重建結(jié)果在視覺上非常接近未量化的版本。

雖然數(shù)據(jù)量大大減少了（源圖像為壓縮圖像的155倍大），但效果是非常好的，不過也引入了一些偽影（比如原圖的心形圖案中不存在偽影）。

有趣的是，這種壓縮方案引入的偽影對圖像內(nèi)容的影響比對圖像質(zhì)量的影響更大，而且以這種方式壓縮的圖像可能包含這些類型的壓縮偽影。

作者還用zlib對調(diào)色板和索引進(jìn)行了無損壓縮，在測試樣本中，大多數(shù)的壓縮結(jié)果都小于5kb，但這種壓縮方法仍然存在更多的優(yōu)化空間。

為了評估該壓縮編解碼器，作者沒有使用任何在網(wǎng)上找到的標(biāo)準(zhǔn)測試圖像，因為網(wǎng)上的圖像都有可能在Stable Diffusion的訓(xùn)練集中出現(xiàn)過，而壓縮這類圖像可能會導(dǎo)致不公平的對比優(yōu)勢。

為了盡可能公平地進(jìn)行比較，作者使用了Python圖像庫中最高質(zhì)量的編碼器設(shè)置，以及使用mozjpeg庫添加了壓縮后的JPG數(shù)據(jù)的無損數(shù)據(jù)壓縮。

值得注意的是，雖然Stable Diffusion的結(jié)果主觀上看起來比JPG和WebP壓縮的圖像要好得多，但在標(biāo)準(zhǔn)測量指標(biāo)(如PSNR或SSIM)方面，它們并沒有明顯更好，但也沒有更差。

只是引入的偽影類型不那么明顯，因為它們對圖像內(nèi)容的影響大于對圖像質(zhì)量的影響。

這種壓縮方法也有一點危險，雖然重建特征的質(zhì)量很高，但內(nèi)容可能會受到壓縮偽影的影響，即使它看起來非常清晰。

例如，在一張測試圖像中，雖然Stable Diffusion作為編解碼器在保持圖像的質(zhì)量方面要好得多，甚至連相機顆粒紋理（camera grain）都能保留下來(這是大多數(shù)傳統(tǒng)壓縮算法難以做到的) ，但其內(nèi)容仍然受到壓縮偽影的影響，像建筑物形狀這樣的精細(xì)特征可能會發(fā)生變化。

雖然在JPG壓縮圖像中當(dāng)然不可能比在Stable Diffusion壓縮圖像中識別出更多的真實值，但是Stable Diffusion壓縮結(jié)果的高視覺質(zhì)量可能具有欺騙性，因為JPG和WebP中的壓縮偽影更容易識別。

如果你也想動手復(fù)現(xiàn)一遍實驗，作者在Colab上開源了代碼。

代碼鏈接：?https://colab.research.google.com/drive/1Ci1VYHuFJK5eOX9TB0Mq4NsqkeDrMaaH?usp=sharing?

最后，作者表示，文章中設(shè)計的實驗仍然是相當(dāng)淺顯的，但效果仍然令人驚喜，未來仍然有很大的改進(jìn)空間。