[Qubitium]

What

MLXTeacherKL now uses the algebraic form of the KL and folds the constant teacher-entropy term out of the optimizer step.

KL(t || s) = sum(t * log t) - sum(t * s_logits) + logsumexp(s_logits)
             ^^^^^^^^^^^^   ^^^^^^^^^^^^^^^^^   ^^^^^^^^^^^^^^^^^^^
             teacher only   student           student

sum(t * log t) depends only on the frozen teacher, so it is computed once per batch while the hidden-state cache is built and reused for every step. The per-step chunk loss then runs two student reductions over the vocabulary instead of materializing the full-vocabulary student log-probabilities.

Result

Qwen3-0.6B, 196 interpolated projections, 3 options (W2/W3/W4), chunk_tokens=128, MLX 0.32.4, Apple M5 Ultra. End-to-end RCO step (mx.value_and_grad of the teacher-KL objective, best of 3):

calibration geometry	tokens	before	after	speedup
16 x 128	2,048	753.3 ms	726.6 ms	1.04x
1 x 4096	4,096	475.8 ms	438.0 ms	1.09x
4 x 4096 (paper)	16,384	1962.6 ms	1798.3 ms	1.09x

Numerics

The rewrite is an identity, not an approximation, so the only difference is floating-point reassociation.

• Objective value at the paper geometry: 0.937571 -> 0.937567 (absolute 4e-6, relative 4e-6).
• Gradient wrt alpha (196x3): max abs 2.87e-4 against a max gradient magnitude of 0.122 (0.24%), mean abs 2.9e-5, norm ratio 1.000632.
• The existing full-vocabulary Torch oracle (tests/test_rco_paper.py::test_teacher_kl_value_and_gradient_match_full_vocabulary_oracle, both dtypes, three chunk sizes) passes unchanged, as do the identical-model-is-zero and resident-vs-disk parity tests.

The identity relies on sum(softmax(t_logits)) == 1, which holds to f32 rounding; logsumexp is subtracted per row, so no large-magnitude term is accumulated across the vocabulary.

Memory

Per chunk (chunk_tokens x 151936), the old form materialized student_logp, teacher_logp and exp(teacher_logp); the new form materializes student_logits and the teacher probabilities. That removes one chunk-sized f32 buffer (78 MB at chunk_tokens=128).

The entropy cache is one f32 scalar per batch. cleanup() now clears it, and test_mlx_search_restores_original_linear_and_exact_budget asserts it is released. No dense or dequantized weight is retained.

Construction cost rises by one teacher pass over the corpus (the entropy pass), amortized over the whole search.

Tests

tests/test_rco_paper.py, tests/test_rco_mlx.py, tests/test_gsq_layer_bit_allocation.py: 62 passed, 1 skipped. ruff check clean.

Scope

MLX only. The Torch LlamaTeacherKL in rco.py still uses the log-softmax form; it is the fallback path on this hardware and is left untouched so this PR stays a single concern.