Source-linked AI summary
Efficient Algorithms for t-distributed Stochastic Neighborhood Embedding
George C. Linderman, Manas Rachh, Jeremy G. Hoskins, Stefan Steinerberger, Yuval Kluger
TL;DR
Large high-dimensional datasets can make existing t-SNE implementations impractical because their computation remains costly and the data may exceed available memory. The paper introduces FIt-SNE, combining FFT-based interpolation and approximate nearest neighbors, alongside late exaggeration and out-of-core PCA. These methods support million-point embeddings, clearer cluster separation, and PCA without fully loading the dataset.
Problem
Existing approximate t-SNE methods can require many hours for datasets with hundreds of thousands to millions of high-dimensional points, while some datasets cannot fit in memory.
Method
FIt-SNE accelerates repulsive-force computation with interpolation and FFT convolution, uses approximate nearest neighbors for input similarities, adds late exaggeration, and provides out-of-core randomized PCA.
Results
For 1 million points, FIt-SNE is 15 times faster than Barnes-Hut in 1D and 30 times faster in 2D, while oocPCA computes 50 principal components in 16 minutes with approximately 10^-9 spectral-norm accuracy using 1GB of memory.
Takeaways & Limitations
The methods allow visualization of datasets with millions of points and datasets too large for memory, including without subsampling on a standard personal computer.
Takeaways & Limitations
Higher accuracy from increasing the interpolation order requires care because equispaced nodes can produce the Runge phenomenon; the implementation uses low-order piecewise interpolation.
Abstract
from arXiv · showhide
t-distributed Stochastic Neighborhood Embedding (t-SNE) is a method for dimensionality reduction and visualization that has become widely popular in recent years. Efficient implementations of t-SNE are available, but they scale poorly to datasets with hundreds of thousands to millions of high dimensional data-points. We present Fast Fourier Transform-accelerated Interpolation-based t-SNE (FIt-SNE), which dramatically accelerates the computation of t-SNE. The most time-consuming step of t-SNE is a convolution that we accelerate by interpolating onto an equispaced grid and subsequently using the fast Fourier transform to perform the convolution. We also optimize the computation of input similarities in high dimensions using multi-threaded approximate nearest neighbors. We further present a modification to t-SNE called "late exaggeration," which allows for easier identification of clusters in t-SNE embeddings. Finally, for datasets that cannot be loaded into the memory, we present out-of-core randomized principal component analysis (oocPCA), so that the top principal components of a dataset can be computed without ever fully loading the matrix, hence allowing for t-SNE of large datasets to be computed on resource-limited machines.
1. Introduction
t-SNE is widely used to visualize and discover structure in large, high-dimensional datasets, but existing approximations can still require many hours at million-point scale. The paper introduces faster computation, late exaggeration for clearer clusters, and out-of-core PCA for datasets exceeding memory capacity.
- FIt-SNE: FIt-SNE accelerates t-SNE by using FFT-based interpolation for repulsive-force computation and approximate nearest neighbors for input similarities.The paper organizes these accelerations around the repulsive and attractive components of the gradient.
- Out-of-core analysis: Out-of-core randomized PCA enables visualization of datasets that cannot be loaded entirely into memory.The paper presents this capability for resource-limited machines handling very large high-dimensional datasets.
- t-SNE motivation: t-SNE embeds high-dimensional data so nearby input points remain nearby in a lower-dimensional representation.The method minimizes the Kullback–Leibler divergence between input and embedding distributions using gradient descent.
- Computational challenge: Exact gradient computation scales as O(n^2), while Barnes-Hut approximates repulsive forces with O(n log n) scaling but can still take many hours on large datasets.Input similarities can be precomputed, whereas repulsive forces are approximated at each iteration.
- Late exaggeration: Late exaggeration sets α > 1 during the final iterations, increasing within-cluster attraction so clusters contract and become easier to distinguish.This extends the earlier use of exaggeration during the initial optimization phase.
2. The repulsive forces Frep,i
FIt-SNE accelerates repulsive-force computation by approximating smooth kernels with piecewise polynomial interpolation on equispaced grids and evaluating the resulting Toeplitz interactions with FFTs. This reduces the computation to linear scaling in the number of points while preserving controllable accuracy.
- Motivation: Direct repulsive-force computation evaluates N^2 pairwise interactions, making it prohibitively expensive for datasets with even a few thousand points.The paper focuses on the repulsive component of the t-SNE gradient, which is an N-body simulation.
- Interpolation strategy: FIt-SNE uses low-order polynomial interpolation of the smooth kernels K1 and K2 to replace direct N-body evaluations.The kernels are smooth, and the interpolant is evaluated at equispaced interpolation points.
- Interpolation strategy: The algorithm spreads points onto equispaced grid nodes, computes interval interactions, and evaluates potentials through piecewise Lagrange interpolation.The three stages compute interpolation coefficients, node values, and the potential at each original point.
- FFT acceleration: Translation invariance and equispaced nodes make the interaction matrix Toeplitz, enabling FFT acceleration of the grid-based convolution.This structure replaces the direct interaction cost in the central evaluation step.
- Complexity: In one dimension, Algorithm 1 has complexity O(N · p + (Nint · p) log(Nint · p)), with interpolation parameters determined by the accuracy tolerance rather than N.The method extends naturally to arbitrary embedding dimensions, although constants vary.
- Accuracy and parameter choice: Increasing the interpolation order can trigger the Runge phenomenon, so the implementation uses low-order piecewise interpolation with p = 3.The paper identifies the t-SNE kernels as archetypical examples of this issue.
- Experiments: For 1 million points, FIt-SNE is 15 times faster than Barnes-Hut in 1D and 30 times faster in 2D for computing 1000 gradient iterations.These timings support applying t-SNE to datasets on the order of millions of points.
2 Dimensional Embedding
Figure 3 presents the time required to compute gradients for 1000 iterations of t-SNE, comparing FI t-SNE with Barnes-Hut t-SNE.
- Figure 3 measures the time required to compute gradients for t-SNE.
- The comparison uses 1000 iterations of t-SNE.
- The figure compares FFT-accelerated Interpolation-based t-SNE with the Barnes-Hut t-SNE implementation.
3. The attractive forces Fattr,i
The attractive-force computation connects points that are close in the original space, but exact or high-dimensional nearest-neighbor computations can be expensive. The paper accelerates this step using randomized approximate nearest neighbors and parallel lookups.
- Attractive forces pull each embedded point toward points that are close in the original space.
- Computing all input interactions is too expensive, so t-SNE restricts each point to interactions with its k nearest neighbors.
- Vantage-point trees are effective in low dimensions but prohibitively expensive for large, high-dimensional datasets.
- Randomized connections to a small subset of approximate neighbors can preserve local geometry while greatly reducing the number of interaction terms.
- FIt-SNE computes approximate nearest neighbors with ANNOY and parallelizes neighbor lookups to accelerate attractive-force computation.
4. Early and Late Exaggeration
Early exaggeration is used during initial optimization, while late exaggeration increases attraction during the final iterations. The paper reports that late exaggeration contracts clusters and makes them easier to distinguish.
- Early exaggeration: Early exaggeration sets α = 12 during the first several hundred iterations before returning it to 1.
- Late exaggeration: For 1 million Infinite MNIST digits, α = 12 during the last 250 of 1000 iterations allows clusters to be more easily distinguished.
- Late exaggeration: Late exaggeration sets α > 1 during the last several hundred iterations to improve cluster interpretability.
- Late exaggeration: Increasing α strengthens attraction within clusters, causing clustered points to contract and become more easily distinguishable.
5. t-SNE Heatmaps
The paper introduces t-SNE Heatmaps, which extend a 1D t-SNE embedding by summarizing gene expression across discretized embedding bins. This supports visualization of gene relationships and their associations with cell clusters.
- The authors hypothesize that 1D t-SNE contains the same information as 2D t-SNE for scRNA-seq cluster analysis.
- In the retinal-cell example, 1D t-SNE and 2D t-SNE contain generally the same information for cluster visualization.
- t-SNE Heatmaps discretize the 1D embedding into p bins and represent each gene by its summed expression across those bins.
- Genes of interest are enriched with nearby genes under the induced Euclidean distance, then displayed as rows in a heatmap.
- Heatmap rows can represent gene clusters, and the induced distance can also be extended to cells for iterative embedding methods.
6. Out-of-Core PCA
Out-of-core randomized PCA enables principal components to be computed for matrices too large to fit in memory by using block-wise matrix operations on disk-stored data. On a 1,000,000 × 30,000 rank-50 matrix, the method used 1GB of memory and achieved approximately 10^-9 spectral-norm accuracy in 16 minutes.
- 6. Out-of-Core PCA: Out-of-core PCA computes principal components without fully loading the input matrix into memory.The algorithm operates on a matrix stored in slow memory and returns orthonormal principal components with a diagonal singular-value matrix.
- 6. Out-of-Core PCA: Randomized PCA approximates a mean-centered matrix with a low-rank matrix by applying it to a small number of random vectors.The resulting vectors approximate the matrix range, after which standard linear algebra computes the principal components.
- 6. Out-of-Core PCA: Block-wise matrix multiplication computes products involving the disk-stored matrix by loading only a fixed number of rows at a time.Each block contributes to the product, so left multiplication by the large matrix never requires loading the full matrix.
- 6. Out-of-Core PCA: The implementation removes an unnecessary renormalization step and applies AA* simultaneously, reducing matrix loads to once per power iteration.Because disk access largely determines runtime, avoiding an extra load is important in the out-of-core setting.
- 6. Out-of-Core PCA: 1GB of memory yielded the top 50 principal components of a 1,000,000 × 30,000 rank-50 matrix in 16 minutes with approximately 10^-9 spectral-norm error.The full matrix would require 240GB merely to store in memory.
7. Summary and Discussion
The paper combines accelerated t-SNE with late exaggeration and out-of-core PCA to address large, high-dimensional datasets. These methods are intended to support visualization without subsampling on a standard personal computer.
- 7. Summary and Discussion: The work presents t-SNE for embedding millions of high-dimensional points in only a few hours.It also includes late exaggeration for cluster identification and out-of-core PCA for datasets that cannot fit in memory.
- 7. Summary and Discussion: Late exaggeration can make clusters easier to identify in t-SNE plots.The feature is presented as a modification of t-SNE alongside the accelerated implementation.
- 7. Summary and Discussion: Out-of-core PCA allows analysis and visualization of datasets that cannot fit into memory.The paper identifies large scRNA-seq datasets as a natural extension and reports applying the approach to 1.3 million mouse brain cells.
8. Software Availability
The paper provides software implementations of FIt-SNE and ooPCA through the KlugerLab GitHub repository, which also includes a script for t-SNE heatmaps.
- 8. Software Availability: FIt-SNE and ooPCA are available through the KlugerLab GitHub repository.The FIt-SNE repository also contains a script for producing t-SNE Heatmap visualizations.
10. Appendix
The appendix expresses the repulsive-force computation as a small set of kernel sums, enabling the required pairwise interactions to be computed efficiently in low-dimensional embeddings.
- 10. Appendix: The repulsive forces are expressed through sums using kernels K1(y, z) and K2(y, z).The kernels are rational functions of the squared distance between y and z.
- 10. Appendix: For a two-dimensional embedding, four kernel sums are computed at each gradient-descent step.The appendix then expresses the repulsive forces in terms of these four sums.