Surprise, cross-entropy, and perplexity

Top-k and top-p fix a rule about the shape of the truncation and let the consequence — how surprising the text turns out — float freely, and that consequence is not stable across a long generation. Drift toward low perplexity is the boredom trap: the sampler keeps picking the safest continuation and the text loops. Drift toward high perplexity is the confusion trap: too much tail survives and the text wanders into incoherence. Mirostat controls perplexity directly to block both.

The unit it controls is surprise. The surprise of a token x with model probability P(x) is S(x) = -log2 P(x) bits: a token rated 1/2 costs 1 bit, one rated 1/32 costs 5. Average surprise over a sequence is the cross-entropy H = -(1/T) Σ_t log2 P(x_t), and perplexity is that exponentiated, PPL = 2^H — the effective branching factor. Mirostat’s whole game is to pin the running average surprise to a chosen τ, which pins perplexity to 2^τ: τ = 3 asks for perplexity 8, the llama.cpp default τ = 5 for 32.

Advertisement

The Zipfian model of the next-token distribution

To connect ‘how many tokens I keep’ to ‘expected surprise’ you need a model of the distribution’s shape. Mirostat uses the empirical regularity that the sorted next-token probabilities follow a Zipf power law: the token at rank i has probability roughly p_i ∝ i^(-s) for some exponent s > 1.

This is the same law that governs word frequencies in natural language, and it holds well enough over the head of a softmax distribution to be useful. Its payoff is analytic tractability: under a Zipf law with exponent s, truncating to the top k tokens and renormalizing gives an expected surprise you can write in closed form as a function of s, k, and vocabulary size N. That closed form lets Mirostat invert the relationship: given a target surprise τ, solve for the k that delivers it — no trial and error inside the step.

Advertisement

Estimating the exponent s from the logits

The exponent s is not fixed — it shifts every step with the context — so Mirostat re-estimates it from the current distribution. In log–log form the Zipf law is log p_i = c - s·log i, a straight line of slope -s, so estimating s is a line fit. Define t_i = ln(p_i / p_{i+1}) and b_i = ln((i+1)/i) over the top handful of tokens (the paper uses about 100). Under a perfect Zipf law t_i = s·b_i, so a least-squares slope through the origin gives:

s_hat = ( Σ_i t_i · b_i ) / ( Σ_i b_i^2 )

This is cheap — a handful of logs and a dot product — and it adapts to whatever the network just produced, sharp or flat. It is noisy per step, but need not be perfect: the outer loop corrects the residual.

Solving for the truncation k

With s_hat in hand and a running budget μ (the current surprise target in bits, of which more below), Mirostat computes the truncation size directly. Writing ε = s_hat - 1, the closed-form choice is:

k = ( ε · 2^μ  /  (1 - N^(-ε)) ) ^ (1 / s_hat)

then rounded to an integer. The intuition: a larger surprise budget μ buys a larger k, keeping more tail so the average token is more surprising; a steeper distribution (larger s_hat) reaches the same surprise with fewer tokens, so the 1/s_hat exponent pulls k down. The term 1 - N^(-ε) corrects for the finite tail. Mirostat then samples from those top k renormalized candidates — an ordinary top-k draw, but with a k that was computed, not fixed.

The feedback loop: the mu update rule

Estimating s and solving for k would, on its own, only hit the target in expectation — the Zipf fit is approximate and each draw is random. What makes Mirostat a genuine controller is the update to μ. Initialize μ = 2τ; after sampling a token x, measure its surprise S(x) = -log2 P(x), form the error, and correct:

error   e = S(x) - τ
μ ← μ - η · e

This is a proportional feedback controller, the same idea as gradient descent on the error. If the sampled token was more surprising than τ (e > 0), μ drops, shrinking the next k and tightening the sampler; if it was too safe (e < 0), μ rises and the next k widens. The learning rate η (default 0.1) sets how aggressively the loop chases its target: too high and μ oscillates, too low and it reacts sluggishly.

A worked step

Take τ = 3 bits (target perplexity 8), η = 0.1, vocabulary N = 32000, and start μ = 2τ = 6. Suppose the current logits give s_hat = 1.2, so ε = 0.2.

2^μ = 2^6 = 64
ε·2^μ = 0.2 × 64 = 12.8
N^(-ε) = 32000^(-0.2) ≈ 0.126   →  1 - 0.126 = 0.874
ratio = 12.8 / 0.874 ≈ 14.6
k = 14.6^(1/1.2) ≈ 9.4  →  k = 9

Sample from the top 9 tokens. Say the drawn token had probability 0.05, so its surprise is -log2(0.05) ≈ 4.32 bits — more surprising than we wanted. The error is e = 4.32 - 3 = 1.32, so μ ← 6 - 0.1×1.32 = 5.87. The next step starts from a slightly smaller budget and a slightly smaller k, steering the running surprise back toward 3. Each token is a small corrective nudge, not a jump.