slow-J commented on code in PR #16790:
URL: https://github.com/apache/lucene/pull/16790#discussion_r4445874681
##########
lucene/core/src/java/org/apache/lucene/util/quantization/OptimizedScalarQuantizer.java:
##########
@@ -192,25 +176,15 @@ public QuantizationResult scalarQuantize(
assert bits > 0 && bits <= 8;
float[] intervalScratch = new float[2];
int points = 1 << bits;
- double vecMean = 0;
- double vecVar = 0;
- float norm2 = 0;
- float centroidDot = 0;
- float min = Float.MAX_VALUE;
- float max = -Float.MAX_VALUE;
- for (int i = 0; i < vector.length; ++i) {
- if (similarityFunction != EUCLIDEAN) {
- centroidDot += vector[i] * centroid[i];
- }
- vector[i] = vector[i] - centroid[i];
- min = Math.min(min, vector[i]);
- max = Math.max(max, vector[i]);
- norm2 += (vector[i] * vector[i]);
- double delta = vector[i] - vecMean;
- vecMean += delta / (i + 1);
- vecVar += delta * (vector[i] - vecMean);
- }
- vecVar /= vector.length;
+ float[] centerStats = new float[5];
+ VectorUtil.osqCenter(vector, centroid, centerStats);
+ float centroidDot = centerStats[0];
+ float min = centerStats[1];
+ float max = centerStats[2];
+ float norm2 = centerStats[3];
+ double vecMean = centerStats[4] / (double) vector.length;
+ // Rounding can push the mean of squares minus the squared mean just below
zero.
+ double vecVar = Math.max(0, norm2 / (double) vector.length - vecMean *
vecMean);
Review Comment:
Tried to hack something up to measure this with results as follows:
A SIMD Welford (per lane, merged at the end) makes `scalarQuantizeVector`
6-8% slower at 384-4096 dims.
On the first 20K Cohere docs, the variance is off by at most 5e-7 both with
or without it. The simple version goes wrong when a vector's mean after
centering is ~1000x its std.
I would like to keep the simple one, WDYT?
##########
lucene/benchmark-jmh/src/java/org/apache/lucene/benchmark/jmh/VectorUtilBenchmark.java:
##########
@@ -118,6 +126,21 @@ public void init() {
floatsB[i] = random.nextFloat();
shortsB[i] = Float.floatToFloat16(floatsB[i]);
}
+ // Rotate through 64 vectors, as the number of descent steps depends on
the vector. They are
+ // centered here and the centroid is zero, so scalarQuantize leaves them
as they are and every
+ // call can be checked against the first.
+ quantizeVectors = new float[64][size];
+ expectedQuantizeResults = new
OptimizedScalarQuantizer.QuantizationResult[64];
+ quantizeCentroid = new float[size];
Review Comment:
Will do in latest commit. The centroid is now random and each call copies
the next vector into a scratch array first. The copy is ~1-2% of
`scalarQuantizeVector` at 384-4096 dims, so I left it in the benchmark.
##########
lucene/core/src/java/org/apache/lucene/internal/vectorization/DefaultVectorUtilSupport.java:
##########
@@ -459,6 +459,80 @@ private float quantizeFloat(float v, byte[] dest, int
destIndex) {
}
}
+ @Override
+ public void osqCenter(float[] vector, float[] centroid, float[] stats) {
+ float dot = 0;
+ float min = Float.MAX_VALUE;
+ float max = -Float.MAX_VALUE;
+ float sumSq = 0;
+ float sum = 0;
+ for (int i = 0; i < vector.length; i++) {
+ dot += vector[i] * centroid[i];
+ float x = vector[i] - centroid[i];
+ vector[i] = x;
+ min = Math.min(min, x);
+ max = Math.max(max, x);
+ sumSq += x * x;
+ sum += x;
+ }
+ stats[0] = dot;
+ stats[1] = min;
+ stats[2] = max;
+ stats[3] = sumSq;
+ stats[4] = sum;
+ }
+
+ @Override
+ public void osqGridStats(float[] vector, float lower, float upper, int
points, double[] stats) {
+ osqGridStatsImpl(vector, lower, upper, points, stats);
+ }
+
+ static void osqGridStatsImpl(
+ float[] vector, float lower, float upper, int points, double[] stats) {
+ float stepInv = (points - 1f) / (upper - lower);
+ float step = (upper - lower) / (points - 1f);
+ long sumK = 0;
+ long sumKK = 0;
+ double sumXK = 0;
+ double sumXD = 0;
+ double sumDD = 0;
+ double sumX = 0;
+ for (float x : vector) {
+ // Adds 0.5 and truncates, as Math.round is costly to match in the
Panama version.
+ int k = (int) ((Math.min(Math.max(x, lower), upper) - lower) * stepInv +
0.5f);
+ float d = x - (lower + k * step);
+ sumK += k;
+ sumKK += (long) k * k;
+ sumXK += x * k;
+ sumXD += x * d;
+ sumDD += d * d;
+ sumX += x;
+ }
+ stats[0] = sumK;
+ stats[1] = sumKK;
+ stats[2] = sumXK;
+ stats[3] = sumXD;
+ stats[4] = sumDD;
+ stats[5] = sumX;
+ }
+
+ @Override
+ public int osqAssign(float[] vector, float lower, float upper, float step,
byte[] dest) {
+ return osqAssignImpl(vector, 0, lower, upper, step, dest);
+ }
+
+ // The Panama version calls this for its tail, with start where its vector
loop stopped.
+ static int osqAssignImpl(
+ float[] vector, int start, float lower, float upper, float step, byte[]
dest) {
+ int sum = 0;
+ for (int i = start; i < vector.length; i++) {
+ int k = Math.round((Math.min(Math.max(vector[i], lower), upper) - lower)
/ step);
Review Comment:
Couldn't find existing assertions, I added some in next commit. With asserts
on, the first 20K Cohere docs at 1, 2, 4, 7 and 8 bits never hit lower > upper
or step < 0, Panama or scalar. Forgot to say earlier that the Panama code has
`x.max(lower).min(upper)` and changing the scalar one would add a difference
since there are some behaviour changes like a 0.0 vs -0.0 edge case that would
throw with `Math.clamp`. WDYT?
--
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]
