Fix UMAP outlier issue by checking for outliers and shuffling

[jinsolp]

Closing #6454

Main difference between out simplicial set embedding and CPU UMAP was in negative sampling.
We should use updated values (value after adding gradients) in the negative sampling stage.

Dispatched to two kernels (and three usages) based on `n_components. Fixed like below.

• optimize_batch_kernel_reg (n_components=2): update the current_reg register value (used later in the negative sampling stage) along with grads
• optimize_batch_kernel (with shared memory): distinguish current_buffer (which used to JUST hold the gradient) from the grad_buffer. Now current_buffer and grad_buffer corresponds to the current_reg and grads registers in the register-approch kernel.
• optimize_batch_kernel (without shared memory): untouched because the grads are applied directly to global memory. This updated value in global memory is read directly for negative sampling later on.

Visualizations 2D

50K samples random selected for plotting.
From the left

• CPU KNN + CPU UMAP
• GPU KNN + CPU UMAP
• GPU KNN + GPU UMAP Before fix
• GPU KNN + GPU UMAP After fix in this PR

Using dataset 639K x 384

Using dataset 1.8M x 384

Visualizations 3D

50K samples random selected for plotting.
Plotting the same dataset with n_components=3 (Which uses the second kernel).
From the left

• GPU KNN + CPU UMAP
• GPU KNN + GPU UMAP Before fix
• GPU KNN + GPU UMAP After fix in this PR

Using dataset 639K x 384 (was already doing pretty well without outliers, still doing well)

Using dataset 1.8M x 384
before fix had outliers.

[viclafargue]

Thanks @jinsolp! It looks great.

[viclafargue]

Importantly, when random_state is set, current != cur_write and other != oth_write as updates accumulate in separate buffer to allow high precision deterministic accumulation of updates. It looks like we may still have outliers in this case? But, I guess that is acceptable for now.

[jinsolp]

Have to add unit test for outlier checking. Plan is to grab a large enough dataset that originally fails (i.e. has outliers).
We can use the all-neighbors python API to build the knn graph quickly, give that as precomputed graphs to both CPU and cuML UMAP, and compare the resulting embeddings.

e.g. get min/max values of the CPU embedding, and check if all values in out embedding is within a certain threshold of that range.

[viclafargue]

Thanks @jinsolp! It looks great! However, I believe that we would have to apply the shuffling before make_epochs_per_sample is called (see comment).

[csadorf]

We should not directly import the cuvs Python API in cuML. If we do, then we need to cuvs (not just libcuvs) to our test dependencies. CC @divyegala

[jinsolp]

we can get rid of this, but if we do, we have to run the full e2e cpu umap on a not-so-small dataset for comparison (because outliers don't show up with small datasets)

[csadorf]

We can consider adding cuvs to the test dependencies, but then let's make sure to guard the cuvs import with pytest.importorskip. I'll let @divyegala chime in on this.

[divyegala]

Yeah, it is fine as a testing dependency. Please point me to the commit that adds cuvs as so, I will verify that we don't leak the dependency by mistake.

[jinsolp]

Using importorskip for now. Left an issue: #7279

[csadorf]

It looks like you are still using a direct import here?

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[csadorf]

Thanks a lot!

[divyegala]

/merge