Parametric UMAP embeddings for representation and semi-supervised learning

Tim Sainburg; Leland McInnes; Timothy Q Gentner

Source-linked AI summary

Parametric UMAP embeddings for representation and semi-supervised learning

Tim Sainburg, Leland McInnes, Timothy Q Gentner

arXiv:2009.12981v4Sep 27, 2020cs.LGcs.CGq-bio.QMstat.ML

TL;DR

The paper addresses the difficulty of optimizing t-SNE with neural networks because normalization requires the entire dataset at each optimization step. It introduces Parametric UMAP by bringing UMAP’s graph-based dimensionality reduction into a network-based embedding and compares it with parametric and non-parametric algorithms.

Problem
t-SNE optimization requires normalization over the entire dataset at each step, making it difficult to use with neural-network training and batch-wise learning.
Method
Parametric UMAP brings UMAP’s non-parametric graph-based dimensionality reduction into a network-based embedding.
Results
The paper reports experiments comparing Parametric UMAP with parametric and non-parametric algorithms.
Takeaways & Limitations
The paper explores UMAP as a regularizer for neural networks, including autoencoders that capture more global structure alongside local structure.
Takeaways & Limitations
The paper identifies t-SNE’s loss function as poorly suited to neural-network training paradigms, particularly because its optimization requires whole-dataset normalization.

Abstract

from arXiv · show

UMAP is a non-parametric graph-based dimensionality reduction algorithm using applied Riemannian geometry and algebraic topology to find low-dimensional embeddings of structured data. The UMAP algorithm consists of two steps: (1) Compute a graphical representation of a dataset (fuzzy simplicial complex), and (2) Through stochastic gradient descent, optimize a low-dimensional embedding of the graph. Here, we extend the second step of UMAP to a parametric optimization over neural network weights, learning a parametric relationship between data and embedding. We first demonstrate that Parametric UMAP performs comparably to its non-parametric counterpart while conferring the benefit of a learned parametric mapping (e.g. fast online embeddings for new data). We then explore UMAP as a regularization, constraining the latent distribution of autoencoders, parametrically varying global structure preservation, and improving classifier accuracy for semi-supervised learning by capturing structure in unlabeled data. Google Colab walkthrough: https://colab.research.google.com/drive/1WkXVZ5pnMrm17m0YgmtoNjM_XHdnE5Vp?usp=sharing

1 Introduction

Parametric UMAP extends UMAP’s graph-based dimensionality reduction by optimizing a neural network over batches instead of directly optimizing point embeddings. The paper evaluates this parametric mapping against existing methods and extends it to semi-supervised learning.

UMAP’s negative-sampling objective is more compatible with batch-wise training than t-SNE’s whole-dataset normalization.
Parametric UMAP brings UMAP into a class of parametric, topologically inspired embedding algorithms.
The paper compares Parametric UMAP with parametric and non-parametric algorithms.
The paper also introduces a semi-supervised-learning extension of Parametric UMAP.

2 Parametric and non-parametric UMAP

UMAP and t-SNE construct neighborhood graphs and optimize low-dimensional embeddings, but UMAP uses unnormalized probabilities, cross-entropy, and sparse negative sampling. Parametric UMAP applies this objective to neural-network weights, enabling batch-wise training and related extensions.

Both algorithms construct a graph of local relationships and optimize a low-dimensional embedding that preserves its structure.
Parametric UMAP optimizes deep-neural-network parameters over batches rather than directly optimizing one embedding per data point.
Graph Construction: UMAP computes local edge probabilities from nearest-neighbor distances under a manifold-uniformity assumption, then combines directional probabilities globally.
Graph Embedding: Unlike t-SNE’s normalized embedding distribution, UMAP uses an unnormalized embedding-space probability distribution and optimizes cross-entropy between input and embedding relationships.
Attraction and repulsion: UMAP reduces computation through nearest-neighbor attraction and negative sampling for repulsion, treating sampled non-neighbor edges as zero-probability relationships.
Parametric optimization: UMAP’s sparse optimization supports batch-wise training, whereas t-SNE requires computing the full probability distribution before each optimization step.
Parametric optimization: Parametric UMAP applies its loss as a regularizer in standard stochastic-gradient-descent deep-learning settings.

3 Related Work

Related work includes parametric dimensionality-reduction methods, structure-preserving autoencoders, and neural architectures that impose geometric or topological constraints. The paper distinguishes direct optimization of an embedding loss from indirect conditioning on non-parametric embeddings.

Recent parametric dimensionality-reduction methods use neural networks and structure-preserving constraints to learn embeddings.
Several approaches regularize autoencoder latent spaces with local structure, diffusion maps, PHATE, or persistent homology.
Other methods use variational-autoencoder objectives or Siamese architectures with structure-preserving losses.
Some parametric approaches condition neural networks on non-parametric embeddings rather than directly optimizing the embedding algorithm’s loss.
The paper contrasts indirect and direct parametric embeddings.

4 UMAP as a regularization

The paper frames UMAP as a regularizer for neural networks, applying its loss through an encoder and combining it with additional objectives such as reconstruction or classification. This supports hybrid models that use UMAP to regularize training and constrain learned embeddings.

4 UMAP as a regularization: UMAP loss can regularize neural networks, while additional losses can regularize the embedding learned by UMAP.The paper presents both directions of regularization as complementary training objectives.
4 UMAP as a regularization: Parametric UMAP applies the UMAP cost function over an encoder network rather than optimizing embeddings directly.This contrasts the non-parametric formulation with the proposed parametric formulation.
4 UMAP as a regularization: A UMAP/AE hybrid combines UMAP loss with reconstruction loss over the network to add autoencoder-based regularization.The paper motivates this combination as a way to capture latent structure.
4 UMAP as a regularization: Combining a locally focused embedding objective with autoencoders aids in capturing more global structure over the dataset.The paper connects this benefit to prior work combining t-SNE with autoencoders.
4 UMAP as a regularization: In semi-supervised learning, classifier loss is applied to labeled data while UMAP loss is applied to unlabeled data in a shared network.The setup jointly learns dataset structure and a supervised objective such as image classification.

A B Embedding Decision Contour

The paper uses UMAP to address the tension between local and global structure preservation and illustrates semi-supervised learning on the moons dataset. Joint classifier and UMAP training separates the moons and yields a clean decision boundary.

A B Embedding Decision Contour: Figure 3 presents an example of semi-supervised learning with UMAP on the moons dataset.The figure caption identifies the dataset and learning setting but does not specify additional visual encodings.
A B Embedding Decision Contour: The moons example trains a Y-shaped network jointly on classifier loss for labeled points and UMAP loss for unlabeled points.The shared latent space links the classifier and UMAP branches.
A B Embedding Decision Contour: The shared latent space pulls apart the two moons, resulting in a decision boundary that divides cleanly between the two distributions.This is the reported outcome of the semi-supervised moons example.
A B Embedding Decision Contour: The paper identifies balancing local and global structure as an open issue while noting that global structure remains desirable.Parametric UMAP addresses this by jointly training with a direct global structure-preservation loss.
A B Embedding Decision Contour: UMAP focuses on local nearest-neighbor structure, whereas global methods such as MDS attempt to preserve all relationships during embedding.Local methods are described as more computationally efficient and able to retain application-relevant cluster structure.

5 Experiments

Experiments compare Parametric UMAP with non-parametric UMAP and several parametric and non-parametric baselines across benchmark and real-world datasets. Parametric UMAP achieves broadly comparable embedding quality, enables faster inference after training, supports global-structure control and reconstruction, and can improve semi-supervised classification when unlabeled data are augmented.

Embedding quality: UMAP-based methods generally preserve clustering better than PCA, AE, and VAE projections, while UMAP and t-SNE variants are broadly similar in 2D.In 64D, UMAP approaches show superior clustering in the reported comparisons, whereas Parametric t-SNE can approach or fall below PCA on silhouette.
Training and embedding speed: Parametric UMAP reaches similar cross-entropy within an order of magnitude of non-parametric UMAP training time, although non-parametric UMAP lowers loss more quickly.The reported training-time comparison makes Parametric UMAP a reasonable alternative in terms of training time.
Training and embedding speed: Parametric mappings make inference faster because embeddings are direct projections through learned networks.Parametric methods are several orders of magnitude faster than non-parametric t-SNE and UMAP for embedding new data, and Parametric UMAP is slightly slower than PCA within the same order of magnitude.
Global structure: Increasing the Pearson-correlation loss weight increases captured global structure while a small weight preserves much of the local structure.The experiments describe a tradeoff between global and local structure as the weight of C_Pearson changes.
Representation learning: Parametric UMAP generally reconstructs better than non-parametric UMAP, and autoencoder regularization slightly improves reconstruction.At 64 dimensions, autoencoder-regularized Parametric UMAP is generally comparable to AE and VAE and better than unregularized Parametric UMAP.
Semi-supervised learning: Semi-supervised classification improves on MNIST, FMNIST, birdsong, and retina when UMAP loss uses augmented data, including CIFAR10.Without augmentation, gains are little to none in the three evaluated datasets; with augmentation, UMAP loss improves performance in each dataset.
Semi-supervised learning: The learned graph largely conforms to existing categorical decisions, so UMAP loss alone adds little information without augmentation.The paper attributes the additional benefit with unlabeled data to regularizing the classifier rather than learning structure already intrinsic to its representations.

P. UMAP MSE

The section presents a one-hidden-layer Parametric UMAP experiment on MNIST, comparing cross-entropy losses and corresponding projections as network width varies.

P. UMAP MSE: Cross-entropy loss and corresponding MNIST projections are compared as the number of neurons in a single hidden layer varies.The dashed grey line denotes the loss for the non-parametric embedding.
P. UMAP MSE: The experiment uses the same network architectures for Parametric UMAP and a network trained to predict non-parametric embeddings with MSE.The x-axis varies the hidden-layer width.
P. UMAP MSE: The figure pairs the loss curves above with the projections produced at the corresponding settings.The lower panels show projections corresponding to the losses in the upper panel.

6 Discussion

Parametric UMAP extends UMAP with a learned mapping between data and embedding spaces, retaining similar embedding quality while enabling faster inference and downstream applications. The discussion also describes improvements in global-structure preservation, autoencoder reconstruction, and semi-supervised classification, with benefits depending on whether categorically relevant structure is present.

Parametric UMAP produces embeddings similar to non-parametric UMAP while adding a learned mapping between data and embedding spaces.The learned mapping supports downstream uses beyond the original embedding optimization.
Inference times for embeddings and reconstructions improve by orders of magnitude while maintaining similar embedding quality to non-parametric UMAP.This speed advantage follows from using learned parametric relationships for inference.
Global structure preservation improves when Parametric UMAP jointly optimizes a global-structure loss rather than imposing global structure only at initialization.The discussion identifies Pearson correlation and broader notions of global relationships as possible components of this objective.
Autoencoder regularization improves reconstruction quality and enables reconstruction of high-dimensional UMAP projections.The autoencoder approach is described as supporting more global structure while UMAP emphasizes local structure.
UMAP loss improves semi-supervised classifier accuracy when unlabeled data contain categorically relevant structure, including stereotyped birdsongs and single-cell transcriptomes.The reported benefit does not extend to cases such as CIFAR10 where categorically relevant structure is not present.
Parametric embedding makes UMAP feasible for continuously generated signals used in real-time analysis and experimental control.The discussion presents this as a downstream implication of the learned mapping and faster inference.

8 Appendix

The appendix describes datasets, model configurations, and evaluation procedures for comparing parametric and non-parametric latent projections. It evaluates neighborhood preservation, supervised neighborhood classification, clustering, and cluster separation.

Datasets: Experiments span datasets including images, single-cell transcriptomes, and stereotyped birdsong, with varying complexity and categorical structure.CIFAR10 is described as less structured than the other datasets, while Cassin’s vireo song is relatively low noise with highly stereotyped vocal elements.
Model configurations: Parametric networks use dataset-specific architectures, including convolutional networks, an LSTM encoder, and a three-layer MLP.MNIST, FMNIST, and CIFAR10 use convolutional networks; Cassin’s vireo uses an LSTM, and Retina uses a three-layer MLP with 100 neurons per layer.
Model configurations: Parametric UMAP and its autoencoder variant train encoder networks with UMAP loss, with the autoencoder optionally adding reconstruction loss.The encoder and decoder may be jointly trained on reconstruction and UMAP losses, while a separate decoder is also trained for an encoder using only UMAP loss.
Evaluation metrics: Trustworthiness compares nearest neighbors between original and projected spaces, while AUC RNX summarizes neighborhood preservation across scales.AUC RNX weights errors in large neighborhoods less than errors in small neighborhoods.
Evaluation metrics: Latent clustering is evaluated with k-Means across varying cluster counts, selecting the best result by silhouette score and comparing normalized mutual information.The selected clustering is evaluated without reference to labels when choosing by silhouette score, then compared using NMI between labels and clusters.

9 Figures

The figures and tables report projection quality, classification, clustering, reconstruction, and classification-accuracy comparisons across datasets and methods. Some visualizations have explicit scope limitations, including unavailable high-dimensional t-SNE results and omitted reconstructions.

Projection and latent-space metrics: Figures 16–20 compare trustworthiness, AUC RNX, KNN generalization error, and silhouette scores across five datasets and projection dimensions.The trustworthiness and silhouette figures explicitly compare 2- and 64-dimensional projections.
Scope limitations: 64-dimensional t-SNE results are omitted because of limitations in high-dimensional t-SNE projections.This is an explicit scope limitation of the displayed comparison rather than a reported performance result.
Clustering: Figure 21 compares clustering results using normalized mutual information between labels and clusters.For each dataset, the reported NMI corresponds to the clustering selected by silhouette score.
Reconstruction: Figure 22 compares reconstruction speed, with reconstructions performed on the same machine.The accompanying description states that values are median times over 10 runs and that comparable neural-network architectures keep parametric implementation speeds equal across methods.
Scope limitations: Reconstructions are omitted for the Retina dataset with PCA using a 64-dimensional projection.The omission is stated specifically for the reconstruction comparison.
Tables: The appendix tables provide method-level values for trustworthiness, AUC RMX, KNN scores, silhouette, clustering, reconstruction error, and classification accuracy.The tables correspond to the reported figures and include classification accuracy for different numbers of labeled training examples.