codebook-metrics

Small, functional toolkit for the math of neural speech tokenizers and discrete codebooks (VQ-VAE / RVQ style). It computes codebook utilization and perplexity, code entropy and effective size, empirical rate-distortion curves, residual-VQ bitrate accounting, and high-rate theoretical bounds — all over the raw code indices and embeddings you already have, with no training loop required. Every function ships with the formula in its docstring and a test that pins a known closed form.

Install

pip install codebook-metrics
# optional extras
pip install "codebook-metrics[plot]"   # matplotlib RD-curve helper
pip install "codebook-metrics[torch]"  # torch embeddings adapter

API

function	computes	formula
perplexity(codes, K)	effective codewords in use	2^{H(p)}
usage_rate(codes, K)	fraction of live codewords	#{k: n_k>0} / K
dead_codes(codes, K)	never-selected codeword indices	{k: n_k = 0}
code_entropy(codes, K)	code distribution entropy (bits)	-Σ p_k log2 p_k
effective_codebook_size(codes, K)	perplexity as a "size"	2^{H(p)}
normalized_mse(x, x_hat)	scale-free distortion	E‖x-x̂‖² / E‖x‖²
code_rate_bits(codes, K)	entropy-coded rate	H(p)
rvq_cumulative_rate(stages, K)	running RVQ bitrate	Σ_s H(p_s)
empirical_rd_curve(x, q, sizes)	sweep K → (rate, distortion)	—
high_rate_distortion_bound(h, R)	asymptotic MSE floor	(1/12) 2^{2(h-R)}
shannon_lower_bound_gaussian(σ², R)	Gaussian RD function	σ² 2^{-2R}
lloyd_optimality_gap(D, h, R)	distance from high-rate optimum	D / D_bound
gini_coefficient(codes, K)	usage imbalance	sorted-p Gini

Examples

Utilization of a batch of codes:

import numpy as np
from codebook_metrics import perplexity, usage_rate, dead_codes

codes = np.random.randint(0, 256, size=10_000)   # one RVQ stage, K = 256
print(perplexity(codes, 256))   # ~ effective codewords used
print(usage_rate(codes, 256))   # fraction of the book that is alive
print(dead_codes(codes, 256))   # indices that never fired

Rate-distortion of a Lloyd codebook plus the high-rate bound:

import numpy as np
from codebook_metrics import empirical_rd_curve, high_rate_distortion_bound
from codebook_metrics.quantize import lloyd_quantize

x = np.random.randn(2000, 8)
curve = empirical_rd_curve(
    x, lambda d, k: lloyd_quantize(d, k, seed=0).quantize(d),
    code_sizes=[2, 4, 8, 16, 32, 64],
)
for p in curve:
    print(f"K={p.num_codes:3d}  R={p.rate:.2f} bits  D={p.distortion:.4f}")

print("floor:", high_rate_distortion_bound(differential_entropy_bits=4.0, rate_bits=6.0))

Cumulative bitrate of a residual VQ:

from codebook_metrics import rvq_cumulative_rate
from codebook_metrics.quantize import residual_vq

x_hat, stage_codes, _ = residual_vq(x, num_codes=64, num_stages=4, seed=0)
print(rvq_cumulative_rate(stage_codes, num_codes=64))  # bits/token after each stage

Guarantees / sanity properties

• perplexity and effective_codebook_size lie in [1, K] (uniform usage ⇒ exactly K).
• code_entropy lies in [0, log2 K]; normalized_entropy / usage_balance in [0, 1].
• code_rate_bits never exceeds log2 K.
• rvq_cumulative_rate is non-decreasing across stages.
• normalized_mse is invariant to a global scaling of the features.
• lloyd_optimality_gap == 1 exactly at the high-rate bound.

Changelog

See CHANGELOG.md.

Citation

If you use this library, please cite:

@misc{kong2026codebookmetrics,
  title  = {Measuring Discrete Speech Tokenizers: Utilization, Rate, and
            High-Rate Distortion Bounds},
  author = {Kong, Tai},
  year   = {2026},
  note   = {preprint},
}