[selimfedakar]

Summary

The int8 path in quantize_embeddings scaled the embeddings through a float16 buffer before rounding to integers:

buf = embeddings.astype(np.float16, copy=True)
np.divide(buf, scale, out=buf)
np.rint(buf, out=buf)

float16 keeps only ~3 significant decimal digits, so embeddings / scale rounds to a different nearest integer than the exact quotient for a fair number of entries. The result is that a non-trivial fraction of int8 codes end up off by ±1 versus the mathematically correct quantization — for ordinary, small-magnitude embeddings, not just extreme inputs.

On a normal 2000 × 256 matrix at std = 0.1, ~0.65% of codes (3337 / 512000) differ from the float32-correct result, touching ~80% of token rows. Measured against the true unrounded quotient, the float32 computation is closer on every differing entry and never worse, so the float16 buffer was strictly less accurate with no accuracy upside — its only benefit was a smaller transient buffer.

As a side effect, casting to float16 first also turned any embedding value above the float16 maximum (65504) into inf before scaling (RuntimeWarning: overflow encountered in cast), with clip happening to rescue the dominant code.

Change

Compute the scaled values in float32 instead of float16. Still a single copy; the transient buffer is 2× larger, but the codes now match the exact computation and the overflow is gone.

Tests

• test_quantize_embeddings_int8_matches_float32_reference — quantizes a 2000 × 256 matrix and asserts the codes equal the float32 reference. Fails on the previous float16 buffer, passes with this change.
• test_quantize_embeddings_int8_handles_values_above_float16_max — covers values above the float16 maximum.

Full tests/test_quantization.py and tests/test_model.py pass (83 tests); ruff check, ruff format --check, and mypy model2vec/quantization.py are clean.

Closes #405.

---

Verification level: measured and reproduced locally on main; regression test fails on revert of the fix. AI-assisted investigation, reviewed and verified by me.