Encrypted Image Watermarking Using Fully Homomorphic Encryption and Zama Concrete ML

Invisible image watermarking is a technique used to embed hidden information within digital images without visibly altering their appearance.

During the Zama Bounty Program Season 7, we challenged the developer community to develop a system that can perform invisible watermarking operations on encrypted images, with Concrete ML, Zama’s confidential AI library using Fully Homomorphic Encryption (FHE). This approach is particularly relevant in light of recent developments in Generative AI and regulatory efforts like the EU AI Act, which push for reliable digital watermarking of AI-generated content.

Indeed, FHE could be used to create a trustless service that allows standardization across all generated images, addressing the growing need for attribution and traceability in GenAI outputs.

Applications include:

• Copyright protection: Prove ownership.
• Authentication: Verify the authenticity of images based on embedded watermarks.
• Tamper detection: Identify and localize manipulations.
• Digital media tacking: Monitor distribution and usage of images across platforms.
  

With FHE, all these applications can be performed without exposing the original content to the watermarking service, ensuring privacy and compliance with emerging regulations.

A focus on the winning submission

The following two images may look exactly the same to the naked eye, but one of them contains an invisible watermark. State-of-the-art watermarking approaches commonly rely on machine learning models, and the most convenient way to deploy such models is through online services. However, there’s a critical question: How can users ensure that images uploaded to an online watermark-embedding service remain private?

In the example above, the image on the right — which contains an invisible watermark — was produced with a privacy-preserving machine learning model built using FHE by Github user Soptq.

Its model contains two parts: 

• An encoder neural network (NN) that privately embeds watermarks.
• A decoder neural network that extracts watermarks from images, if they are present. 

The watermark is a 11-bit code, and to increase its robustness, it is augmented with an error-correcting scheme, bringing the total to 16 bits. As quantization-aware-trained (QAT) models have the best accuracy-latency tradeoffs when using Concrete ML, the winning solution adapted the encoder network to use QAT with Brevitas. Typically, the decoder is applied to detect any watermark on a public image, therefore it was not converted to FHE.

Encoder overview.

The encoder model consist of two main modules:

• The watermark image creation module
• The watermark embedding module

def forward(self, image: torch.Tensor, watermark=None):
# quant input
watermark = self.quant_watermark(watermark)
watermark = self.watermark2image(watermark)

image = self.quant_image(image)
watermark = self.quant_image(watermark)
inputs = torch.cat([image, watermark], dim=1)

In the code snippet above, the watermark image is created from the watermark bits, and then concatenated with the input image. The [.c-inline-code]quant_image[.c-inline-code] member is a [.c-inline-code]QuantIdentity[.c-inline-code] Brevitas quantizer. Because the two inputs are concatenated, they must share the same quantization, as shown in this code.

The second part of the [.c-inline-code]forward[.c-inline-code] function creates an MUNIT image-translation network with skip connections. Again, quantization is carefully handled for layer outputs that are concatenated.

enc = []
x = self.pre(inputs)
for layer in self.enc:
	enc.append(x)
	x = layer(x)

enc = enc[::-1]
for i, (layer, skip) in enumerate(zip(self.dec, enc)):
	if i < self.config.num_down_levels - 1:
		x = layer(x, skip)
	else:
		skip = self.quant_dec(skip)
		inputs = self.quant_dec(inputs)
		x = layer(x, torch.cat([skip, inputs], dim=1))

Unsampling and reflection-padding.

As several types of layers such as [.c-inline-code]Upsample [.c-inline-code], [.c-inline-code]Repeat[.c-inline-code] and [.c-inline-code]Reflection Padding[.c-inline-code] were not available in Concrete ML, the bounty winner implemented them elegantly using other PyTorch operators:

class QuantUpsample(nn.Module):
    def __init__(self, scale_factor):
        super(QuantUpsample, self).__init__()
        self.s = scale_factor

    def forward(self, x):
        n, c, h, w = x.shape
        out = x.reshape(-1, c, h, 1, w, 1)
        out = torch.cat([out] * self.s, dim=-3)
        out = torch.cat([out] * self.s, dim=-1)
        out = out.reshape(-1, c, h * self.s, w * self.s)
        return out


class Reflection1xPad2d(nn.Module):
    def __init__(self):
        super(Reflection1xPad2d, self).__init__()

    def forward(self, x):
        x_pad_right = x[:, :, :, [-2]]
        x_pad = torch.cat([x, x_pad_right], dim=3)

        x_pad_bottom = x_pad[:, :, [-2], :]
        x_pad = torch.cat([x_pad, x_pad_bottom], dim=2)

        return x_pad

Model compilation and execution.

The model is compiled with optimized quantization parameters: [.c-inline-code]rounding_treshold_bits[.c-inline-code] that ensure the best tradeoff between accuracy and latency. 

quant_encoder = compile_brevitas_qat_model(
	encoder,
  (train_sub_set, secret_compile_set),
  rounding_threshold_bits={"n_bits": 7, "method": "approximate"},
  configuration=config,
  verbose=False,
  output_onnx_file="tmp.onnx",
)

In the code above, the [.c-inline-code]train_sub_set[.c-inline-code] is a representative set of images and the [.c-inline-code]secret_compile_set[.c-inline-code] is a set of 16-bit watermarks with error-codes. These two sets are used by Concrete ML to determine the cryptographic parameters of the FHE-compatible compilation model. 

Finally, the model can be applied to an encrypted image [.c-inline-code]input[.c-inline-code] with a new watermark in the [.c-inline-code]secret[.c-inline-code] variable:

encoded_input = quant_encoder.forward(
	input.numpy(), 
	np.expand_dims(secret.numpy(), 0), 
	fhe="execute"
)

The watermark is originally an 11-bit string, it is extended to 16 bits by adding error correction bits:

secret: tensor([[1., 0., 0., 0., 1., 0., 1., 0., 0., 1., 1.]])
secret+ECC: tensor([[1., 0., 0., 1., 0., 0., 0., 0., 0., 1., 0., 1., 0., 0., 1., 1.]])

A note on performance: on a desktop CPU, the full execution process takes around 20 minutes, while on a more powerful server, it runs several times faster.

Watermark extraction.

Finally, to extract the watermark from an image, a second neural network is applied. In a typical use case, the image rights owner would check if an image they find online was a copy of the one they own the rights to. 

noised_decoded_secret = decoder(noised_input) > 0.5
noised_secret_bits = tensor2bitarray(noised_decoded_secret[0])
noised_secret = hamming_decode(noised_secret_bits)
noised_secret = bitarray2tensor(noised_secret).unsqueeze(0).float()
print("Original secret", original_secret)
print("Decoded secret", noised_secret)
print("Secret recovered? ", torch.all(noised_secret == original_secret).item())

# Original secret tensor([[1., 0., 0., 0., 1., 0., 1., 0., 0., 1., 1.]])
# Decoded secret tensor([[1., 0., 0., 0., 1., 0., 1., 0., 0., 1., 1.]])
# Secret recovered?  True

‍Conclusion

Soptq, author of the winning submission, successfully implemented a state-of-the-art neural network for embedding watermarks. This approach is highly resilient to various image transformations, including blurring, noise, resampling, cropping, and JPEG compression. The robustness is achieved through a unique training process where transformations are randomly applied, and hard negatives are mined to focus the model on the most challenging scenarios.

The runner-up solution leveraged a DCT-based decomposition approach. While it offers faster performance compared to the winning solution, its resistance to image transformations is lower. You can explore the second-place solution on this Hugging Face space.

Looking ahead, private watermarking has the potential to prove image ownership and authenticity. It can also play a crucial role in detecting image tampering—a growing concern in the era of generative AI, where misinformation and privacy risks are more prevalent than ever.

For Season 8 of the Zama Bounty Program, we invite the community to explore how FHE can enhance privacy in biological age estimation using machine learning models.

Additional links

• Star Zama's Concrete ML GitHub repository to endorse our work.
• Review the Concrete ML documentation.
• Get support on our community channels.
• Participate in the Zama Bounty Program to get rewards in cash!