TokenTrace: Multi-Concept Attribution through Watermarked Token Recovery

TokenTrace is a novel proactive watermarking framework that embeds signatures into text prompt embeddings and initial latent noise to enable robust, query-based disentanglement and attribution of multiple concepts within a single generative AI image, outperforming existing methods in both accuracy and visual quality.

The Gist

Imagine you are a famous artist who has spent years developing a unique painting style. Suddenly, a powerful AI starts painting pictures that look exactly like your work, but it doesn't give you credit. Worse yet, the AI might mix your style with a picture of a specific dog you trained it on, creating a "hybrid" image. Current tools can't tell who owns the style and who owns the dog; they just see one big, messy picture.

This is the problem TokenTrace solves. It's like a high-tech, invisible detective system for AI-generated art.

Here is how it works, broken down into simple analogies:

1. The Problem: The "Blended Smoothie"

Think of an AI image as a smoothie made from different ingredients.

• Ingredient A: A specific object (like a "golden retriever").
• Ingredient B: A specific style (like "Van Gogh's brushstrokes").

Old watermarking methods were like putting a single, giant sticker on the outside of the smoothie cup. If you wanted to know if the "dog" or the "Van Gogh style" was inside, the sticker couldn't tell you. It just said, "This cup has a sticker." If you tried to peel the sticker off to check one ingredient, the whole cup might get ruined.

2. The Solution: Invisible DNA in the Recipe

TokenTrace changes the game. Instead of putting a sticker on the outside (the final picture), it secretly rewrites the recipe (the instructions the AI follows).

Imagine you are baking a cake.

• The Secret: You want to prove you baked a "Chocolate Cake" and a "Red Velvet Cake" in the same batch.
• The Old Way: You write "Made by Bob" on the frosting.
• The TokenTrace Way: You secretly change the flour to have a tiny bit of "Chocolate DNA" and the sugar to have a tiny bit of "Red Velvet DNA."

When the cake is baked, the DNA is baked into the cake itself. You can't see it, and it doesn't change the taste (the picture looks perfect). But because the DNA is in the ingredients, it survives even if you slice the cake, freeze it, or wrap it in foil.

3. How It Works: The "Dual-Conditioning" Strategy

TokenTrace does this by tweaking two things at the same time before the AI starts drawing:

1. The Text Prompt (The Recipe Card): It slightly alters the words you type (e.g., "a dog") to hide a secret code.
2. The Initial Noise (The Raw Dough): It slightly changes the random static the AI starts with to hide another part of the code.

By hiding the secret in both the words and the starting noise, the watermark becomes deeply woven into the fabric of the image. It's not just sitting on top; it's part of the image's DNA.

4. The Detective: The "Query-Based" Module

This is the coolest part. When you want to check who owns the image, you don't just scan the whole picture blindly. You use a Query.

Imagine you have a locked box with a mix of gold and silver coins.

• Old Method: You shake the box and hope to hear a "gold" sound.
• TokenTrace Method: You hold up a magnet specifically for gold and ask, "Show me the gold!" The system then uses a special lens (the TokenTrace module) to look at the image and say, "Yes, there is a gold coin here!" Then you hold up a magnet for silver and it says, "Yes, there is a silver coin here too!"

In technical terms:

• You feed the image into the system.
• You type a question: "Is the dog in this picture?"
• The system checks the "Dog DNA" hidden in the recipe and confirms it.
• You then ask: "Is the Van Gogh style in this picture?"
• The system checks the "Style DNA" and confirms that too.

It can separate the two concepts perfectly, even though they are mixed together in one image.

5. Why This Matters

• For Artists: It protects your unique style and your specific creations. Even if an AI mixes your style with a celebrity's face, you can prove, "That style is mine," and "That face is theirs."
• For Everyone: It stops AI from stealing credit. It ensures that when we see a beautiful image, we know exactly who (or what) contributed to its creation.
• Robustness: Because the secret is baked into the "ingredients" (the text and noise), it survives common tricks like cropping the image, compressing it for the web, or even trying to hack it with AI attacks.

Summary

TokenTrace is like giving every AI-generated image a set of invisible, unbreakable barcodes hidden inside its very DNA. Unlike old methods that put a sticker on the surface, TokenTrace writes the ownership story into the recipe itself. And the best part? You can ask the image specific questions ("Who made the dog?" "Who made the style?") and get a clear, separate answer for each, ensuring everyone gets the credit they deserve.

Technical Summary

1. Problem Statement

Generative AI models, particularly text-to-image diffusion models, pose a significant threat to intellectual property (IP) by replicating unique artistic styles and concepts without attribution. While proactive watermarking (embedding signatures during generation) offers a solution, existing methods face critical limitations in multi-concept scenarios:

• Pixel-Domain Methods: Often fragile and easily broken by common transformations (compression, cropping).
• Latent-Domain Methods: More robust but typically embed a single, holistic watermark. They fail to disentangle multiple overlapping concepts (e.g., a specific object rendered in a specific artistic style) within a single image.
• The Gap: Current systems cannot independently verify the presence of multiple distinct concepts in a generated image, leaving a major gap in protecting creator IP for compositional generations.

2. Methodology: TokenTrace

TokenTrace is a proactive watermarking framework designed for robust, multi-concept attribution. It introduces a dual-conditioning strategy that embeds secret signatures into both the semantic domain (text prompts) and the latent domain (initial noise), rather than just the pixel space.

A. Concept Encoding (Generation Phase)

The framework uses two parallel networks to process a concept secret ($S$) before the diffusion model generates the image:

1. Concept Encoder ($f_{enc}$): Perturbs the text prompt embedding. It takes the concept secret and the specific target token embedding (e.g., the token for "style" or "object") to generate a perturbation. This is fused element-wise with the target token embedding to create a perturbed prompt ($\hat{E}_{prompt}$).
2. Secret Mapper ($f_{map}$): Perturbs the initial latent noise. It generates a structured Gaussian noise pattern based on the secret, which is added to the original noise ($z_T$) to create perturbed noise ($\hat{z}_T$).

The diffusion model is then conditioned on both $\hat{E}_{prompt}$ and $\hat{z}_T$ to generate the watermarked image ($I_{wm}$). This dual embedding ensures the signature is deeply woven into the image structure, providing robustness against transformations.

B. Concept Decoding (Retrieval Phase)

To retrieve the watermark, TokenTrace employs a query-based mechanism:

• Input: The watermarked image ($I_{wm}$) and a specific textual query ($P_{query}$) specifying which concept to retrieve (e.g., "a photo of ").
• TokenTrace Module: A parameter-efficient module (using frozen CLIP encoders) that fuses image features and text query features via trainable projection and attention layers. It predicts the concept embedding ($\tilde{e}_c$) corresponding to the query.
• Secret Decoder: A lightweight linear network that translates the predicted embedding back into the original bit-secret ($\tilde{S}$).

This query-based approach allows the system to selectively isolate and verify specific concepts from a single image, solving the multi-concept disentanglement problem.

C. Training Objective

The model is trained end-to-end using a composite loss function ($L_{total}$) balancing four constraints:

1. Cross-Entropy Loss ($L_{CE}$): Ensures the retrieved secret matches the ground truth.
2. Contrastive Style Descriptor Loss ($L_{CSD}$): Maintains high-level semantic and style consistency between clean and watermarked images.
3. L2 Loss ($L_{L2}$): Minimizes pixel-level differences to ensure imperceptibility.
4. Regularization Loss ($L_{reg}$): Ensures the predicted concept embedding is close to the ground-truth embedding.

3. Key Contributions

1. Dual-Domain Embedding: A novel framework that ties watermarks to concept semantics by perturbing both text prompts and latent noise, avoiding the spatial overlap issues of pixel-based methods.
2. Query-Based Disentanglement: A retrieval module that uses textual queries to independently verify multiple overlapping concepts (objects and styles) within a single image.
3. State-of-the-Art Performance: Demonstrated superior performance in both single-concept and complex multi-concept attribution tasks while maintaining high visual fidelity.

4. Experimental Results

The authors evaluated TokenTrace on WikiArt (styles), ImageNet (objects), and custom multi-concept benchmarks.

• Single-Concept Attribution:
  • Style (WikiArt): Achieved 91.67% attribution accuracy and 98.33% bit accuracy, outperforming ProMark (87.19%) and CustomMark (89.25%).
  • Object (ImageNet): Achieved 90.43% attribution accuracy and 95.82% bit accuracy on 1,000 object classes.
• Multi-Concept Attribution:
  • Successfully disentangled and attributed multiple concepts (e.g., object + style) from a single image.
  • In a 4-concept scenario, TokenTrace achieved 81.57% attribution accuracy, significantly outperforming CustomMark (72.78%).
  • The enhanced TokenTraceP variant (using prompt weighting) pushed accuracy to 86.08%.
• Robustness: The method maintained high accuracy (>82%) against severe transformations including JPEG compression, rotation, cropping, Gaussian noise, and adversarial attacks.
• Visual Fidelity: The watermarked images were visually indistinguishable from clean originals, with FID scores (14.98) comparable to CustomMark and significantly better than pixel-based ProMark (17.63).
• Scalability: The system showed graceful degradation as the concept library scaled from 10 to 1,000 concepts and supported sequential learning with minimal retraining.

5. Significance

TokenTrace addresses a critical bottleneck in generative AI attribution: the inability to trace specific components within a composite generation. By moving the watermarking process into the semantic and latent domains and utilizing a query-based retrieval mechanism, it provides a robust, scalable, and IP-protective solution for the creative industry. This work enables the verification of ownership for complex, multi-concept AI generations, ensuring that artists and creators can trace their unique styles and objects even when they are combined in new ways.