How Watermarking Text Works | Typebulb

Typebulb requires JavaScript to run bulbs.

How Watermarking Text Works<br>Learn how text watermarking works in your browser: a secret key runs a tournament between words the model drew, so the winner is one it would say.<br>format: typebulb/v1<br>name: How Watermarking Text Works

**code.tsx**

```tsx<br>import {<br>Module, Linear, LayerNorm, compileForward, checkWebGPU,<br>add, mul, matmul, sum, reshape, swapAxes,<br>splitHeads, mergeHeads, softmaxCausal, gelu,<br>type Tensor,<br>} from 'tensorgrad'<br>import {<br>App, Component, a, div, h1, h2, h3, p, span, strong, em, button, inputRange, inputTextArea,<br>table, thead, tbody, tr, th, td,<br>type VElement,<br>} from 'domeleon'

// ============================================================================<br>// The language model<br>// ============================================================================

// TinyStories-1M (roneneldan/TinyStories-1M), a GPT-Neo decoder trained only on<br>// synthetic three-year-old-vocabulary stories. 3.75M parameters, of which the 50257-row<br>// embedding table — also the tied output head — is 3.2M and the entire eight-layer<br>// transformer stack only 0.4M. The copy shipped here is 3.63M: the exporter keeps the<br>// first 256 rows of the 2048-row position table, which is all a 256-token window reaches.<br>const D = 64, L = 8, HEADS = 16, VOCAB = 50257

// The compiled graph is one fixed length. 256 is also GPT-Neo's local-attention window:<br>// this checkpoint alternates global and local(256) attention layers, and below 256 tokens<br>// the two are the same function, so one causal attention is exact. Above it they diverge.<br>const CTX = 256

const INPUTS = { embed: [1, CTX, D], sel: [1, CTX] } as const

class Block extends Module {<br>ln1 = new LayerNorm(D)<br>q = new Linear(D, D, { bias: false })<br>k = new Linear(D, D, { bias: false })<br>v = new Linear(D, D, { bias: false })<br>attnOut = new Linear(D, D)<br>ln2 = new LayerNorm(D)<br>fc = new Linear(D, 4 * D)<br>proj = new Linear(4 * D, D)

class TinyStories extends Module {<br>// Present as a parameter for the output head only. The input-side lookup happens on the<br>// CPU: tensorgrad's `embedding` composes to `oneHot @ table`, and a [1,256,50257] one-hot<br>// is 51MB of tensor per step to express what is really a gather of 256 rows.<br>wte = this.param([VOCAB, D])<br>blocks: Block[]<br>lnf = new LayerNorm(D)<br>constructor() {<br>super()<br>this.blocks = Array.from({ length: L }, () => new Block())

// GPT-Neo is PRE-norm, and its attention does NOT divide qk by sqrt(headDim). That missing<br>// scale is a Mesh-Tensorflow inheritance and the single easiest way to get this port wrong:<br>// adding it still yields fluent stories, drawn from a distribution ~45% off the real one.<br>// scripts/pack-tinystories.mjs checks the unscaled form against HuggingFace's own logits on<br>// every build — though against its own CPU reference pass, so what that proves is the<br>// weights and the arithmetic, not this graph: the two are separate implementations that<br>// share the checkpoint and a 1e-5 layernorm epsilon.<br>function block(b: Block, h: Tensor): Tensor {<br>const a = b.ln1.fwd(h)<br>const q = splitHeads(b.q.fwd(a), HEADS)<br>const k = splitHeads(b.k.fwd(a), HEADS)<br>const v = splitHeads(b.v.fwd(a), HEADS)<br>const ctx = mergeHeads(matmul(softmaxCausal(matmul(q, swapAxes(k, -1, -2)), -1), v))<br>const h2 = add(h, b.attnOut.fwd(ctx))<br>const m = b.ln2.fwd(h2)<br>return add(h2, b.proj.fwd(gelu(b.fc.fwd(m), { approximate: 'tanh' })))

// `sel` is one-hot over positions: it picks the last real token's row out of the padded<br>// window. Padding needs no mask of its own — attention is causal, so a real position can<br>// never see a later pad.<br>function forward(m: TinyStories, { embed, sel }: { embed: Tensor; sel: Tensor }): Tensor {<br>let h = embed<br>for (const b of m.blocks) h = block(b, h)<br>h = m.lnf.fwd(h)<br>const last = sum(mul(h, reshape(sel, [1, CTX, 1])), 1)<br>return matmul(last, swapAxes(m.wte, -1, -2))

// ============================================================================<br>// The watermark: tournament sampling<br>// ============================================================================

// Dathathri et al., "Scalable watermarking for identifying large language model outputs",<br>// Nature 634 (2024) — the SynthID-Text scheme, deployed in Gemini and open-sourced in<br>// HuggingFace Transformers as SynthIDTextWatermarkLogitsProcessor.<br>//<br>// The construction in one paragraph. At each step the model produces its distribution over<br>// the vocabulary. Draw N^m candidate words from THAT distribution, independently and with<br>// replacement, and run them through m knockout layers: in each layer the survivors are split<br>// into matches of N, and the winner of a match is the candidate whose coin came up heads,<br>// ties broken uniformly at random. The coins come from the key. The last word standing is<br>// emitted.<br>//<br>// What that buys, and it is the whole reassurance: every candidate in the bracket was drawn<br>// from the model's own distribution before the key saw it, so the key can never promote...