文件更小,質量更高,大火的Stable Diffusion還能壓縮圖像?
近來,Stable Diffusion 成為一個新興的研究方向。一位名為 Matthias Bühlmann 的博主嘗試實驗探究這種模型的強大功能,結果發現 Stable Diffusion 是一個非常強大的有損圖像壓縮編解碼器。他撰寫了一篇博客描述了這個實驗分析過程,以下是博客原文。
首先 Matthias Bühlmann 給出在高壓縮因子條件下,Stable Diffusion 方法與 JPG、WebP 的壓縮結果,所有結果都是 512x512 像素的分辨率:
舊金山風景圖,從左至右:JPG (6.16kB), WebP (6.80kB), Stable Diffusion: (4.96kB)。
糖果店,從左至右:JPG (5.68kB), WebP (5.71kB), Stable Diffusion (4.98kB)。
動物照片,從左至右:JPG (5.66kB), WebP (6.74kB), Stable Diffusion (4.97kB)。
這些例子明顯表明,與 JPG 和 WebP 相比,使用 Stable Diffusion 壓縮圖像可以在更小的文件大小下保留更出色的圖像質量。
探究實驗
Matthias Bühlmann 分析了一下其中的工作原理,Stable Diffusion 使用三個串聯的訓練好的人工神經網絡:
- 變分自編碼器(Variational Auto Encoder,VAE)
- U-Net
- 文本編碼器(Text Encoder)
VAE 將圖像空間中的圖像編碼和解碼為某種潛在的空間表征。源圖像(512 x 512,3x8 或 4x8 bit)的潛在空間表征會分辨率更低(64 x 64)、精度更高(4x32 bit)。
VAE 在訓練過程中自行學習,隨著模型的逐步訓練,不同版本模型的潛在空間表征看起來可能會有所不同,例如 Stable Diffusion v1.4 的潛在空間表征如下(重映射為 4-channel 彩色圖像):
當重新擴展和將潛在特征解釋為顏色值(使用 alpha channel)時,圖像的主要特征仍然可見,并且 VAE 還將更高分辨率的特征編碼到像素值中。
例如,通過一次 VAE 編碼 / 解碼 roundtrip 得到如下結果:
值得注意的是,這種 roundtrip 不是無損的。例如,圖中藍色帶子上白色的字在解碼后可讀性稍差了一些。Stable Diffusion v1.4 模型的 VAE 一般不太擅長表征小型文本和人臉。
我們知道,Stable Diffusion 的主要用途是根據文本描述生成圖像,這就要求該模型要對圖像的潛在空間表征進行操作。該模型使用經過訓練的 U-Net 迭代地對潛在空間圖像進行去噪,輸出它在噪聲中「看到」(預測)的內容,類似于我們有時把云看成某種形狀或面孔。在迭代去噪步驟中,第三個 ML 模型(文本編碼器)指導 U-Net 來嘗試看到不同的信息。
Matthias Bühlmann 分析了 VAE 生成的潛在表征(latent representation)是如何進行有效壓縮的。他發現對 VAE 中的潛在表征進行采樣或對潛在表征應用已有的有損圖像壓縮方法,都會極大地降低重構圖像的質量,而 VAE 解碼過程似乎對潛在表征的質量魯棒性較高。
Matthias Bühlmann 將潛在表征從浮點數量化為 8-bit 無符號整數,結果發現只有非常小的重構誤差。如下圖所示,左:32-bit 浮點潛在表征;中:ground truth;右:8-bit 整數潛在表征。
他還發現通過 palette 和抖動算法進一步量化,得到的結果會出乎意料的好。然而,當直接使用 VAE 解碼時,palettized 表征會導致一些可見的偽影:
左:32-bit 潛在表征;中:8-bit 量化潛在表征;右:帶有 Floyd-Steinberg 抖動的 palettized 8-bit 潛在表征
帶有 Floyd-Steinberg 抖動的 palettized 表征引入了噪聲,使解碼結果失真。于是 Matthias Bühlmann 使用 U-Net 來去除抖動帶來的噪聲。經過 4 次迭代,重構結果在視覺上非常接近未量化的版本:
重構結果(左:帶有 Floyd-Steinberg 抖動的 palettized 表征;中:經過四次迭代去噪;右:Ground Truth)。
雖然結果非常好,但還是會引入一些偽影,例如上圖中心形符號上的光澤陰影。
雖然從主觀上看,Stable Diffusion 壓縮圖像的結果比 JPG 和 WebP 好很多,但從 PSNR、SSIM 等指標看,Stable Diffusion 并沒有明顯的優勢。
如下圖所示,雖然作為編解碼器的 Stable Diffusion 在保留圖像粒度方面比其他方法要好得多,但受壓縮偽影的影響,圖像中物體形狀等特征可能會發生變化。
左:JPG 壓縮;中:Ground Truth;右:Stable Diffusion 壓縮。
值得注意的是,當前的 Stable Diffusion v1.4 模型在壓縮過程中無法很好地保留字體很小的文本信息和人臉特征,但 Stable Diffusion v1.5 模型在人臉生成方面有所改進。
左:Ground Truth;中:經過 VAE roundtrip (32-bit 潛在特征) ;右:從 palettized 去噪 8-bit 潛在特征解碼的結果。
博客發布后,Matthias Bühlmann 的實驗分析引起了大家的討論。
Matthias Bühlmann 自己認為 Stable Diffusion 的圖像壓縮效果比預期好,U-Net 似乎能夠有效消除抖動引入的噪聲。不過,Stable Diffusion 模型未來的版本可能不會再有這種圖像壓縮特性。
然而有網友質疑道:「VAE 本身就被用于圖像壓縮」,例如基于 Transformer 的圖像壓縮方法 TIC 就用到了 VAE 架構,所以 Matthias Bühlmann 的實驗似乎是大材小用了。
對此,你有什么看法?