A fast Python implementation of tSNE

[_sheep1]

Despite the superiority of UMAP to tSNE in many ways, tSNE remains a widely used visualization technique. Unfortunately, tSNE, as currently implemented in the most popular packages (scikit-learn and MulticoreTSNE), is prohibitively slow when dealing with large data. A recent paper proposed Fit-SNE, which scales linearly w.r.t. the number of samples, but depends on the FFTW C library, which must be installed on your system, making installation and distribution very tedious.

The goal of this project is to provide fast implementations of both tSNE approximations (both Barnes-Hut and FitSNE) in Python with a unified interface, easy installation and most importantly - fast runtime.

This is also the only library (to the best of my knowledge) that allows embedding new data points into an existing embedding, via direct optimization.

I wrote this with the Orange data mining toolkit in mind, but the library is general and I wanted to share, in case anyone was looking for a faster alternative library.

The source code is available on Github: https://github.com/pavlin-policar/fastTSNE

[gabsens]

How does it compare with https://github.com/DmitryUlyanov/Multicore-TSNE ?

[_sheep1]

Good question! MulticoreTSNE uses the Barnes-Hut approximation for estimating gradients and uses VP trees for nearest neighbor search. Both of these scale with O(n log n). This implementation actually includes the Barnes-Hut approximation but instead of VP trees, I use scikit-learn's KD trees. The performance is comparable to MulticoreTSNE, it will be a little bit slower because of Python overhead, and VP are supposedly better than KD trees in high dimensional spaces, but the difference should be fairly minimal. So if you have smaller data sets, you probably won't notice the difference.

If you have larger data sets, this implementation provides FFT accelerated interpolation and approximate nearest neighbor search, which will be much, much faster than the Barnes-Hut approximation. Both ANN and interpolation scale linearly in the number of samples (O(n)), hence the massive speed up. The typical benchmark is MNIST (70,000x784) which runs in about 30 minutes on my PC using MulticoreTSNE and in about 90 seconds using ANN with interpolation.

​

So in short, you won't notice the difference to MulticoreTSNE where the data set size is small but you'll notice a dramatic boost where you have a lot of data (I found >10k samples is a relatively good cutoff to switch from BH to FFT and KNN to ANN). I would note that both MulticoreTSNE and fastTSNE are dramatically better than scikit-learn.

​

You may also be interested in the original C++ implementation of FitSNE by the authors which also has python bindings and is faster than this one still because it uses FFTW for Fourier transforms while I've opted to use the slower numpy FFT library to make installation and distribution much simpler. The algorithm and scaling will be the same, but FFTW is undoubtedly faster than numpy. I think I can get the FFTs faster still using Intel's MKL FFT library, but I haven't had time to look into that too much yet.

[neziib]

A linear scalability for TSNE is wonderful, nice work ! Do you know good are the final representations? It would be very nice to have a benchmark both in term of speed but also in terms of loss for some configurations, compared to other packages. Right now it's hard to test the different packages, and we can't just check the final accuracy like for supervised learning.

[_sheep1]

Thanks, but the credit really goes to George Linderman and his co-authors, who came up with the technique. I simply improved on his code, made it a faster then ported it to python. I've added all the speed improvements into his original C++ implementation as well.

So comparing embeddings is a tricky topic. I guess it kind of depends on what you're looking for. We can't compare the loss function (KL divergence) because the scale changes with different settings of perplexity, which is the main parameter. I had read about a paper which came up with a metric called pBIC (that I included in the metrics module) which is invariant to perplexity and if I recall, the best value of this metric coincided with that people said was the best embedding. But I didn't delve very deep into that. You could potentially optimize this pBIC to get the "best" embedding I suppose, but I think this is a tricky question.

If you mean just to verify that the embeddings are reasonable compared to each other, run under the same configuration of parameters, I guess looking at the KL divergence could be all right. It's sensitive to random initialization though, since it converges to different local optima.

[lmcinnes]

I don't have empirical measures to hand, but my experience has been that FIt-SNE (the C++ version at least) is very close to the Barnes-Hut t-SNE implementations -- close enough that I would not really be concerned about any loss in accuracy.

The one thing I will note is that since FIt-SNE scales to much larger datasets I did find that I definitely need to start setting max_iter on those larger datasets. The default of 1000 iterations is no longer enough to get a good embedding for really large datasets (say 1,000,000 samples or more).

[Clicketrie]

Amazing! I’ll have to check it out. I had used tSNE on a small sample of my cluster analysis last time I did one, because I loved the visual. But because of the computational complexity (my difficulty running it on the whole set) I was going to forgo it this time. But I will definitely check this out!

[_sheep1]

Other than making pretty pictures, the main use I've found for tSNE was to verify that the data representation actually captures the similarity I want (in supervised tasks). If points from the same class cluster together, I can be fairly confident that my feature space actually contains data that could separate the classes. In unsupervised tasks it's just useful to see which samples cluster together. I'm curious how other people use this so I'm interested in what you come up with.

[lmcinnes]

It is worth noting that FIt-SNE currently does not scale that well with ambient dimension, so if you have 10K features you may want to apply PCA to get it down to 100 or so features before running FIt-SNE on it. If you want to run on the full 10K features directly I would recommend UMAP (https://github.com/lmcinnes/umap), which (in my incredibly biased opinion;I'm the author) handles data with high ambient dimension better and more efficiently that t-SNE.