🤖 machine learning

Traceable Black-box Watermarks for Federated Learning

This paper proposes **TraMark**, a novel server-side watermarking method for Federated Learning that enables traceable black-box verification of model leakage by partitioning the model parameter space and assigning unique, client-specific watermarks to each local model without compromising main task performance.

Original authors: Jiahao Xu, Rui Hu, Olivera Kotevska, Zikai Zhang

Published 2026-02-10

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Jiahao Xu, Rui Hu, Olivera Kotevska, Zikai Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are part of a secret cooking club. Instead of sharing your secret recipes (your private data), everyone works on a single "Master Cookbook" (the global model). Every week, you each write down a few small improvements to the recipes and send them to a Head Chef (the server). The Head Chef combines everyone's ideas into a new version of the Master Cookbook and sends it back to you.

The Problem: The Recipe Thief
The problem is that once you get a copy of the Master Cookbook, you might be tempted to photocopy it and sell it to a rival restaurant. If that happens, the Head Chef will know someone stole the book, but they won't know which club member was the thief.

Current security methods are like putting a tiny, invisible ink signature on every page. But to check if a stolen book is yours, you have to tear the pages apart and look at the ink under a microscope (this is "white-box" access). In the real world, if a thief is selling the book, you can't just grab it and put it under a microscope—you can only look at the recipes themselves (this is "black-box" access).

The Solution: TraMark (The "Secret Ingredient" Method)
The researchers created a system called TraMark. Here is how it works using a simple analogy:

1. The Divided Kitchen (Partitioning the Model)
Instead of putting a signature on every single page, the Head Chef decides that 99% of the cookbook is for the actual recipes (the Main Task Region). The remaining 1% is a special "Secret Ingredient Section" (the Watermarking Region).

2. Personalized Spice Blends (Masked Aggregation)
When the Head Chef sends the cookbook back to you, they don't send the exact same book to everyone.

For the recipes, everyone gets the same great, averaged-out instructions.
But for the Secret Ingredient Section, the Chef gives you a unique blend of spices that no one else has.

If you are Member A, your secret ingredient might be a pinch of saffron. If you are Member B, it’s a dash of cardamom. Because these spices are tucked away in that tiny 1% section, they don't ruin the taste of the actual recipes.

3. The Taste Test (Black-Box Verification)
Now, imagine a rival restaurant is suddenly selling a cookbook that looks suspiciously like yours. To catch the thief, the Head Chef doesn't need to tear the book apart. They just perform a "Taste Test."

The Chef asks: "If I add a specific amount of saffron to this dish, does it taste exactly like Member A's special version?"

If the dish tastes like saffron $\rightarrow$ Member A leaked it!
If it tastes like cardamom $\rightarrow$ Member B leaked it!

Because the "spices" (the watermarks) are unique to each person, the Chef can identify the thief just by looking at the final result, without ever needing to see the internal "ink" or "notes" inside the book.

Summary of the "Magic"

It’s Traceable: Every member has a unique "flavor" so you can point fingers.
It’s Black-Box: You can catch the thief just by "tasting" the output, not by dissecting the model.
It’s High Quality: Because the "spices" are kept in a tiny, separate section, the actual recipes (the AI's performance) stay delicious and accurate.
It’s Tough to Remove: If a thief tries to "wash out" the spices (an attack), the flavors are so deeply tied to the structure of that section that it’s very hard to get rid of them without ruining the whole recipe.

Technical Summary: Traceable Black-Box Watermarks for Federated Learning (TraMark)

Published at: ICLR 2026

1. Problem Statement

Federated Learning (FL) systems face a significant risk of model leakage, where malicious clients may illegally distribute the global model for personal profit. While watermarking is a known defense for Intellectual Property (IP) protection, existing methods suffer from two major limitations in FL:

Lack of Black-Box Traceability: Many existing watermarking methods (e.g., parameter-based) require "white-box" access (access to model weights) for verification. In real-world scenarios, a leaked model might only be accessible via an API (black-box setting).
The Traceability Gap: Even if a model is verified as stolen, current black-box methods often fail to identify which specific client leaked it. Conversely, methods that attempt traceability often require access to clients' private local data or modify the local training protocol, making them vulnerable to tampering or privacy violations.

The authors formalize the Traceable Black-box Watermarking problem: how to inject unique, verifiable, and non-colliding watermarks into models shared with clients, such that the server can identify the leaker using only model outputs (black-box) without accessing private client data.

2. Methodology: TraMark

The authors propose TraMark, a novel server-side watermarking framework that integrates seamlessly into standard FL (e.g., FedAvg). The core innovation lies in parameter space partitioning and masked aggregation.

Parameter Space Partitioning: The server divides the model parameters into two distinct, non-overlapping regions using binary masks:
1. Main Task Region ( $M_m$ ): Responsible for the primary learning task.
2. Watermarking Region ( $M_w$ ): Reserved exclusively for embedding the watermark.
Masked Aggregation: To prevent "watermark collision" (where watermarks from different clients fuse and become indistinguishable during the next round of aggregation), TraMark uses a specialized aggregation rule. It averages updates only within the Main Task Region but preserves the unique, client-specific updates in the Watermarking Region. This ensures each client receives a personalized global model.
Server-Side Watermark Injection: The server performs the watermarking process itself. For each client $i$ , the server takes the personalized model and performs a few iterations of training using a unique watermark dataset ( $D^w_i$ ) that only affects the watermarking region. This prevents the watermark from degrading the main task performance.
Distinct Watermark Datasets: To ensure traceability, each client is assigned a unique trigger set and a unique predefined output distribution. The authors suggest using Out-of-Distribution (OOD) samples (e.g., using MNIST digits as triggers for a CIFAR-trained model) to ensure the watermark is independent of the main task.
Warmup Phase: To ensure the partition is effective, the model undergoes a "warmup" period of standard FL training. This allows the server to identify "unimportant" parameters (those with small magnitudes) to assign to the watermarking region, minimizing impact on accuracy.

3. Key Contributions

Formalization: The first work to formally define the problem of traceable black-box watermarking in FL.
Novel Mechanism: Introduction of masked aggregation and parameter partitioning to prevent watermark entanglement and collision.
Server-Side Implementation: A method that requires zero access to client data and zero modification to the client's local training protocol, making it highly practical and resistant to client tampering.
Robustness: A design that ensures watermarks are coupled with the model, making them resistant to common attacks like pruning and fine-tuning.

4. Experimental Results

The authors evaluated TraMark across various datasets (FMNIST, CIFAR-10/100, Tiny-ImageNet) and settings (IID and non-IID).

High Traceability: TraMark achieved an average Verification Rate (VR) of 99.58%, meaning it successfully identified the correct leaker in nearly all cases.
Minimal Utility Loss: The impact on the main task was negligible, with an average accuracy drop of only 0.54% compared to standard FedAvg.
Robustness to Attacks:
- Pruning/Fine-tuning: The watermarks remained effective even when clients attempted to remove them via magnitude-based pruning or local fine-tuning.
- Quantization: The method proved resilient to INT8 and FP16 quantization.
- Adaptive Attacks: It showed resistance to Reconstructive Neuron Pruning (RNP), a state-of-the-art backdoor removal technique.
Scalability: Experiments on Tiny-ImageNet with 100 clients showed that the method scales effectively without significant computational or communication overhead.

5. Significance

TraMark provides a robust solution for accountability in Federated Learning. By enabling the server to act as a verifier that can pinpoint leakers through black-box queries, it protects the intellectual property of FL developers. Its ability to operate without client cooperation or data access makes it a highly viable candidate for real-world, privacy-preserving, and secure distributed learning systems.