# TokenTrace: Multi-Concept Attribution through Watermarked Token Recovery

arXiv: https://arxiv.org/abs/2602.19019 · HTML: https://mlsanity.com/paper/2602.19019/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Watermarking generated content and models, one of 12 directions in LLM safety and security.
- Builds on 39 of its 62 references that are in the map.
- First posted on arXiv 22 Feb 2026.

## Placed in

- LLM safety and security / [Watermarking generated content and models](https://mlsanity.com/safety/watermarking-generated-content-and-models/)

## Builds on (most cited first, 30 of 39)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding](https://mlsanity.com/paper/2205.11487/) (2022)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation](https://mlsanity.com/paper/2208.12242/) (2022)
- [An Image is Worth One Word: Personalizing Text-to-Image Generation using Textual Inversion](https://mlsanity.com/paper/2208.01618/) (2022)
- [Visual Prompt Tuning](https://mlsanity.com/paper/2203.12119/) (2022)
- [Conditional Prompt Learning for Vision-Language Models](https://mlsanity.com/paper/2203.05557/) (2022)
- [Multi-Concept Customization of Text-to-Image Diffusion](https://mlsanity.com/paper/2212.04488/) (2022)
- [ELITE: Encoding Visual Concepts into Textual Embeddings for Customized Text-to-Image Generation](https://mlsanity.com/paper/2302.13848/) (2023)
- [AnyDoor: Zero-shot Object-level Image Customization](https://mlsanity.com/paper/2307.09481/) (2023)
- [HiDDeN: Hiding Data With Deep Networks](https://mlsanity.com/paper/1807.09937/) (2018)
- [PhotoMaker: Customizing Realistic Human Photos via Stacked ID Embedding](https://mlsanity.com/paper/2312.04461/) (2023)
- [The Stable Signature: Rooting Watermarks in Latent Diffusion Models](https://mlsanity.com/paper/2303.15435/) (2023)
- [Forget-Me-Not: Learning to Forget in Text-to-Image Diffusion Models](https://mlsanity.com/paper/2303.17591/) (2023)
- [Diffusion Model-Based Image Editing: A Survey](https://mlsanity.com/paper/2402.17525/) (2024)
- [A Self-Supervised Descriptor for Image Copy Detection](https://mlsanity.com/paper/2202.10261/) (2022)
- [Gaussian Shading: Provable Performance-Lossless Image Watermarking for Diffusion Models](https://mlsanity.com/paper/2404.04956/) (2024)
- [Invisible Image Watermarks Are Provably Removable Using Generative AI](https://mlsanity.com/paper/2306.01953/) (2023)
- [EasyControl: Adding Efficient and Flexible Control for Diffusion Transformer](https://mlsanity.com/paper/2503.07027/) (2025)
- [Measuring Style Similarity in Diffusion Models](https://mlsanity.com/paper/2404.01292/) (2024)
- [Watermark Anything with Localized Messages](https://mlsanity.com/paper/2411.07231/) (2024)
- [Black-Box Forgery Attacks on Semantic Watermarks for Diffusion Models](https://mlsanity.com/paper/2412.03283/) (2024)
- [FakeTagger: Robust Safeguards against DeepFake Dissemination via Provenance Tracking](https://mlsanity.com/paper/2009.09869/) (2020)
- [ROBIN: Robust and Invisible Watermarks for Diffusion Models with Adversarial Optimization](https://mlsanity.com/paper/2411.03862/) (2024)
- [Evaluating Data Attribution for Text-to-Image Models](https://mlsanity.com/paper/2306.09345/) (2023)
- [DesignDiffusion: High-Quality Text-to-Design Image Generation with Diffusion Models](https://mlsanity.com/paper/2503.01645/) (2025)
- [Proactive Image Manipulation Detection](https://mlsanity.com/paper/2203.15880/) (2022)

## Abstract (arXiv)

Generative AI models pose a significant challenge to intellectual property (IP), as they can replicate unique artistic styles and concepts without attribution. While watermarking offers a potential solution, existing methods often fail in complex scenarios where multiple concepts (e.g., an object and an artistic style) are composed within a single image. These methods struggle to disentangle and attribute each concept individually. In this work, we introduce TokenTrace, a novel proactive watermarking framework for robust, multi-concept attribution. Our method embeds secret signatures into the semantic domain by simultaneously perturbing the text prompt embedding and the initial latent noise that guide the diffusion model's generation process. For retrieval, we propose a query-based TokenTrace module that takes the generated image and a textual query specifying which concepts need to be retrieved (e.g., a specific object or style) as inputs. This query-based mechanism allows the module to disentangle and independently verify the presence of multiple concepts from a single generated image. Extensive experiments show that our method achieves state-of-the-art performance on both single-concept (object and style) and multi-concept attribution tasks, significantly outperforming existing baselines while maintaining high visual quality and robustness to common transformations.