chore: fix cosine residual calculation - #2015
Conversation
| <T as ArrowPrimitiveType>::Native: Float + Sum, | ||
| { | ||
| let v = arr.as_primitive::<T>(); | ||
| Ok(Arc::new(PrimitiveArray::<T>::from_iter_values(normalize(v.values()))) as ArrayRef) |
There was a problem hiding this comment.
FYI I plan to fix this soon, but normalize doesn't dispatch to the optimized SIMD kernels we have yet.
| let num_rows = data.len() / dimension; | ||
| let num_rows = data.num_rows(); | ||
|
|
||
| let (data, metric_type) = if self.metric_type == MetricType::Cosine { |
There was a problem hiding this comment.
should we move this higher? maybe to search?
There was a problem hiding this comment.
THe shuffler is still using this piece separately, tho
| .collect::<Vec<_>>(); | ||
| let data = T::ArrayType::from(values); | ||
| FixedSizeListArray::try_new_from_values(data, self.dimension as i32)? | ||
| normalize_fsl(&fsl)? |
There was a problem hiding this comment.
Should we remove cosine type PQ entirely? And let IVF handle normalization and resid?
There was a problem hiding this comment.
^^ can we add a panic! that make sure we don't create a cosine PQ?
7ef3c0c to
9d8a2e5
Compare
9d8a2e5 to
2be5e30
Compare
| let ivf_transform = Arc::new(IvfTransformer::new( | ||
| centroids.clone(), | ||
| metric_type, | ||
| MetricType::L2, |
There was a problem hiding this comment.
what about dot product?
| let mut best_stddev = f32::MAX; | ||
|
|
||
| let rng = rand::rngs::SmallRng::from_entropy(); | ||
| let rng = SmallRng::from_entropy(); |
There was a problem hiding this comment.
nit: add a todo to use seeds
| .collect::<Vec<_>>(); | ||
| return compute_partitions_l2(centroids_array, &normalized, dimension) | ||
| .collect(); | ||
| panic!("KMeans: should not use cosine distance to train kmeans, use L2 instead."); |
There was a problem hiding this comment.
add this check in the constructor?
There was a problem hiding this comment.
The match arm needs a branch to handle MetricType::Cosine tho.
| let partition_ids = | ||
| self.ivf | ||
| .find_partitions(&query.key, query.nprobes, self.metric_type)?; | ||
| let mut query = query.clone(); |
There was a problem hiding this comment.
how cheap is this copy, can we avoid it for the L2 path?
There was a problem hiding this comment.
This is very cheap, the query vector is zero copy (copied by pointer), other than that, there are just 6 other string/int fields.
| // vector_column, | ||
| // ))); | ||
| // }; | ||
| let mt = if metric_type == MetricType::Cosine { |
There was a problem hiding this comment.
nit: make a macro for this?
| }); | ||
| } | ||
| pre_filter.wait_for_ready().await?; | ||
| println!("PQIndex::search: metric type: {:?}", self.metric_type); |
| training_data = normalize_fsl(&training_data)?; | ||
| } | ||
|
|
||
| println!("PQ Training, ivf: {:?} ", ivf); |
| span!(Level::INFO, "compute residual for PQ training") | ||
| .in_scope(|| ivf2.compute_residual(&training_data, None)) | ||
| .await? |
There was a problem hiding this comment.
I think we still need to make sure we don't residulize for dot.
So, maybe something like pq_params.should_residulize?
There was a problem hiding this comment.
So the design lets the caller decide whether to run residual or not. If provided IVF, it runs residual PQ.
So PQ is orthogonal to distance type as much as possible.
There was a problem hiding this comment.
but we always residualize currently.
There was a problem hiding this comment.
need a flag for dot distance to not residualize
There was a problem hiding this comment.
Ok, checked it on the caller
chebbyChefNEQ
left a comment
chebbyChefNEQ
left a comment
There was a problem hiding this comment.
just a few println! and pq residual needs to be fixed
…cking The training and assignment dispatches accepted any distance type for a float column, and the membership pass then panicked with "not supported" unless the centroid HNSW happened to be built, which bypasses that match. Narrow both dispatches to L2 and Dot so the existing descriptive error arms handle the rest. lance.util.KMeans documents cosine, so the binding now normalizes its input and clusters with l2, the way the index build path does. That makes cosine behave the same at every k instead of panicking below the HNSW threshold. The module and new_with_params doc comments claimed cosine inputs are normalized each iteration. That was implemented in lance-format#1723 and removed in lance-format#2015 while the comment was carried forward.
3 participants
No description provided.