使用 Stable Signature + IMATAG 的 BZH 進行魯棒影像水印

2023 年 11 月 28 日，Stability AI 釋出了 SDXL Turbo，一款先進的 AI 影像生成模型。僅兩天後，IMATAG 在 🤗 Hugging Face 上提供了一個魯棒且不可見的水印解決方案，用於識別由該 AI 生成的影像。在該演示中，使用者可以透過提供描述所需影像的提示來生成影像。在影像生成過程中，會插入一個不可感知的水印，以證明其為合成影像。該演示還允許在對生成的影像進行各種攻擊（例如壓縮或重新著色）後進行水印檢測。

傳統上，水印是在內容建立之後進行的。然而，對於生成式 AI，更合適的方法是在生成過程中引入水印。演示中使用的技術是基於由 INRIA 和 Meta 的科學家開發的一個過程的改進，該過程被稱為 StableSignature。它涉及修改生成式 AI 模型（在本例中為 SDXL Turbo）的權重，使其自然地生成帶有水印的影像。因此，它在生成過程中不產生額外成本，並且更難從處理管道中移除。為了增強魯棒性，IMATAG 將此方法與其內部的零位解碼器（指示水印的存在與否，不含有效載荷）結合，該解碼器即使在水印被高度篡改後也能解碼。

此外，透過將該方法改編到 🤗 Hugging Face 的 🧨 diffusers 庫，IMATAG 使得只需 5 行程式碼即可將魯棒和不可見水印新增到生成管道中！

from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionXLPipeline

model = "stabilityai/sdxl-turbo"

vae = AutoencoderKL.from_pretrained("imatag/stable-signature-bzh-sdxl-vae-medium")

pipe = StableDiffusionXLPipeline.from_pretrained(model, vae=vae)

... 我們是這樣做的！

目錄

1. Stable Signature
2. BZH
3. 貢獻
4. 基準測試
5. 安全
6. 結論

Stable Signature

Stable Signature 的總體思想是微調 Stable Diffusion 的 VAE 解碼器，使其直接為固定金鑰生成可微分檢測器期望的特定水印訊號。其原理與針對性對抗攻擊非常相似，只不過我們希望在固定金鑰下增強水印檢測器的響應，而不是增強分類器對目標類別的響應。此外，與 SSL 水印中直接修改影像畫素不同，梯度資訊被反向傳播到 VAE 解碼器的權重中（參見 Meta 的部落格文章圖 2 (b)）。

讓我們停下來思考一下其中的含義。首先，水印被合併到 VAE 解碼器的權重中，使得其難以去除，除非原始權重公開（Stable Diffusion 不幸如此）。這也意味著水印與非水印生成相比，不產生額外的計算成本。然而，VAE 解碼器結構對水印施加了一些限制。其中之一是它會平鋪輸出影像，這意味著水印的粗略特徵不會被複制。此外，調整生成的輸出大小會改變檢測器可以依賴的補丁數量，但不會改變它們的比例，後者是固定的。這意味著總體而言，水印天生對裁剪具有魯棒性，但對縮放不具有魯棒性：原始論文表明，即使只有 10% 的原始畫素也足以高置信度地找到水印，但“調整大小和 JPEG 50 轉換似乎是最具挑戰性的，有時會低於 0.9 [位元精度]”，因為它僅依賴於水印檢測器的預訓練。另一點是，水印的可見性取決於初始模型的微調程度，可能難以控制。水印的強度由學習率和微調過程中使用的損失函式的 lambda 控制，而不是固定的 PSNR/SSIM 預算。最後，即使檢測器保持不變，每次檢測秘密金鑰都需要從頭開始進行微調過程。

BZH

IMATAG 的 BZH（盲零位元隱藏）水印系統源自 HiDDeN，與 Stable Signature 中使用的檢測器類似。然而，我們不是解碼二進位制訊息，而是依賴於零位元水印，並提取一個高維向量，我們將其與從金鑰生成的向量進行關聯。對於隨機金鑰，該期望向量均勻分佈在維度為 $dd$ 的單位超球面上。在查詢內容上檢測水印就相當於計算偶然情況下，在未加水印的影像上執行檢測器時，我們獲得比在查詢內容上觀察到的相關性 $CC$ 更高的機率 $cc$。這被稱為拒絕零假設 (H0) 的 p 值，並在給定閾值的情況下，導致錯誤地聲稱內容已加水印的機率（假陽性）。事實證明，它可以透過評估正則化不完全 beta 函式 $II$ 來分析計算，該函式表示超球面帽的面積。

$$P(C>c)=I_{1-c^2}(\frac{d-1}{2},\frac{1}{2})P(C>c)\; =\; I\_\{1-c^2\}(\backslash frac\{d-1\}\{2\},\; \backslash frac\{1\}\{2\})$$

然而，由於在此用例中金鑰是固定的，我們需要仔細考慮當我們在未加水印的隨機影像資料集上執行檢測器時，檢測器的輸出行為。事實上，正如 Stable Signature 的二進位制輸出所指出的，在這種情況下，檢測器的輸出可能具有相關性且不居中，從而無法使用上述公式計算 p 值，因為它並未均勻分佈在超球面上。然而，透過在影像資料集上學習輸出的白化變換，可以高置信度地恢復此假設。在這裡，我們使用 Flickr100k 來訓練此線性變換並執行 ZCA 白化。在 H0 下，p 值應均勻分佈。我們透過視覺化 p 值的直方圖（應為平坦的）並執行 Kolmogorov-Smirnov 檢驗來檢查我們是否可以拒絕此零假設，從而在 5000 張 AI 生成的影像上測試了此假設

與其他 HiDDeN 相比的改進包括：針對裁剪/調整大小/JPEG 以及重捕獲攻擊的再訓練、更好的聚合程式、嚴格控制水印 MSE、檢測器中支援蒙版輸入等……

貢獻

我們對 Stable Signature 過程進行了一些修改，以適應我們的需求。首先，由於我們的檢測器是零位元的，最佳化損失函式從預期訊息和預測訊息之間的二進位制交叉熵改為預期向量 $vv$ 與提取向量 $v^{\mathrm{\prime }}v\text{'}$ 之間的負餘弦相似度，兩者都已歸一化。因此，使用與論文相同的符號，

$${\mathcal{L}}_m=-v\cdot v^{\mathrm{\prime }}\backslash mathcal\{L\}\_m\; =\; -v\; \backslash cdot\; v\text{'}$$

然後，我們沒有使用感知損失 (LPIPS) 來約束解碼後的補丁與未加水印的補丁相似，而是迴歸到 KL 自動編碼器的初始訓練損失 (Stable Diffusion 論文，附錄 G，公式 (25))，該損失僅關注以合理的方式重建輸入影像

$${\mathcal{L}}_i(x)=L_{rec}(x,\mathcal{D}(\mathcal{E}(x)))-L_{adv}(\mathcal{D}(\mathcal{E}(x)))+log(D(x))+L_{reg}(x;\mathcal{E},\mathcal{D})\backslash mathcal\{L\}\_i(x)\; =\; L\_\{rec\}(x,\backslash mathcal\{D\}(\backslash mathcal\{E\}(x)))\; -\; L\_\{adv\}(\backslash mathcal\{D\}(\backslash mathcal\{E\}(x)))\; +\; log(D(x))\; +\; L\_\{reg\}(x;\backslash mathcal\{E\},\backslash mathcal\{D\})$$

我們認為這是一個更好的目標，因為未加水印的模型不應釋出，這使我們能夠更自由地生成補丁。此外，重新引入判別器項以確保解碼補丁的分佈與原始補丁的分佈難以區分，這很重要，尤其是在高失真率下。總的來說，我們希望利用解碼器的自由度來重建合理的補丁，同時包含水印，而不是適應一個本不應可用的預先存在的模型。不幸的是，在 Stable Diffusion 的 KL 自動編碼器訓練期間學習到的判別器是損失的一部分，並且（據我們所知）未釋出，因此在微調過程中必須從頭重新學習。我們透過簡單地改變最終損失的 $\lambda \backslash lambda$ 引數，微調了多個模型，以在感知度和魯棒性之間實現各種折衷

$$\mathcal{L}={\mathcal{L}}_i+\lambda {\mathcal{L}}_m\backslash mathcal\{L\}\; =\; \backslash mathcal\{L\}\_i\; +\; \backslash lambda\; \backslash mathcal\{L\}\_m$$

最後，我們預處理影像以檢測具有固定寬高比的 256 畫素邊長矩形縮放。這與 BZH 訓練設定相對應，並允許水印自然地抵抗縮小。然而，這意味著 VAE 解碼器是針對特定生成解析度訓練的，如果輸出大小改變，則必須重新訓練。我們認為這是可以接受的，因為大多數生成模型以固定大小（512px、768px、1024px）工作，並且以不同解析度生成可能會產生偽影。

基準

我們從魯棒性和感知性方面比較了該水印解決方案與其他幾種解決方案。我們使用 COCO 2017 驗證標題作為提示，使用 SDXL-turbo 生成了約 5000 張影像。作為基線，我們使用了 invisible-watermark 包中的 DCTDWT，因為它是 StableDiffusion 原始程式碼和 🤗 HuggingFace 的 🧨 diffusers 庫管道中使用的預設水印。我們還比較了生成後使用我們內部的 BZH 水印器（與我們微調的檢測器對應的那個）和 IMATAG 的生產水印（lamark）的水印。透過計算解碼器 p 值並繪製相應的 ROC 曲線來評估解決方案。假陽性率軸使用對數刻度，以更好地評估極低速率下的效能，這通常是我們感興趣的範圍。p 值閾值 1e-12（每個測試幾乎沒有錯誤）或 1e-3（文獻中常用）用淺灰色垂直線表示。StableSignature 變體用實線表示，而後水印解決方案用虛線表示。以下是我們沒有攻擊時獲得的結果

在所有水印中，DCTDWT（平均 PSNR 42.6dB）和弱模型表現最差。它們仍然比“無”水印（對應於隨機檢測）表現更好。DCTDWT 的 p 值透過計算其 48 位訊息中匹配的位數 (m) 並假設在 H0 下這些位是隨機且等機率的來計算。那麼 p 值由以下公式給出

$$P(M\ge m)=I_{0.5}(m,48-m+1)P(M\; \backslash geq\; m)\; =\; I\_\{0.5\}(m,48-m+1)$$

並且總有 2^48 中的一次機會程式碼是正確的。因此，假陽性率低於 1/2^48 的效能是未定義的，這就是為什麼灰色虛線在此值處停止的原因。但是請注意，儘管 Stable Diffusion 的 DCTDWT 金鑰的位是隨機且等機率繪製的，但 DCTDWT 的輸出並未進行白化，這與我們的模型和原始 StableSignature 論文相反。因此，上述假設實際上在實踐中並不成立，並且該系統的 p 值僅供參考。使用 DCTDWT 來判斷內容是否帶有水印應極其謹慎。BZH1、BZH2 和 BZH3 對應三種不同級別的水印強度，平均 PSNR 分別為 46.8dB（幾乎不可見）、43.8dB 和 41.9dB（略微可見）。它們都與 StableSignature 對應版本表現良好，“極端”模型非常可見，但效能與 BZH1 持平。請注意，在沒有攻擊的情況下，IMATAG 的生產水印（lamark）在平均 PSNR 45.0dB 下完美檢測，因此我們沒有在此圖中顯示它。

這似乎表明後期水印處理效果更好……但如果我們在檢測前開始更改影像呢？

上圖顯示了 StableSignature 論文中使用的“組合”攻擊後的效能：亮度增加 1.5，中心裁剪 50%，然後進行 80% 質量的 JPEG 壓縮。首先，DCTDWT 對此攻擊不具備魯棒性，表現甚至比隨機猜測還差。基於 StableSignature 的方法與後水印方法相比仍具有競爭力，BZH1 現在與弱模型持平。對於 lamark 來說，這種攻擊仍然非常容易，它以很大的優勢擊敗了所有其他方法。如果我們攻擊更強，亮度增加 1.5，縮小 2 倍，裁剪 50%，JPEG 壓縮 50%，我們會得到以下結果

現在，大多數水印都在努力抵抗這種攻擊，BZH2 和 lamark 在效能上並駕齊驅，只有 BZH3 保持了不錯的效能水平，但代價是相當顯眼。以下是影像的樣子

安全性

釋出水印器存在安全風險，因為它允許任何人給內容新增水印。在生成式 AI 的背景下，這意味著任何人都可以拿一張真實圖片併為其新增水印，假裝它是生成的。在 Stable Signature 的特殊情況下，這是透過使用帶有水印的 VAE 對影像進行編碼和解碼來完成的。然而，有人可能會爭辯說，影像已經被處理過，因此不再是真實的了。無論如何，檢測水印只是意味著影像已經通過了這個特定的帶水印的 VAE 解碼器，而任何人都可以訪問這個解碼器。此外，透過使用相同的水印器給許多影像新增水印，人們可以從中學習到一個檢測器。

釋出完整的檢測器使任何人都可以從任何內容中提取任何水印。即使水印器保密，這也可以僅讀取期望訊號並重新訓練一個新的水印器來生成它。此外，即使僅提供二進位制決策，也仍然可以訓練代理檢測器並執行對抗性攻擊。順便說一句，在對抗性術語流行之前，水印社群很久以前就將這些稱為神諭攻擊。儘管存在一些技術可以緩解此問題，但對檢測器決策的訪問（例如透過 API）應始終受到控制和限制。

此外，我們尚未討論的一個方面是秘密金鑰。僅僅因為它固定就使得該解決方案不安全。對於安全的解決方案，需要有多個金鑰，經常更改，並在之前使用的所有金鑰上執行檢測。理想情況下，金鑰應在每次使用後更改，就像一次性密碼一樣，但這會使檢測變得難以進行且魯棒性降低。此外，加金鑰的水印器和檢測器應保密。

最後，對於 Stable Signature 的特殊情況，Stable Diffusion 的原始 VAE 是公開的，這使得水印很容易受到強烈攻擊（正如論文中指出的那樣），或者提取水印並將其移動到另一張影像。

那麼，接下來我們該怎麼做呢？

首先，我們不認為我們演示中展示的水印解決方案是安全的。此外，我們知道它對高階攻擊（例如使用原始 VAE 進行擴散淨化）甚至一些簡單攻擊（例如翻轉，儘管在這種情況下可以簡單地執行兩次檢測器）或旋轉不具備魯棒性。對於更安全的解決方案、對特定攻擊的魯棒性、有效載荷支援等高階功能，您仍然應該聯絡 IMATAG。

然而，好訊息是，由於演示水印系統在固定金鑰下安全性無論如何都被破解了，也許我們可以釋出不僅僅是水印器。事實上，只要我們能檢測到演示水印並且只檢測到這個水印，我們就不會削弱系統對另一個秘密金鑰的安全性。為此，我們在水印影像上訓練了一個 Resnet-18 代理分類器，就像對抗性攻擊者在黑盒分類系統上所做的那樣。但我們實際上可以白盒訪問我們的完整檢測器，所以我們可以使用知識蒸餾來獲得更好的檢測器。由於我們不想洩露太多關於教師的資訊，所以我們不想在影像未加水印時模仿它。因此，我們最終混合使用了分類損失（在未加水印的影像上）和知識蒸餾損失（在加水印的影像上）來訓練 resnet18 學生。

這使得一個二進位制分類器能夠檢測演示水印。與完整檢測器不同，我們只有二進位制決策，沒有 p 值估計。為了重新校準檢測器，我們計算了 Flickr1M 資料集中 1M 張影像的 logits 並存儲了它們。然後透過簡單地計算這些樣本中低於影像觀測到的 logit 的 logits 比例來獲得近似的 p 值。儘管這是一個粗略的估計，但它提供了一種快速檢測大多數水印影像的方法，具有中等置信度和魯棒性。

我們很高興這使得演示水印系統在基本用例中實現了自給自足，並且比之前的公共解決方案更能抵禦無意攻擊。這是為生成式 AI 提供開放水印系統邁出的正確一步，現在我們的重點應該是提高其安全性！

結論

總而言之，DCTDWT（隱形水印）不足以抵抗即使是無意攻擊。由於訊息不具備等機率性，也很難評估其假陽性檢測率。如果希望受益於其在擴散模型生成中易於整合且無需額外計算成本，以及難以去除的特性（無法透過註釋一行程式碼來完成），則基於 StableSignature 的方法具有競爭力。將其與 BZH 結合使用，可實現魯棒且自給自足的水印解決方案。然而，在兩種情況下，安全性都受到固定（和公開）金鑰的阻礙。使用相應水印器進行後期水印處理可以更好地控制失真，更不明顯，並提供更簡單的金鑰管理。最後，IMATAG 的生產水印，適用於真實和生成的影像，如果希望最大限度地兼顧安全性和感知性/魯棒性，它仍然是最佳選擇。