xande commented on PR #15903: URL: https://github.com/apache/lucene/pull/15903#issuecomment-4184783958
Thanks everyone for the incredibly thorough feedback. This is exactly the kind of review I was hoping for. I spent some time with Kiro going back through the paper, the Elastic OSQ blog posts, and the actual codec source. Hopefully this will help the community to decide on the path forward (which could be to be improving OSQ with learning from TQ). Also the Lucene communmity might find discussions and findings in this reference implementation helpful - https://github.com/tonbistudio/turboquant-pytorch ## On TurboQuantMSE vs TurboQuantProd (@mccullocht) You're correct — this implements TurboQuantMSE only. The [paper](https://arxiv.org/abs/2504.19874)'s TurboQuantProd variant (Algorithm 2, Section 3.2) applies a (b-1)-bit MSE quantizer followed by a 1-bit [QJL](https://arxiv.org/abs/2402.09644) transform on the residual vector r = x - dequant(quant(x)), yielding an unbiased inner product estimator at total bit-width b. Based on the community comments on https://github.com/tonbistudio/turboquant-pytorch?tab=readme-ov-file#v3-improvements-community-informed, MSE alone is enough. ## On the non-uniform grid and performance (@tveasey, @mccullocht) @tveasey's analysis is spot-on and this is the most important tradeoff. I had Kiro pull up the actual centroid values from [`BetaCodebook.java`](https://github.com/apache/lucene/pull/15903/files#diff-BetaCodebook) and the OSQ initialization intervals from the [Elastic blog](https://www.elastic.co/search-labs/blog/scalar-quantization-optimization) to compare side by side. The TQ centroids are Lloyd-Max optimal for the Gaussian distribution, which means non-uniform spacing. For b=2, TQ places centroids at [-1.510, -0.453, 0.453, 1.510] vs OSQ's initial uniform grid at [-1.493, -0.498, 0.498, 1.493] (before OSQ's per-vector interval refinement). The spacing difference is small but the consequence is large: you can't decompose the dot product into integer arithmetic plus scalar corrections. The current implementation uses a float gather-multiply-accumulate loop: for each packed byte, extract indices, look up centroid values from a 2^b-entry table, multiply by the rotated query coordinate, accumulate. At b=4 this is a 16-entry LUT. The JMH numbers (313K ops/s at d=4096, ~3.2µs per candidate) reflect JVM auto-vectorization of this loop, not explicit SIMD [benchmak data has to be corss-checked]. For comparison, @tveasey notes OSQ gets 4-8x performance vs float arithmetic with well-crafted SIMD integer dot products. That's a real and significant gap. The question is whether the recall advantage at low bits justifies the scoring cost, or whether the other properties (no calibration, byte-copy merge, streaming) matter enough for specific workloads. The framing: TQ is not going to beat OSQ on scoring throughput (even though it looks like there are ways to make it faster). Its value proposition is the combination of (a) no calibration overhead, (b) merge-friendly architecture, and (c) high-dimension support beyond 1024d. If scoring throughput is the bottleneck, OSQ wins. On the Panama path forward: @mccullocht is right that `Vector.shuffle()` is the key operation for LUT-based scoring. For b=4, the 16-entry centroid table fits in a single 512-bit register (16 × float32), and `vpermps` (AVX-512) or `tbl` (NEON) can do the gather. Whether the JVM's Panama implementation of `VectorShuffle` actually emits these instructions efficiently is an open question. ## On the PQ-style codebook approach (@mccullocht, @tveasey) @mccullocht suggested a PQ-like approach with precomputed per-query lookup tables. For TQ at b=4, the current implementation already does something similar: the centroid table is the LUT, and the inner loop is `sum += query[i] * centroids[packed_index[i]]`. The 16-entry table is small enough to stay in L1 cache. The larger question @tveasey raises about whether the better route is residual quantization with centering is interesting. I had Kiro review the [OSQ blog](https://www.elastic.co/search-labs/blog/scalar-quantization-optimization)'s "Refining the quantization interval" section in detail. It thinks "OSQ's per-vector interval optimization — the coordinate descent that minimizes dot product error weighted toward the document vector direction (with λ=0.1) — is clever (Kiro thinks so, I agree:)) and exploits structure that TQ's data-oblivious approach deliberately ignores. For data that has exploitable per-dimension structure, OSQ should win on recall-per-bit." ## On the global transform and per-segment overhead (@mikemccand, @mccullocht) Yes, this PR uses one global transform per field (seed derived from field name hash). This means: - No per-segment rotation storage - Merge can byte-copy quantized data (not yet implemented, but architecturally supported since all segments share the same rotation) - Query rotation is done once per query, not per-segment @mccullocht asked whether per-segment transforms would increase randomness. Theoretically yes — different random rotations would give independent quantization errors across segments, which could improve recall after merging results from multiple segments. But the cost is losing byte-copy merge and needing a separate query rotation per segment. ## On the AI collaboration process (@mikemccand) The original prompts and iteration history are preserved in the commit history and in the TURBOQUANT_*.md files in the repo root. The session log (`TURBOQUANT_SESSION_LOG.md`) documents the full interaction timeline including the expert review simulation rounds that shaped the architecture. The implementation plan (`TURBOQUANT_IMPLEMENTATION_PLAN.md`) shows the phased approach with gate conditions. I agree these could be treated like source artifacts, though I could not find a good place for them and deleted in one of the latest commits. In fact I'm using Kiro right now to research these responses, cross-referencing those TURBOQUANT*.md plans, the paper, the OSQ blog posts, and the actual source code to fact-check my claims before posting. ## On adding random rotation to OSQ (@benwtrent, @tveasey) @benwtrent mentioned that adding random rotation to OSQ for non-Gaussian components would be straightforward. This could be the most interesting direction (and MAY explain our low recall values on internal data sets): OSQ with a Hadamard pre-rotation would get the best of both worlds — the rotation homogenizes coordinate distributions (helping with non-Gaussian embeddings), and OSQ's per-vector interval optimization + integer arithmetic scoring handles the rest. The rotation would add some latency per query at d=4096 but OSQ's scoring would remain fast. Worth exploring? -@xande, co-authored with Kiro/Opus 4.6 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
