`cuml.UMAP` embeddings result in outliers

[jcrist]

This came out of running some analysis on hyperparameters in our nn-descent UMAP implementation, but seems to affect both build algos (indicating it's an issue in the steps after building the knn graph).

On sufficiently wide datasets, the embedding_ produced by cuml.UMAP includes some significant outliers that aren't found in the umap.UMAP implementation.

Reproducer (this takes a few minutes to run due to the CPU implementation):

import time

import umap
import cuml
from sklearn.datasets import make_blobs


def fit_and_evaluate(label, umap, X, cpu=None):
    start = time.time()

    umap.fit(X)

    duration = time.time() - start
    mins = int(duration // 60)
    secs = duration % 60

    print(f"{label}:")
    print(f"- time: {mins:02d}:{secs:05.2f}")
    print(f"- min: {umap.embedding_.min(axis=0).tolist()}")
    print(f"- max: {umap.embedding_.max(axis=0).tolist()}")
    if cpu is not None:
        # semi-arbitrarily define outliers as points outside 2 stds of cpu bounds
        cpu_min = cpu.embedding_.min(axis=0)
        cpu_max = cpu.embedding_.max(axis=0)
        cpu_std = cpu.embedding_.std(axis=0)
        lower_bound = cpu_min - (2 * cpu_std)
        upper_bound = cpu_max + (2 * cpu_std)
        n_outliers = (
            ((umap.embedding_ < lower_bound) | (umap.embedding_ > upper_bound))
            .any(axis=1)
            .sum()
        )
        print(f"- n_outliers: {n_outliers}")


# Fit 3 different UMAPs
common = {"n_neighbors": 15, "init": "random"}
umap_nnd = cuml.UMAP(build_algo="nn_descent", **common)
umap_bf = cuml.UMAP(build_algo="brute_force_knn", **common)
umap_cpu = umap.UMAP(**common)

X, _ = make_blobs(centers=10, n_features=512, n_samples=500_000)

fit_and_evaluate("CPU", umap_cpu, X)
fit_and_evaluate("NN-Descent", umap_nnd, X, cpu=umap_cpu)
fit_and_evaluate("Brute Force KNN", umap_bf, X, cpu=umap_cpu)

Output:

CPU:
- time: 04:41.46
- min: [-7.943415641784668, -8.329208374023438]
- max: [17.919870376586914, 15.174100875854492]
NN-Descent:
- time: 00:05.88
- min: [-646.9306030273438, -815.4732055664062]
- max: [782.5948486328125, 438.62982177734375]
- n_outliers: 72
Brute Force KNN:
- time: 00:35.99
- min: [-263.966796875, -427.2183837890625]
- max: [260.0986328125, 11.626998901367188]
- n_outliers: 33

NN-Decent does worse than Brute Force KNN in all runs I've done, but the difference is run dependent. Some runs brute force doesn't appear as bad. I think this probably means the input KNN graph has an influence here, but the actual bug is probably downstream of the knn graph construction.

The actual plots generated here appear pretty identical across all 3 versions once the GPU versions are trimmed to exclude things outside the 0.005/0.955 quantiles (a stricter definition of "outlier" than used above). Without this trimming the plots have a much larger bounds since the min/max on the GPU versions are sometimes orders-of-magnitude larger than the CPU version.

[csadorf]

@divyegala @jinsolp Can you advise on this?

[jcrist]

Barring other suggestions, I intend to do a bit more debugging to better isolate the location of the divergence. We can isolate the knn step off by passing in a precomputed_knn to cuml.UMAP from umap.UMAP to see if the divergence still occurs (I expect it will). After that it's picking through the downstream steps, which will be slower going. Modifying fussy_simplicial_set to support precomputed_knn would let us isolate if the issue is in the generation of fss_graph or afterwards in the embedding step (this might be a useful change regardless of helping debug).

[viclafargue]

We can isolate the knn step off by passing in a precomputed_knn

I tried something along these lines this morning. Generating a dataset with lower number of clusters and lower value for cluster_std as it may impact the quality of the KNN graph. And tried to run a precomputed KNN graph as a separate step : knn_graph = kneighbors_graph(X, n_neighbors=15, metric='euclidean', mode='distance', include_self=True). This does not seem to have an effect on the number of outliers, therefore I guess that the issue is indeed part of the downstream steps.

The fuzzy_simplicial_set function allows to provide a precomputed graph with the arguments knn_indices and knn_dists. It could indeed be interesting to test things out by comparing the results with cuML and reference fuzzy_simplicial_set simplicial_set_embedding functions.

[viclafargue]

Could test it with 300K points. Surprising results... It looks like there is an issue when running the simplicial_set_embedding function on a fuzzy simplicial set graph trained with the reference implementation.

CPU fss - CPU embedding:
- min: [-11.508538246154785, -10.999808311462402]
- max: [20.184934616088867, 19.740995407104492]
- n_outliers: 0
GPU fss - CPU embedding:
- min: [-9.22128963470459, -11.480342864990234]
- max: [19.570436477661133, 19.142061233520508]
- n_outliers: 0
CPU fss - GPU embedding:
- min: [-634.0408935546875, -664.8304443359375]
- max: [110.56887817382812, 197.0208740234375]
- n_outliers: 25
GPU fss - GPU embedding:
- min: [-17.3656005859375, -15.5699462890625]
- max: [36.01678466796875, 35.7691650390625]
- n_outliers: 0

[viclafargue]

Here is a script demonstrating that most outliers seem to stem from cuML's simplicial_set_embedding function. The script runs every versions of estimators/functions with explicitly specified constant parameters to account for potentially different defaults. This points to possible discrepancies in the iterative embedding optimization algorithm.

from sklearn.datasets import make_blobs
import numpy as np
import umap.distances as dist
from cuml.manifold.umap import UMAP
from umap import UMAP as refUMAP

from umap.umap_ import fuzzy_simplicial_set as ref_fuzzy_simplicial_set
from umap.umap_ import simplicial_set_embedding as ref_simplicial_set_embedding
from cuml.manifold.simpl_set import fuzzy_simplicial_set, simplicial_set_embedding


def fit_and_evaluate(label, cuml_embeddings, ref_embeddings):

    print(f"{label}:")
    print(f"- min: {cuml_embeddings.min(axis=0).tolist()}")
    print(f"- max: {cuml_embeddings.max(axis=0).tolist()}")
    # semi-arbitrarily define outliers as points outside 2 stds of cpu bounds
    cpu_min = ref_embeddings.min(axis=0)
    cpu_max = ref_embeddings.max(axis=0)
    cpu_std = ref_embeddings.std(axis=0)
    lower_bound = cpu_min - (2 * cpu_std)
    upper_bound = cpu_max + (2 * cpu_std)
    n_outliers = (
        ((cuml_embeddings < lower_bound) | (cuml_embeddings > upper_bound))
        .any(axis=1)
        .sum()
    )
    print(f"- n_outliers: {n_outliers}")

X, _ = make_blobs(centers=10, n_features=512, n_samples=300_000)

n_components = 2
n_neighbors = 15
random_state = 42
metric = "euclidean"
initial_alpha = 1.0
a, b = UMAP.find_ab_params(1.0, 0.1)
gamma = 1.0
negative_sample_rate = 5
n_epochs = 500
init = "random"
metric_kwds = {}
densmap = False
densmap_kwds = {}
output_dens = False
output_metric = "euclidean"
output_metric_kwds = {}

def run_UMAP(library):
    if library == "umap-learn":
        estimator_class = refUMAP
    elif library == "cuML":
        estimator_class = UMAP
    else:
        raise ValueError("Invalid library")
    model = estimator_class(n_neighbors=n_neighbors,
                            n_components=n_components,
                            metric=metric,
                            n_epochs=n_epochs,
                            init=init,
                            learning_rate=initial_alpha,
                            a=a,
                            b=b,
                            repulsion_strength=gamma,
                            negative_sample_rate=negative_sample_rate,
                            random_state=random_state)
    embeddings = model.fit_transform(X)
    return embeddings

def run_fuzzy_simplicial_set(library):
    if library == "umap-learn":
        func = ref_fuzzy_simplicial_set
    elif library == "cuML":
        func = fuzzy_simplicial_set
    else:
        raise ValueError("Invalid library")
    result = func(X, n_neighbors, random_state, metric)
    return result.get() if library == "cuML" else result[0]

def run_simplicial_set_embedding(library, graph):
    if library == "umap-learn":
        func = ref_simplicial_set_embedding
        additional_kwargs = {
            'densmap': densmap,
            'densmap_kwds': densmap_kwds,
            'output_dens': output_dens
        }
        metric_val = dist.named_distances_with_gradients[metric]
    elif library == "cuML":
        func = simplicial_set_embedding
        additional_kwargs = {}
        metric_val = metric
    else:
        raise ValueError("Invalid library")

    result = func(
        X,
        graph,
        n_components=n_components,
        initial_alpha=initial_alpha,
        a=a,
        b=b,
        gamma=gamma,
        negative_sample_rate=negative_sample_rate,
        n_epochs=n_epochs,
        init=init,
        random_state=np.random.RandomState(random_state),
        metric=metric_val,
        metric_kwds=metric_kwds,
        output_metric=output_metric,
        output_metric_kwds=output_metric_kwds,
        **additional_kwargs
    )
    return result.to_output("numpy") if library == "cuML" else result[0]


ref_embeddings = run_UMAP("umap-learn")
cuml_embeddings = run_UMAP("cuML")

ref_fss_graph = run_fuzzy_simplicial_set("umap-learn")
cu_fss_graph = run_fuzzy_simplicial_set("cuML")

cpu_fss_cpu_embedding = run_simplicial_set_embedding("umap-learn", ref_fss_graph)
cpu_fss_gpu_embedding = run_simplicial_set_embedding("cuML", ref_fss_graph)

gpu_fss_cpu_embedding = run_simplicial_set_embedding("umap-learn", cu_fss_graph)
gpu_fss_gpu_embedding = run_simplicial_set_embedding("cuML", cu_fss_graph)

fit_and_evaluate("cuML", cuml_embeddings, ref_embeddings)
fit_and_evaluate("CPU fss - CPU embedding", cpu_fss_cpu_embedding, ref_embeddings)
fit_and_evaluate("GPU fss - CPU embedding", gpu_fss_cpu_embedding, ref_embeddings)
fit_and_evaluate("CPU fss - GPU embedding", cpu_fss_gpu_embedding, ref_embeddings)
fit_and_evaluate("GPU fss - GPU embedding", gpu_fss_gpu_embedding, ref_embeddings)

[jcrist]

Could #6539 be responsible (or partly responsible) here?

[viclafargue]

Thanks for investigating this @jcrist. On my side, I looked into the specific case that generates many outliers, namely the cuML embeddings generated from a umap-learn fuzzy simplicial set. After printing almost every values and comparing with the reference implementation, there wasn't anything truly significant. However, after serializing a fuzzy simplicial set and generating many cuML embeddings with it I noticed that the outliers are very often the same samples (even when changing seed) and these are the samples that have the most connections. It also looks like the umap-learn fuzzy simplicial set is even more imbalanced than the cuML one. Whereas the average sample is connected bi-directionally to roughly 25 others, some samples are connected to thousands.

I believe that the explanation is that during gradient update these highly connected samples accumulate much more gradients and the sum of them causes sudden moves that are more related to the distribution of connected points than to the position of the point being updated itself. I could not yet pinpoint the exact difference between implementations though. It looks like umap-learn implementation updates the positions of the samples directly in place and I believe also do the updates in parallel with Numba CPU.

[viclafargue]

After spending a some amount of time investigating potential inconsistencies in implementations, I’ve explored a few things and have some updates on the issue:

Enabling deterministic behavior stores a buffer of updates, which makes it easier to identify the largest embedding updates. I observed some extremely large updates—on the order of hundreds. I then experimented with disabling deterministic mode (i.e., allowing in-place modifications) and running the simplicial set embedding optimization in batches of 1000, rather than processing all updates in parallel. This significantly reduced the presence of outliers.

It seems the type of embeddings produced by the reference library is influenced by Numba's limited CPU parallelism and its in-place update behavior. As embeddings evolve, the updates naturally shrink, which constrains the overall embedding range.

In terms of addressing these outliers, we obviously don't want to give up parallelism. However, it appears that artificially reducing the learning rate (or initial_alpha value) helps achieve more consistent results. Fine-tuning the learning rate might allow us to better align the embeddings with those produced by umap-learn.