veloxquant_mlx.codebooks.adaptive_codebook module¶
Data-driven codebook that refits its centroids from observed key statistics.
Unlike the closed-form distributions in strategies
(Gaussian, Beta, uniform), AdaptiveScalarCodebook buffers real
post-rotation vectors during a calibration window, builds an empirical
histogram, and runs Lloyd-Max against that histogram to fit centroids to the
actual data distribution — proxying to a default closed-form codebook until
enough vectors have been observed. Drop-in API-compatible with
ScalarCodebook.
- class veloxquant_mlx.codebooks.adaptive_codebook.AdaptiveScalarCodebook(b, d, n_calib=64, default_codebook=None, n_hist_bins=64, seed=0)[source]¶
Bases:
objectScalarCodebook that refits its centroids from observed key vectors.
- Phase 1 (calibration, first n_calib batches passed through observe()):
Buffer raw post-rotation vectors.
quantize() / dequantize() proxy to a default N(0, 1/d) ScalarCodebook so encode/decode keep working before calibration completes.
- Phase 2 (after n_calib vectors observed):
Build an empirical histogram from the buffer.
Run Lloyd-Max with that empirical PDF to fit centroids.
Replace the proxy codebook with a ScalarCodebook of fitted centroids.
The wrapper is API-compatible with ScalarCodebook (quantize/dequantize/ centroids_mx/centroids_numpy/k/b), so existing quantizers can use it as a drop-in replacement.
- Parameters:
b (
int) – Bit-width.d (
int) – Vector dimension (used for default codebook shape).n_calib (
int) – Number of vectors to buffer before fitting.default_codebook (
ScalarCodebook|None) – Initial ScalarCodebook used during calibration.n_hist_bins (
int) – Histogram bin count for empirical PDF estimation.seed (int)
- get_codebook()[source]¶
Return (centroids, boundaries) of the current codebook.
Boundaries are recomputed as midpoints between sorted centroids.
- observe(y)[source]¶
Accumulate post-rotation vectors during calibration.
The full batch is buffered (not truncated to the first remaining rows) so a single large call – e.g. an entire prefill batch – does not silently bias calibration toward whichever rows happened to be first. _fit() randomly subsamples the buffer down to n_calib.