使用 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 的生产水印，适用于真实和生成的图像，如果希望最大限度地兼顾安全性和感知性/鲁棒性，它仍然是最佳选择。