What Does a Diffusion Model Remember?

tl;dr My graduate student (Tue Do) has developed DIME, an efficient membership test based on the theory of the exact optimal denoiser for diffusion models. DIME gives some impressive results: across a range of datasets, DIME consistently outperforms prior attacks at comparable or substantially lower query cost, improving certain metrics by up to 3×; remarkably, its two-query variant can outperform existing 30-query baselines. Check it out!

Diffusion models have become one of the dominant approaches to generative modeling because they can produce remarkably high-quality outputs while remaining relatively stable to train. They underpin many modern image generation and editing systems (e.g., Stable Diffusion, Midjourney, DALL-E, etc.), and related diffusion techniques have also been applied to video, audio, scientific data, inverse problems, and other modalities. Their success comes from a simple but powerful training objective: a model learns to reverse a gradual noising process, transforming corrupted data back toward samples that resemble the training distribution. This denoising formulation scales well to high-dimensional data and gives practitioners considerable control over generation, including conditioning on text, class labels, or other signals.

The widespread use of diffusion models also makes it important to understand what they learn about their training data, not just how good their generated samples look. State-of-the-art models are often trained on enormous collections of images gathered from the web or from sensitive domains such as medicine, biometrics, and personal photography. A model that behaves well at generation may nevertheless retain information about particular training examples. This raises questions about privacy, memorization, copyright, data provenance, and whether a person’s data has actually been removed after a deletion or machine-unlearning request. Membership inference, the problem of determining whether a particular record was used during training, is therefore a useful way to probe the boundary between learning general statistical structure and retaining example-specific information.

Diffusion models are especially interesting from a privacy perspective because their core denoising task is closely related to reconstruction. Given a corrupted input, the model repeatedly estimates information needed to move that input back toward likely clean data. If the model has learned unusually detailed information about a particular training example, its denoising behavior near that example may differ from its behavior near unseen data. Studying these differences helps us understand where memorization lives inside a diffusion model and how it can be detected. This is valuable both offensively, for developing stronger privacy audits, and defensively, for evaluating whether techniques such as differential privacy, deduplication, or regularization actually prevent training examples from leaving detectable traces.

There’s the diffusion model training question: Given a noisy version of an image, what noise was added, or equivalently, what clean image should the model reconstruct? And the membership inference question: Given a candidate image, was this image part of the model’s training set?

In DIME, we exploit the fact that these two questions are, in fact, tightly connected. A perfectly trained diffusion denoiser does not merely learn a generic rule for “what noise looks like.” On a finite training set, its optimal prediction can be written exactly as a posterior-weighted combination of the training examples. The model is therefore carrying out a kind of soft retrieval: given a noisy input, it asks which training points could plausibly have generated it and blends those points according to their posterior responsibilities.

That observation turns membership inference from a search for ad hoc statistics into a derivation. The model’s implicit reconstruction of a candidate has an exact squared error, and that error splits into two terms:

  1. Reconstruction bias: how far the local posterior mean is from the candidate; and
  2. Local crowding: how dispersed the plausible training points are around that mean.

The first signal resembles what earlier norm and reconstruction-based attacks measure. The second is the new ingredient: a non-member can be reconstructed surprisingly well simply because it lies near the center of several training points. Bias alone can therefore be fooled by local geometry, while crowding exposes the ambiguity.

Both signals can be estimated from forward queries to the denoising network. They share the same randomly perturbed queries, so the attack needs only q+1 model evaluations and its cheapest setting uses just two total queries. In experiments on CIFAR-10, CIFAR-100, STL10-U, CelebA, and ImageNet, DIME improves the low-false-positive detection rate substantially, sometimes beating 30-query baselines with only two queries. We also test the attack against differentially private mechanisms and finds that differentially private training drives DIME and the evaluated baselines back toward chance.


A minimal diffusion refresher

Let the finite training set be

\displaystyle \mathcal D_{\mathrm{train}}=\{x_1,\ldots,x_n\}\subseteq\mathbb R^d.

At timestep t, the forward process forms

\displaystyle g=\sqrt{\bar\alpha_t}\,x+\sigma_t z, \qquad z\sim\mathcal N(0,I_d), \qquad \sigma_t=\sqrt{1-\bar\alpha_t}.

The noise-prediction network is trained by mean squared error:

\displaystyle \mathcal L_t(\theta) = \mathbb E_{x\leftarrow\mathcal D_{\mathrm{train}},\,z\sim\mathcal N(0,I_d)} \left[ \left\| \widehat\varepsilon_\theta(\sqrt{\bar\alpha_t}x+\sigma_tz,t)-z \right\|_2^2 \right].

For squared loss, the population-optimal prediction is a conditional expectation. Here, because the data distribution is the empirical distribution over a finite training set, that conditional expectation can be calculated explicitly.


The first main result: the exact finite-set optimal denoiser

For a query g, define the unnormalized compatibility of training point x_i by

\displaystyle s_i(g) = \exp\!\left( -\frac{\|g-\sqrt{\bar\alpha_t}x_i\|_2^2}{2\sigma_t^2} \right),

and normalize these scores into responsibilities

\displaystyle r_i(g)=\frac{s_i(g)}{\sum_{j=1}^n s_j(g)}.

We prove that the MSE-optimal denoiser is

\displaystyle \boxed{ \widehat\varepsilon^*(g,t) = \frac{g}{\sigma_t} - \frac{\sqrt{\bar\alpha_t}}{\sigma_t} \sum_{i=1}^n r_i(g)x_i. }

This formula has a simple Bayesian derivation. If training point x_i generated g, the corresponding noise would be

\displaystyle z_i=\frac{g-\sqrt{\bar\alpha_t}x_i}{\sigma_t}.

Conditioned on seeing g, the posterior probability that the latent training index was i is precisely r_i(g). The optimal squared-loss estimate is therefore

\displaystyle \mathbb E[z\mid g] = \sum_i r_i(g)z_i,

which expands to the boxed expression.

The denoiser is a soft decoder over the training set

The weighted point

\displaystyle m(g)=\sum_i r_i(g)x_i

is the model’s implicit reconstruction of the clean training example that could have produced g. Training examples close to g/\sqrt{\bar\alpha_t} receive exponentially larger responsibility. The noise scale \sigma_t acts like a temperature:

  • at high noise, many records can plausibly explain the query, so responsibilities are diffuse;
  • at low noise, the posterior sharpens around the closest training record.

As \sigma_t\to0, the soft responsibilities converge almost everywhere to the indicator of the nearest training point’s Voronoi cell. In other words, the ideal denoiser approaches a hard nearest-neighbor decoder.

This explains why membership leakage is strongest in the low-noise regime. For a member x^*=x_i, the nearest point is the record itself, so the posterior can concentrate on x_i. For a non-member, the best the empirical model can do is retrieve or blend nearby training points.


Why reconstruction bias alone is not enough

The obvious attack would compare the model’s implicit reconstruction with the candidate. A member should reconstruct accurately; a non-member should not. This intuition is useful but incomplete.

A non-member can have small reconstruction bias when it lies near the mean of several training points; the dispersion of those points (the crowding term) distinguishes it from an isolated member.

We separate three cases.

Case A: isolated member

The candidate is itself a training point and has no close competitors. Responsibility concentrates on the candidate. The posterior mean is close to x^*, and the responsible neighborhood has almost no spread. Both bias and crowding are small.

Case B: isolated non-member

The closest training point or cluster lies away from the candidate. The posterior mean is displaced from x^*. Bias is large, even if the nearby training points are tightly grouped.

Case C: crowded non-member

The candidate lies near the center of several training points. Their weighted mean can accidentally be close to x^*, producing small bias. But responsibility is divided among separated records, so the local variance is large.

The one-dimensional example \mathcal D=\{-1,+1\} and x^*=0 makes the blind spot exact. If the two records have equal responsibility, the posterior mean is zero, so reconstruction bias is zero even though zero is not in the training set. The variance is one, correctly signaling ambiguity.


What the score distributions and timestep plots show

DIME score distributions from the paper.

Member and held-out DIME score distributions are visibly separated across the five checkpoints; dashed lines show calibrated thresholds.

The distribution plots make the aggregate metrics easier to interpret. DIME is not succeeding because of a handful of extreme outliers; the member and non-member score distributions shift relative to one another across every tested checkpoint. The amount of overlap explains why STL10-U permits near-total detection while CelebA remains harder.

Timestep stability from the paper.

The four smaller DDPM checkpoints remain effective across a broad timestep range, whereas ImageNet Guided Diffusion is most vulnerable at early timesteps.

The ideal denoiser theory predicts stronger concentration at lower noise. Guided Diffusion follows that prediction clearly: attack performance is highest at early timesteps and declines as noise increases. The smaller DDPM checkpoints exhibit a broad plateau instead. This suggests that the roughly 550M-parameter Guided Diffusion model may approximate the ideal denoiser more faithfully than the roughly 35M-parameter DDPM models. This is a plausible interpretation of the observed scale difference, not a formal proof that parameter count is the cause.

An additional practical lesson is that exact timestep selection may not be fragile for the smaller checkpoints. A broad plateau gives an attacker or auditor some tolerance to imperfect calibration.


Synthetic experiments isolate the role of crowding

The real-network experiments inevitably mix two phenomena:

  1. properties of the exact finite-set optimal denoiser; and
  2. approximation and optimization behavior of the trained neural network.

The synthetic appendix removes the second factor by evaluating the closed-form denoiser directly on small, known training sets.

Synthetic bias and variance decomposition from the paper.

The crowded non-member can have relatively small bias, but its variance rises much earlier and separates it from the member.

The left panel shows why a bias-only attack can fail: the crowded non-member’s posterior mean can remain close to the candidate. The middle panel shows the missing signal: responsibility spread creates a nonzero variance well before the member’s variance grows. Their sum in the right panel recovers the intended reconstruction error.

Further synthetic experiments compare DIME with SimA-MC while varying dataset size, noise level, and query count. They show that the extra crowding signal makes DIME’s membership performance decay more slowly as noise increases. These experiments are diagnostic rather than a substitute for the neural-network results, but they confirm that the claimed effect exists in the exact theoretical object itself.


Conclusion

DIME’s central insight can be stated as follows:

A finite-set diffusion denoiser implicitly performs soft retrieval over its training examples, and the error of that retrieval contains both a reconstruction signal and a geometric crowding signal.

We turn this insight into an exact formula, decompose the corresponding error, link hidden local covariance to the denoiser’s Jacobian, and estimate both terms with q+1 forward queries. The resulting attack is not merely more accurate in aggregate; its largest gains occur at the low-false-positive operating point that matters most for credible membership decisions. It also survives a dramatic jump from small unconditional DDPMs to a large class-conditional ImageNet model.

Moreover, the defenses we apply are reassuring: differential privacy suppresses not only earlier heuristic attacks but also the stronger signal produced by the ideal denoiser derivation. The broader message is therefore balanced: Diffusion models can leak membership through richer local geometry than prior attacks measured, but formal stability remains a meaningful line of defense.

DIME provides both a practical attack and a new lens for future work: privacy leakage in diffusion models can be studied as the geometry of a soft decoder over a finite codebook.


References

Tue Do and Daniel Alabi, “DIME: Query-Efficient Framework for Membership Inference on Diffusion Models”, arXiv:2608.22824, 2026.

From Proof Scarcity to Proof Abundance

tl;dr anyone who previously used or currently uses mathematical tools to solve problems in their field (whether in sciences, pure or applied mathematics, engineering, economics, and so on) needs to pay attention to discussions around the use of AI to solve mathematically precise problems. The revolution has begun!

Here’s what I believe: once AI-assisted proofs become abundant, their scarcity value will fall and attention will shift toward the harder tasks of formulating compelling conjectures, building theories, and creating definitions that organize rapidly expanding bodies of results. The analogy is economic: just as an influx of currency can produce inflation rather than lasting wealth, an influx of proofs can reduce the value attached to any individual proof without necessarily producing a comparable increase in understanding.

Firsthand Experience

Honestly, for me, the question of how useful AI could be in solving math problems did not begin in 2026. In November of 2022, ChatGPT was released. I had just begun my postdoc. With support from the Simons Foundation, I started thinking seriously in 2022 about what mathematics (and intellectual work more broadly) might look like in a world with advanced AI. At the time, much of this remained speculative: What would happen if mathematical reasoning became inexpensive? Which parts of research would retain their value? Would the bottleneck shift from proving statements to identifying the right questions, definitions, and conceptual frameworks?

At ICM (International Congress of Mathematicians) 2026, Terence Tao gave a lecture on Mathematics in the Age of AI. I’d encourage everyone to, at least, read the slides! Tao disusses the implications of the abundance of proofs (not necessarily correct). Recent events have made this transition feel concrete rather than hypothetical. While I was watching the 2026 World Cup final, Levent Alpöge announced a counterexample (credited to the AI system Fable) to the Jacobian conjecture. More precisely, the construction disproves the conjecture in dimension three, and therefore in every dimension n\geq 3, while the two-dimensional case remains open. The counterexample has a nonzero constant Jacobian determinant but maps three distinct points to the same output, directly contradicting the conjectured invertibility.

In recent days have also tested advanced AI systems on several problems from my own active research agenda. These were not merely exercises with known answers: in multiple cases, the systems made progress that appeared genuinely novel, uncovering arguments or directions that I had not previously considered. The important question is no longer whether AI can occasionally imitate mathematical reasoning. It is now clear that AI can perform at least some research-level mathematics. The remaining questions concern the breadth and reliability of that ability, the amount of human guidance and computation it requires, and, most importantly, how the mathematical community should evaluate, explain, verify, and build upon the resulting work.

Terence Tao on Mathematics in the Age of AI

In the rest of this blog post, I will highlight a few parts of his talk which I found illuminating!

Tao identifies the central challenge posed by mathematical AI: not merely generating correct proofs, but transforming those proofs into shared human understanding. That is, the most important question is no longer whether artificial intelligence can prove difficult theorems. I believe that Tao deliberately set that question aside. Instead, he asked what may be a more consequential question:

If AI systems become capable of performing a substantial fraction of research-level mathematical work, what exactly should the mathematical community be trying to accomplish?

We all need to think of that shift (from technological capability to institutional purpose). Most discussions of AI and mathematics focus on benchmark performance: Can a model solve Olympiad problems? Can it formalize a proof in Lean? Can it make progress on an open conjecture? Tao’s lecture asks what happens after the answer to some of those questions becomes yes. Tao asks the audience to provisionally assume that reasonably capable mathematical AI is coming. He then investigates what follows for proof, exposition, reviewing, education, authorship, and the organization of mathematical knowledge.

What I appreciate from the talk is that the resulting picture is neither utopian nor apocalyptic. (Although a large fraction of folks see it as a warning.) Clearly, we are headed towards an era in which proofs are scarce to one in which proofs are abundant; our present institutions are designed almost entirely for the former world, where proofs are not abundant.


A crisis of values rather than a crisis of logic

In the late nineteenth and early twentieth centuries, mathematicians were forced to examine assumptions that had previously remained largely implicit. Russell’s paradox, the emergence of competing foundational programs, and Gödel’s incompleteness theorems created a period of uncertainty about sets, numbers, infinity, axioms, and formal proof. This led to the development of more explicit and standardized foundations that gave mathematicians a trusted environment in which to work.

Tao suggests that AI is producing a comparable period of turbulence, although the foundations now at issue are not primarily logical. They are the foundations of mathematical values and practices: what counts as progress, who deserves credit, what authors are responsible for, why proofs matter, and what the profession is ultimately trying to produce. A disruptive development can expose assumptions that were previously invisible because they rarely came into conflict. Before AI, several goals of mathematics tended to advance together. A mathematician who solved a difficult problem often developed a new technique, trained students, wrote an influential paper, helped organize a research community, and contributed to a larger theory. Because those outputs were correlated, institutions could reward one visible quantity, such as the publication of a paper with one or more important theorems, and assume that the others would follow.

AI threatens to break those correlations.


The two questions that are often confused

Tao separates the AI debate into two questions.

The first is a capability question: Will AI tools, at an acceptable cost and with some level of human supervision, be able to perform meaningful research-level mathematical tasks with a nontrivial rate of success?

Tao formulates this as a family of conjectures because almost every term requires qualification. Which tools? Which fields? What counts as research-level? How much supervision? At what cost? What success rate? What standard of correctness? What standard of exposition? A system that autonomously proves 80 percent of new theorems in algebraic geometry is very different from one that occasionally completes a technical lemma after several days of expert prompting. Both might satisfy some version of an AI capability claim. Tao also stresses that public evidence is affected by reporting bias, commercial incentives, undisclosed compute budgets, selective demonstrations, and ambiguous levels of human assistance. The truth of a capability claim must also be separated from whether that capability is desirable.

The second question is a question about goals and values: What are the mathematical community’s actual objectives, including the implicit objectives encoded in publication, hiring, funding, prizes, and professional prestige? Tao’s key methodological move is to condition on a working hypothesis: suppose that a reasonably strong version of the capability claim becomes true reasonably soon. He does not ask the audience to endorse or welcome the hypothesis. He asks what the community should do if it holds. Very reasonable! This is a useful form of reasoning under uncertainty. To continue to do mathematics, we cannot wait until every capability forecast is resolved before examining whether our institutions are robust to rapid automation. Even a moderate probability of large changes can justify preparing publishing policies, authorship standards, review infrastructure, and educational norms.


First Proof

Tao briefly points to the First Proof project as one of the better-controlled pieces of evidence about current AI capabilities. For its second batch, First Proof assembled ten unpublished research-level problems originating in the work of research mathematicians. Four systems were tested in a controlled environment. The outputs were evaluated using a double-blind journal-review model involving approximately thirty experts, with each submission examined by multiple referees. Across the four systems, seven of the ten problems received at least one solution judged essentially correct or in need of only minor revisions. One system found a novel solution to a stochastic partial differential equations problem that differed from the human proof and impressed the referees. Other problems saw little or no meaningful progress.

The results are substantial, but their limitations are equally informative. The systems often handled routine portions of an argument in meticulous detail while moving too quickly through the genuinely difficult step. Referees found unsupported appeals to supposedly standard arguments, incorrect or missing citations, unnecessary notation, and lengthy explanations of facts that an expert would normally omit. Some submissions reproduced distinctive language and notation from earlier work without appropriate attribution (i.e., behavior that would raise serious plagiarism concerns in a human submission).

The computational costs also varied dramatically. In the reported runs, per-problem costs ranged from single-digit dollar amounts to nearly one thousand dollars, depending on the system and problem. This reinforces Tao’s point that “AI solved the problem” is not a complete scientific statement without information about models, harnesses, prompts, search procedures, human intervention, tokens, time, and cost.

Most importantly, First Proof evaluated only one part of mathematical research: generating proofs of questions selected and formulated by humans. Its organizers explicitly distinguish solving questions within existing frameworks from the equally important work of asking new questions and developing the frameworks in which those questions become meaningful.


Mathematics has more than one objective

What is mathematical research for?

Tao identifies several answers:

  • solving pure and applied problems;
  • building theories and techniques;
  • understanding the natural and social world;
  • training future mathematicians (part of my job);
  • sustaining a mathematical community;
  • contributing to a shared network of knowledge;
  • creating works of enduring intellectual or aesthetic value.

Historically, progress along these dimensions was positively correlated. Solving an important problem often required theory building. Theory building produced reusable techniques. Writing and teaching those techniques helped train students and build a community. This made it tempting to treat theorem production as a proxy for mathematical progress. Tao invokes Goodhart’s law: once a proxy becomes an explicit target, optimization can destroy the relationship that made the proxy useful. AI makes this danger especially severe because it can lower the cost of producing objects that look like the outputs rewarded by current metrics.

Suppose an institution rewards:

  • the number of problems solved;
  • the number of papers published;
  • being first to announce a proof;
  • benchmark scores;
  • citation counts;
  • the number of formalized theorems.

An AI-intensive workflow may increase each of these without proportionately increasing understanding, theory, pedagogy, reliability, or long-term usefulness.

A thousand isolated proofs are not necessarily more valuable than one conceptual framework explaining why all thousand theorems are true. A technically correct argument is not necessarily a useful argument. A formally verified statement is not necessarily the statement that mathematicians intended to prove. A highly cited result is not necessarily well understood.

The deeper point is that mathematical knowledge is not a warehouse of theorem statements. It is an organized system of concepts, explanations, analogies, techniques, examples, counterexamples, and judgments about what matters.


Tao’s proposed “oral defense” standard

Tao offers a rule of thumb: a result should not be published unless its human authors can convincingly present a clear, correct, properly attributed expert-level explanation of it. The proposal addresses several problems simultaneously.

An author capable of defending the result is more likely to have:

  • checked that the theorem says what it is supposed to say;
  • located the crucial step;
  • understood the relationship to prior work;
  • detected misleading AI-generated exposition;
  • assumed meaningful responsibility for correctness.

The idea resembles a thesis defense. Producing a document is not enough; the author must demonstrate command of its content. As a literal universal publication requirement, however, the rule would need careful implementation. It is bound to be controversial: Oral presentation ability varies, and a rigid talk requirement could disadvantage researchers with disabilities, non-native speakers, or authors whose strengths lie in writing rather than performance. Large collaborations may also have genuinely distributed understanding.

The underlying principle is stronger than any particular format: publication should require evidence of human intellectual ownership and accountability. That evidence might take the form of a seminar, an annotated proof map, a technical interview, an expository companion, or an independent human audit. (I’m hoping to implement the oral defense standard within my research group.)


Selected references and further reading

Terence Tao, Mathematics in the Age of AI, ICM public-lecture slides, July 24, 2026. (Teorth)

Mohammed Abouzaid et al., First Proof: Second Batch Report, 2026. (First Proof Project)

William P. Thurston, On Proof and Progress in Mathematics, Bulletin of the American Mathematical Society, 1994. (arXiv)

Tanya Klowden and Terence Tao, Mathematical Methods and Human Thought in the Age of AI, 2026. (arXiv)

Johan Commelin, Mateja Jamnik, Rodrigo Ochigame, Lenny Taelman, and Akshay Venkatesh, Shaping the Future of Mathematics in the Age of AI, 2026. (arXiv)

Jeremy Avigad, Mathematicians in the Age of AI, 2026. (arXiv)

Leiden Declaration on Artificial Intelligence and Mathematics, June 2026. (Leiden AI & Math Declaration)

Multi-Bit Watermarking

tl;dr Recently, my research group and I have been collaborating with scientists at Amazon to design, implement, and test multi-bit watermarking schemes. We will likely release our pre-print soon. But until then, in this blog post, I present some thoughts on multi-bit text watermarking.

Text watermarking can be introduced as a digital provenance tool: a model provider slightly changes generation so that later, with the right detector, the model provider can tell (preferably with overwhelmingly high confidence) whether some generated text came from its model. That is the zero-bit version: the hidden payload or message is just “watermark present.” The more interesting version is multi-bit watermarking, where the generated text carries a message: a user ID, model ID, timestamp, content-policy tag, licensing tag, or audit trail.

That extra payload is useful, but it creates a fundamental tension. The process of embedding more bits requires the output distribution to depend more strongly on the hidden message; stronger dependence makes decoding easier, but also makes the text easier to distinguish from ordinary model output. Our latest work (to be released soon!) states this tradeoff directly: higher payloads tend to increase detectability, while stronger distortion-free requirements reduce achievable rates.

This is not a brand new conceptual problem. It is the LLM version of a much older question in information hiding, data hiding, steganography, and digital watermarking: how many bits can be hidden in a host object while preserving some notion of distributional fidelity, stealth, or robustness? Moulin and O’Sullivan’s information-theoretic analysis describes information hiding as hiding information in a host data set so that it can be reliably communicated to a receiver, covering watermarking, fingerprinting, steganography, and data embedding [1]. Cachin’s information-theoretic steganography model casts the adversary’s task as a hypothesis test between innocent cover messages and stego messages, with security quantified through distributional divergence [2]. Chen and Wornell’s quantization-index modulation work studies embedding a signal, such as a digital watermark, inside a host signal to form a composite signal, with provable rate-distortion-robustness behavior [3].

The LLM twist is that the “host” is not a fixed image, audio clip, or document. It is a next-token distribution. At each generation step, the base model gives a distribution (Q_t) over the vocabulary, and the watermarker chooses a nearby distribution (Q_t^\star) from which it samples. This turns multi-bit text watermarking into a channel coding problem with a distributional constraint.

From zero-bit detection to multi-bit communication

A zero-bit watermark asks:

H_0: X_{1:n} \sim ordinary model versus H_1: X_{1:n} \sim watermarked model.

A multi-bit watermark asks a stronger question about the following relation:

M \in \{0,1\}^k \quad \longrightarrow \quad X_{1:n} \quad \longrightarrow \quad \widehat M,

where M is the hidden message, X_{1:n} is the generated text, and \widehat M is the decoded message. The receiver may know a secret key, a public codebook, and perhaps the base model. The adversary may try to detect, remove, forge, or corrupt the watermark.

Early LLM watermarking work focused largely on zero-bit detection. Kirchenbauer et al.’s watermark, for example, randomly selects “green” tokens before generation and softly promotes them during sampling, giving an efficient statistical detector [4]. SynthID-Text is another prominent watermarking system for LLM outputs, described as preserving text quality while enabling efficient detection with minimal latency overhead. [5] Multi-bit schemes move beyond detection into payload extraction; for example, multi-bit text watermarking methods have been proposed for traceability, robust extraction, and paraphrase-resilient embedding.

Our work frames this shift using an information-theoretic channel view: reliable message recovery is governed by the conditional mutual information of the induced watermark channel, and different distortion regimes create different capacity-detectability frontiers.

The basic channel model

Let \mathcal X be the vocabulary. At time t, the base LLM gives

Q_t \in \Delta(\mathcal X),

a next-token distribution conditioned on the previous context. A multi-bit watermark has:

M \in \{0,1\}^k

as the message, and a secret key or side information S_t. The encoder maps (M,S_t) into a state

Z_t = f_t(M,S_t),

then samples

X_t \sim P_{X_t|Z_t=z}.

The key design constraint is that each conditional law P_{X_t|Z_t=z} should remain close to the base law Q_t. To quantify closeness, one could use total variation distance: d_{\mathrm{TV}}(P,Q) = \frac12 \sum_{x \in \mathcal X} |P(x)-Q(x)|. A watermarked distribution is called uniformly per-key \varepsilon-distortion-free when, for every time, message, and realized key, the conditional token distribution stays within \varepsilon total variation of the base distribution.

This definition matters because total variation has an operational meaning: it is exactly the largest possible distinguishing advantage of any detector trying to decide whether a sample came from P or Q.

Total variation is maximum distinguishing advantage

Let P and Q be distributions on a finite alphabet \mathcal X. A detector is a function \phi:\mathcal X \to \{0,1\}. Its distinguishing advantage is \Pr_{X\sim P}[\phi(X)=1] - \Pr_{X\sim Q}[\phi(X)=1].

For a fixed detector,

\mathbb E_P[\phi(X)]-\mathbb E_Q[\phi(X)] = \sum_x \phi(x)(P(x)-Q(x)).

Let

A = \{x : P(x) \ge Q(x)\}.

Since 0 \le \phi(x) \le 1, \sum_x \phi(x)(P(x)-Q(x)) \le \sum_{x\in A} (P(x)-Q(x)).

But \sum_{x\in A} (P(x)-Q(x)) = \frac12 \sum_x |P(x)-Q(x)| = d_{\mathrm{TV}}(P,Q).

The upper bound is achieved by the detector \phi^\star(x)=1\{x\in A\}. Taking the absolute value allows either P or Q to be the larger distribution on the chosen set. Therefore, \sup_\phi \left| \Pr_{P}[\phi(X)=1]-\Pr_Q[\phi(X)=1]\right| = d_{\mathrm{TV}}(P,Q).

So a per-token TV budget is a bound on the best possible one-token test. In our work, we have been using the TV budget to implement watermarking schemes and test the detectability and message-recoverability of such schemes.

How this relates to information hiding, data hiding, and steganography

The terminology across communities is inconsistent, but the underlying mathematical formulation is (fairly) stable.

Information hiding is the broad umbrella. There is a host object, a hidden message, a distortion or detectability constraint, and a receiver. The goal is reliable communication through the host.

Data hiding often emphasizes embedding payload bits into media. The fidelity constraint may be perceptual: image distortion, audio quality, semantic preservation, or edit distance.

Digital watermarking often emphasizes provenance, ownership, authentication, or tracing. The hidden message may be a copyright mark, model identifier, user identifier, or policy tag. Robustness against attacks is often central.

Steganography emphasizes secrecy of the very existence of communication. The adversary’s detection problem is primary. Cachin’s model, with a passive adversary distinguishing cover from stego distributions, is especially close to modern distributional formulations of LLM watermark stealth.

Multi-bit LLM watermarking sits at the intersection. It is data hiding because it embeds a payload. It is watermarking because the payload is usually provenance or attribution metadata. It is steganography when the watermarked output must be statistically indistinguishable from ordinary model output. It is channel coding because reliable recovery is governed by mutual information and error-correcting codes.

Our work explicitly connects LLM watermarking to the information-theoretic literature on watermarking, data hiding, and steganography, noting that classical work models watermarking as communication over a constrained channel and that the LLM setting replaces perceptual host distortion with distributional constraints on the base next-token law.

The central frontier

In my opinion, a good multi-bit watermark should satisfy at least three of these:

  1. Payload: many bits per token.
  2. Reliability: low message error after generation and possible edits.
  3. Distortion-freeness/Stealth/Low-detectability: small statistical distance from ordinary model output.
  4. Utility: low degradation of fluency, factuality, reasoning, and user experience.

Core results in information theory imply that these cannot all be maximized simultaneously. Stronger distortion-free constraints reduce mutual information. More robustness requires redundancy. More redundancy lowers rate. More aggressive perturbations improve decoding but increase detectability and may degrade quality.

The right question is not “Can we make a perfect multi-bit watermark?” It is: What is the best achievable rate at a given detectability, robustness, and quality budget?

That is a capacity question!

References

[1] Pierre Moulin and Joseph A. O’Sullivan. Information-theoretic analysis of information hiding. IEEE Trans. Inf. Theory, 49(3):563–593, 2003.

[2] Christian Cachin. An information-theoretic model for steganography. Information and Computation, 192(1):41–56, 2004.

[3] B. Chen and G. W. Wornell. Quantization index modulation: a class of provably good methods for digital watermarking and information embedding. IEEE Trans. Inf. Theory, 47(4):1423–1443, September 2006.

[4] John Kirchenbauer, Jonas Geiping, Yuxin Wen, Jonathan Katz, Ian Miers, and Tom Goldstein. A watermark for large language models. In Proceedings of the 40th International Conference on Machine Learning, 2023.

[5] Sumanth Dathathri, Abigail See, Sumedh Ghaisas, Po-Sen Huang, Rob McAdam, Johannes Welbl, Vandana Bachani, Alex Kaskasoli, Robert Stanforth, Tatiana Matejovicova, Jamie Hayes, Nidhi Vyas, Majd Merey, Jonah Brown-Cohen, Rudy Bunel, Borja Balle, Ali Cemgil, Zahra Ahmed, Kitty Stacpoole, and Pushmeet Kohli. Scalable watermarking for identifying large language model outputs. Nature, 634:818–823, 2024.