🚨 Clip out-of-range values in int8/uint8 quantization

[anjaliy11]

Hello!

Pull Request overview

• Clip quantized values to the valid range before casting to int8/uint8 in quantize_embeddings

Details

quantize_embeddings computes (embeddings - starts) / steps and then casts straight to np.uint8/np.int8. When an embedding value falls outside the range used to compute starts/steps (either the provided ranges, a calibration_embeddings set, or the embeddings themselves), the intermediate value lands outside [0, 255] before the cast. NumPy doesn't clip on an out-of-range float-to-int cast, it wraps. So a value that should saturate at 255 can silently come out as e.g. 44 instead.

This is easiest to see with a calibration set that doesn't cover the data being quantized:

calibration_embeddings = np.array([[1, 20, -3], [4, 5, -60]], dtype=np.float32)
dataset = np.array([[-1, 15, 1]], dtype=np.float32)

quantize_embeddings(dataset, precision="int8", calibration_embeddings=calibration_embeddings)
# Before: [[ -42   42 -112]]
# After:  [[-128   42  127]]

The fix adds np.clip(q_vals, 0, 255) before the final cast for both uint8 and int8, so out-of-range values saturate at the dtype bounds instead of wrapping.

Originally reported in #3159, with a fix proposed in #2865 by @meiravgri in 2024. That PR was never merged, and quantize_embeddings has since moved from sentence_transformers/quantization.py to sentence_transformers/util/quantization.py. I've applied the fix to the current location and added a regression test (test_quantize_clips_out_of_range_values) covering both precisions.

Fixes #3159

Anjali Yadav

---

Maintainer note

I pushed a small simplification on top, hoisting the shared np.clip(np.floor(...), 0, 255) out of the two branches. No behavior change.

Worth calling out for users: beyond the out-of-range clipping, this PR also introduces np.floor before the cast, which changes int8 output for in-range values too. The old int8 path cast x - 128 with truncation-toward-zero, so it rounded toward bucket 128 (up below it, down above it). Flooring first makes int8 a consistent uniform quantizer, matching what uint8 already did.

In practice, int8 values in the lower half of the range now shift down by one level (one bucket out of 256). This is a small fidelity improvement (measured against the original float32, it cuts reconstruction error by roughly 2.5x under center-of-bucket dequantization, and is a wash for the raw int8 scoring used in retrieval today), but it does mean int8 embeddings are no longer bit-identical to those produced by older versions. If you need exact reproducibility against an already-quantized corpus, re-quantize it with this version. uint8 is unaffected, as floor and truncation agree for positive values.

• Tom Aarsen

[tomaarsen]

Hello!

Thanks for picking this up, and for digging back through #3159 and #2865 for the context. This is a real bug and the fix is clean. I pushed a small simplification on top (hoisting the shared np.clip(np.floor(...), 0, 255) out of the two branches, no behavior change).

As far as I can tell, the fix is correct, and there was indeed a bug previously. The regression test is also good.

One heads-up: np.floor also shifts int8 in-range values. It's worth noting since it is not in the PR description: the np.floor is not needed for the wraparound fix itself, and it also changes int8 output for in-range values (the old path truncated x - 128 toward zero), shifting the lower half of the range by one level. I checked whether that is churn or a real change, and it is a modest improvement: floor makes int8 a consistent uniform quantizer, matching what uint8 already did. That roughly halves reconstruction error against the original fp32 under center-of-bucket dequant (about 0.0010 to 0.0004 MAE on random embeddings) and is a wash for the raw int8 retrieval scoring we do today. uint8 is unchanged.

Since it does move existing int8 outputs by one level, so this is a bit breaking technically. I think it's still worth pushing.

Thanks again for the clear write-up and the test!

• Tom Aarsen

[Copilot]

Pull request overview

This PR addresses incorrect wrap-around behavior when scalar-quantizing embeddings to int8/uint8 by ensuring out-of-range bucket values saturate at dtype bounds instead of wrapping during NumPy casts.

Changes:

• Add clipping of computed quantization bucket values to [0, 255] in quantize_embeddings before casting to uint8/int8.
• Add a regression test covering out-of-range saturation for both int8 and uint8 precisions.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
sentence_transformers/util/quantization.py	Clips computed bucket values prior to casting in int8/uint8 quantization to prevent wrap-around.
tests/util/test_quantization.py	Adds a regression test ensuring out-of-range values saturate instead of wrapping for int8/uint8 quantization.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[tomaarsen]

As I mentioned in my comment, this might be fine.

[anjaliy11]

Thanks Tom, really appreciate you taking the time to dig into this and push the simplification.
Makes sense on the rounding floor giving a uniform quantizer for int8 (matching what uint8 already did) is a sensible improvement, and the MAE numbers back it up. Agreed it's worth taking even though it shifts in-range values by one level; the old truncate-after-shift behavior wasn't something users should have been relying on anyway.
Thanks again for reviewing so thoroughly, and for closing the loop on the Copilot comment. Let me know if there's anything else you'd like me to adjust before merging.

• Anjali