Skip to content

feat: Add KL Divergence command#1146

Open
spicyneuron wants to merge 5 commits intoml-explore:mainfrom
spicyneuron:kl-divergence
Open

feat: Add KL Divergence command#1146
spicyneuron wants to merge 5 commits intoml-explore:mainfrom
spicyneuron:kl-divergence

Conversation

@spicyneuron
Copy link
Copy Markdown
Contributor

This PR adds a new mlx_lm.kld command for comparing a candidate model against a baseline. Main use case is comparing quantizations of the same model.

Some design decisions worth mentioning:

  • The cache only stores the baseline top-K logprobs at each position, plus a single tail-mass bucket for the rest of the distribution. This is similar in spirit to the DWQ approach and avoids creating massive full-logit caches while still preserving the part of the distribution that matters most.
  • The metric direction is explicit as KL(baseline || model).
  • The cache records tokenizer metadata and validates candidate compatibility before evaluation. That helps prevent accidentally comparing models whose vocab sizes or token-id mappings do not line up.
  • Eval token loading now goes through a shared load_eval_tokens() helper so perplexity and kld use the same sampling path. That keeps the evaluation setup more consistent and easier to reason about across tools.

This PR also wires the new command into the CLI and adds tests covering cache creation, cache reuse, KL computation, and tokenizer mismatch handling.

@spicyneuron spicyneuron changed the title Add KL Divergence command feat: Add KL Divergence command Apr 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant