UMAP#

class torchdr.UMAP(n_neighbors: int = 30, n_components: int = 2, min_dist: float = 0.1, spread: float = 1.0, a: float | None = None, b: float | None = None, lr: float = 1.0, optimizer: str | Type[Optimizer] = 'SGD', optimizer_kwargs: Dict | str | None = None, scheduler: str | Type[LRScheduler] | None = 'LinearLR', scheduler_kwargs: Dict | str | None = 'auto', init: str = 'pca', init_scaling: float = 0.0001, min_grad_norm: float = 1e-07, max_iter: int = 1000, device: str = 'auto', backend: str | FaissConfig | None = 'faiss', verbose: bool = False, random_state: float | None = None, max_iter_affinity: int = 100, metric: str = 'sqeuclidean', negative_sample_rate: int = 5, check_interval: int = 50, exclude_neighbors_from_negative_sampling: bool | None = None, discard_NNs: bool | None = None, compile: bool = False, distributed: bool | str = 'auto', input_layout: str = 'replicated', **kwargs)[source]#

Bases: NegativeSamplingNeighborEmbedding

UMAP introduced in [McInnes et al., 2018] and further studied in [Damrich and Hamprecht, 2021].

It uses a UMAPAffinity as input affinity \(\mathbf{P}\) and output affinity \(Q_{ij} = (1 + a \| \mathbf{z}_i - \mathbf{z}_j \|^{2b})^{-1}\) where \(a, b\) are fitted from min_dist and spread.

The loss function is defined as:

\[-\sum_{ij} P_{ij} \log Q_{ij} + \sum_{i,j \in \mathrm{Neg}(i)} \log (1 - Q_{ij})\]

where \(\mathrm{Neg}(i)\) is the set of negatives samples for point \(i\).

Note

This implementation supports multi-GPU training when launched with torchrun. Set distributed='auto' (default) to automatically detect and use multiple GPUs. It also supports the shared non-parametric transform path implemented in NegativeSamplingNeighborEmbedding.

Parameters:

  • n_neighbors (int, optional) – UMAP neighbor count. During fit this includes the sample itself, so the fuzzy graph uses n_neighbors - 1 other samples, matching umap-learn. Transform is bipartite and therefore uses n_neighbors training samples because a new query has no self-neighbor in the training set.

  • n_components (int, optional) – Dimension of the embedding space.

  • min_dist (float, optional) – Minimum distance between points in the embedding space.

  • spread (float, optional) – The effective scale of the embedded points. Used to configure the UMAPAffinityOut.

  • a (float, optional) – Parameter for the Student t-distribution.

  • b (float, optional) – Parameter for the Student t-distribution.

  • lr (float, optional) – Initial learning rate for the embedding optimization, by default 1.0. This matches the numerical default in umap-learn.

  • optimizer (str or torch.optim.Optimizer, optional) – Name of an optimizer from torch.optim or an optimizer class. Default is “SGD”.

  • optimizer_kwargs (dict, optional) – Additional keyword arguments for the optimizer. Default is None.

  • scheduler (str or torch.optim.lr_scheduler.LRScheduler, optional) – Name of a scheduler from torch.optim.lr_scheduler or a scheduler class. Default is “LinearLR”.

  • scheduler_kwargs (dict, 'auto', or None, optional) – Additional keyword arguments for the scheduler. Default is ‘auto’, which corresponds to a linear decay from the learning rate to 0 for LinearLR.

  • init ({'normal', 'pca'} or torch.Tensor of shape (n_samples, output_dim), optional) – Initialization for the embedding Z, default ‘pca’.

  • init_scaling (float, optional) – Scaling factor for the initialization, by default 1e-4.

  • min_grad_norm (float, optional) – Precision threshold at which the algorithm stops, by default 1e-7.

  • max_iter (int, optional) – Number of maximum iterations for the descent algorithm. by default 2000.

  • device (str, optional) – Device to use, by default “auto”.

  • backend ({"keops", "faiss", None} or FaissConfig, optional) – Which backend to use for handling sparsity and memory efficiency. Can be: - “keops”: Use KeOps for memory-efficient symbolic computations - “faiss”: Use FAISS for fast k-NN computations with default settings - None: Use standard PyTorch operations - FaissConfig object: Use FAISS with custom configuration Default is “faiss”.

  • verbose (bool, optional) – Verbosity, by default False.

  • random_state (float, optional) – Random seed for reproducibility, by default None.

  • max_iter_affinity (int, optional) – Number of maximum iterations for the input affinity computation.

  • metric ({'euclidean', 'manhattan'}, optional) – Metric to use for the input affinity, by default ‘sqeuclidean’.

  • n_negatives (int, optional) – Number of negative samples for the noise-contrastive loss, by default 10.

  • check_interval (int, optional) – Check interval for the algorithm, by default 50.

  • exclude_neighbors_from_negative_sampling (bool, optional) – Whether to exclude nearest neighbors from negative sampling. Default is False.

  • discard_NNs (bool, optional) – Deprecated alias for exclude_neighbors_from_negative_sampling.

  • compile (bool, optional) – Whether to compile the algorithm using torch.compile. Default is False.

  • distributed (bool or 'auto', optional) – Whether to use distributed computation across multiple GPUs. - “auto”: Automatically detect if running with torchrun (default) - True: Force distributed mode (requires torchrun) - False: Disable distributed mode Default is “auto”.

  • input_layout ({'replicated', 'sharded'}, optional) –

    How the input rows are laid out across distributed ranks. - “replicated” (default): every rank holds the full input. - “sharded”: each rank holds a distinct contiguous shard whose rows

    concatenate, in rank order, into the global dataset, so no rank ever materializes the whole input. Requires an exact Flat FAISS backend and init in {“random”, “normal”, “hyperbolic”, “pca”}; the embedding stays replicated (one coordinate per global point on every rank). A high-level FAISS plan must set distribution="shard".

    Default is “replicated”.

Notes

The default learning rate and its linear decay match umap-learn’s learning_rate and initial_alpha schedule, but the resulting updates are not step-for-step equivalent. umap-learn applies sampled-edge updates sequentially with its own optimizer, whereas TorchDR accumulates vectorized gradients and applies them simultaneously with PyTorch SGD. For fit-time attractive updates, the symmetric graph lets TorchDR recover the contribution of moving both endpoints by scaling the reduced gradient; this is intentionally disabled during transform because training points are frozen.

on_affinity_computation_end()[source]#

Build per-row exclusion indices for negative sampling.

on_training_step_start()[source]#

Defer negative sampling until the current positive mask is known.

Examples using UMAP:#

Neighbor Embedding on genomics & equivalent affinity matcher formulation

Neighbor Embedding on genomics & equivalent affinity matcher formulation

UMAP Non-Parametric Transform on Handwritten Digits

UMAP Non-Parametric Transform on Handwritten Digits