[Paper Review] From $t$-SNE to UMAP with contrastive learning

Sebastian Damrich, Jan Niklas Böhm+2 more|arXiv (Cornell University)|Jun 3, 2022

Data Visualization and Analytics4 citations

TL;DR

This paper establishes a conceptual link between t-SNE and UMAP through contrastive learning, showing that t-SNE uses noise-contrastive estimation (NCE) while UMAP uses negative sampling (NEG), which introduces a distortion leading to tighter, more compact clusters. The authors propose a unified spectrum of embeddings by interpolating between these methods, enabling controlled trade-offs between local and global structure and improving robustness against over-interpretation of single-embedding artifacts.

ABSTRACT

Neighbor embedding methods $t$-SNE and UMAP are the de facto standard for visualizing high-dimensional datasets. Motivated from entirely different viewpoints, their loss functions appear to be unrelated. In practice, they yield strongly differing embeddings and can suggest conflicting interpretations of the same data. The fundamental reasons for this and, more generally, the exact relationship between $t$-SNE and UMAP have remained unclear. In this work, we uncover their conceptual connection via a new insight into contrastive learning methods. Noise-contrastive estimation can be used to optimize $t$-SNE, while UMAP relies on negative sampling, another contrastive method. We find the precise relationship between these two contrastive methods and provide a mathematical characterization of the distortion introduced by negative sampling. Visually, this distortion results in UMAP generating more compact embeddings with tighter clusters compared to $t$-SNE. We exploit this new conceptual connection to propose and implement a generalization of negative sampling, allowing us to interpolate between (and even extrapolate beyond) $t$-SNE and UMAP and their respective embeddings. Moving along this spectrum of embeddings leads to a trade-off between discrete / local and continuous / global structures, mitigating the risk of over-interpreting ostensible features of any single embedding. We provide a PyTorch implementation.

Motivation & Objective

• To clarify the conceptual relationship between t-SNE and UMAP, which have long been considered unrelated despite their shared use in dimensionality reduction.
• To identify the mathematical origin of the structural differences between t-SNE and UMAP by analyzing their underlying contrastive learning methods: NCE and NEG.
• To develop a generalized embedding framework that interpolates between t-SNE and UMAP, enabling controlled exploration of local versus global data structure.
• To address an instability in UMAP by identifying and correcting the distortion introduced by negative sampling.
• To provide a unified, open-source PyTorch implementation for contrastive neighbor embeddings.

Proposed method

• Reformulate t-SNE as noise-contrastive estimation (NCE), enabling direct comparison with UMAP’s negative sampling (NEG) approach.
• Derive the exact mathematical relationship between NCE and NEG, showing that NEG introduces a systematic distortion in the similarity landscape.
• Introduce a generalized embedding framework using a tunable normalization constant $\bar{Z}$ to interpolate between t-SNE ($\bar{Z} = Z^{t\text{-SNE}}$) and UMAP ($\bar{Z} = |X|/m$).
• Implement a sampling-based optimization method that generalizes both t-SNE and UMAP, allowing for smooth transitions across the spectrum of embeddings.
• Apply early exaggeration-like strategies in the new framework to stabilize training, especially for low $\bar{Z}$ values.
• Provide a PyTorch-based implementation with support for both non-parametric and parametric variants, including NC- and InfoNC-t-SNE.

Experimental results

Research questions

• RQ1 What is the fundamental conceptual relationship between t-SNE and UMAP, given their distinct loss functions and visual outputs?
• RQ2 How does negative sampling in UMAP differ mathematically from noise-contrastive estimation in t-SNE, and what is the resulting distortion?
• RQ3 Can a continuous spectrum of embeddings be constructed that interpolates between t-SNE and UMAP, enabling controlled exploration of local and global structure?
• RQ4 Does the proposed framework mitigate instability issues in UMAP, particularly in cluster separation and compactness?
• RQ5 How does the new framework compare to existing contrastive methods in preserving neighborhood structure and global topology?

Key findings

• The distortion introduced by negative sampling in UMAP causes more compact, tightly separated clusters compared to t-SNE, which uses noise-contrastive estimation.
• When $\bar{Z} = Z^{t\text{-SNE}}$, the generalized embedding closely approximates t-SNE; when $\bar{Z} = |X|/m$, it matches UMAP’s behavior.
• The NC- and InfoNC-t-SNE implementations show improved stability and reduced dependence on early exaggeration compared to standard t-SNE.
• Longer optimization times significantly improve the quality of Neg-t-SNE embeddings, with KL divergence decreasing from 5.52 to 4.65 over 10,000 epochs on the Kuzushiji-49 dataset.
• Higher noise sample counts ($m$) in NC- and InfoNC-t-SNE improve k-NN recall without degrading Spearman correlation, indicating better preservation of local structure.
• The proposed spectrum of embeddings allows users to systematically explore trade-offs between discrete, local structure and continuous, global structure, reducing the risk of over-interpreting individual visualizations.