NegativeSamplingNeighborEmbedding#

class torchdr.NegativeSamplingNeighborEmbedding(affinity_in: Affinity, affinity_out: Affinity | None = None, kwargs_affinity_out: Dict | None = None, n_components: int = 2, lr: float | str = 1.0, optimizer: str | Type[Optimizer] = 'SGD', optimizer_kwargs: Dict | str = 'auto', scheduler: str | Type[LRScheduler] | None = None, scheduler_kwargs: Dict | str | None = 'auto', min_grad_norm: float = 1e-07, max_iter: int = 2000, init: str = 'pca', init_scaling: float = 0.0001, device: str = 'auto', backend: str | FaissConfig | None = None, verbose: bool = False, random_state: float | None = None, early_exaggeration_coeff: float = 1.0, early_exaggeration_iter: int | None = None, repulsion_strength: float = 1.0, n_negatives: int = 5, check_interval: int = 50, exclude_neighbors_from_negative_sampling: bool | None = None, discard_NNs: bool | None = None, compile: bool = False, **kwargs)[source]#

Bases: NeighborEmbedding

Neighbor embedding that approximates the repulsive term via negative sampling.

This class extends NeighborEmbedding for methods that avoid the \(O(n^2)\) cost of computing the repulsive term over all point pairs. Instead, a fixed number of negative samples (n_negatives) are drawn uniformly per point at each iteration, reducing the repulsive cost to \(O(n)\).

Negative sampling details:

  • At each iteration, n_negatives indices are sampled uniformly (excluding the point itself) for each point in the local chunk.

  • When exclude_neighbors_from_negative_sampling is True, nearest neighbors are also excluded from the negative samples to avoid conflicting gradients.

  • The sampled indices are stored in neg_indices_ and refreshed every iteration via on_training_step_start().

Non-parametric transform support:

This family also provides the shared machinery for out-of-sample non-parametric transform. The base implementation handles the generic transform lifecycle:

  • find nearest neighbors from new points to the reference training set;

  • build a bipartite affinity from new points to training points;

  • initialize new embeddings from that bipartite graph;

  • optimize only the new points while keeping the fitted training embedding frozen.

The design is intentionally split so the base class owns the generic transform pipeline, while each algorithm only provides the method-specific bipartite affinity through _compute_bipartite_affinity(). This keeps the out-of-sample logic centralized and prevents each subclass from reimplementing the same transform scaffolding.

Inherits distributed multi-GPU support from NeighborEmbedding.

Subclasses must implement _compute_attractive_loss() and _compute_repulsive_loss() (or the gradient equivalents). Subclasses that support non-parametric transform must additionally implement _compute_bipartite_affinity().

Direct subclasses: UMAP, LargeVis, InfoTSNE, PACMAP.

Parameters:

  • affinity_in (Affinity) – The affinity object for the input space.

  • affinity_out (Affinity, optional) – The affinity object for the output embedding space. Default is None.

  • kwargs_affinity_out (dict, optional) – Additional keyword arguments for the affinity_out method.

  • n_components (int, optional) – Number of dimensions for the embedding. Default is 2.

  • lr (float or 'auto', optional) – Learning rate for the optimizer. Default is 1e0.

  • optimizer (str or torch.optim.Optimizer, optional) – Name of an optimizer from torch.optim or an optimizer class. Default is “SGD”. For best results, we recommend using “SGD” with ‘auto’ learning rate.

  • optimizer_kwargs (dict or 'auto', optional) – Additional keyword arguments for the optimizer. Default is ‘auto’, which sets appropriate momentum values for SGD based on early exaggeration phase.

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

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

  • min_grad_norm (float, optional) – Tolerance for stopping criterion. Default is 1e-7.

  • max_iter (int, optional) – Maximum number of iterations. Default is 2000.

  • init (str, optional) – Initialization method for the embedding. Default is “pca”.

  • init_scaling (float, optional) – Scaling factor for the initial embedding. Default is 1e-4.

  • device (str, optional) – Device to use for computations. Default is “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 None.

  • verbose (bool, optional) – Verbosity of the optimization process. Default is False.

  • random_state (float, optional) – Random seed for reproducibility. Default is None.

  • early_exaggeration_coeff (float, optional) – Coefficient for the attraction term during the early exaggeration phase. Default is 1.0.

  • early_exaggeration_iter (int, optional) – Number of iterations for early exaggeration. Default is None.

  • repulsion_strength (float, optional) – Strength of the repulsive term. Default is 1.0.

  • n_negatives (int, optional) – Number of negative samples to use. Default is 5.

  • check_interval (int, optional) – Number of iterations between two checks for convergence. Default is 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, default=False) – Whether to use torch.compile for faster computation.

  • **kwargs – All other parameters (including distributed) are forwarded to NeighborEmbedding.

on_affinity_computation_end()[source]#

Build per-row exclusion indices for negative sampling.

on_training_step_start()[source]#

Sample negatives using a unified path for single- and multi-GPU.

Examples using NegativeSamplingNeighborEmbedding:#

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