# Microsoft Research Blog

The Microsoft Research blog provides in-depth views and perspectives from our researchers, scientists and engineers, plus information about noteworthy events and conferences, scholarships, and fellowships designed for academic and scientific communities.

## Optimizing Barnes-Hut t-SNE

May 30, 2018 | By Tavian Barnes, Software Development Engineer II

Barnes-Hut optimization highlighted as points converge to their t-SNE positions

Ten years ago, while writing a physics engine, I learned about the Barnes-Hut algorithm for the gravitational n-body problem. Normally, computing the Newtonian gravitational forces between n bodies requires $$O(n^2)$$ evaluations of Newton’s law of universal gravitation, as every body exerts a force on every other body in the system. Barnes-Hut is an approximation that brings that complexity down to $$O(n \log n)$$, by treating clusters of faraway objects as single particles. An octree is used to subdivide the space, and when a certain threshold is met, an internal node may be treated as a point particle at its children’s center of mass, rather than recursing into it.

I didn’t use the Barnes-Hut algorithm at the time. But years later I was surprised to see it being used in the seemingly unrelated context of deep learning. t-SNE is a dimensionality reduction technique often used for visualizing high-dimensional data like word embeddings. It attempts to map points that are close in the high-dimensional space to points that are still close in a low-dimensional space, preserving clusters in the high-dimensional representation. It turns out that t-SNE can be accelerated by applying the same Barnes-Hut technique as gravitational simulation, since they both involve an inverse square law computed between pairs of points. In this post, I’m going to take the Barnes-Hut t-SNE implementation by the original inventor of t-SNE, Laurens van der Maaten (with Geoffrey Hinton), and show how its performance can further be increased while keeping the algorithm the same.

The first step is to check out and build the code:

$git clone git@github.com:lvdmaaten/bhtsne.git Cloning into 'bhtsne'... remote: Counting objects: 205, done. remote: Total 205 (delta 0), reused 0 (delta 0), pack-reused 205 Receiving objects: 100% (205/205), 83.59 KiB | 2.70 MiB/s, done. Resolving deltas: 100% (110/110), done.$ cd bhtsne
$g++ -O2 -g -Wall -Wno-sign-compare sptree.cpp tsne.cpp tsne_main.cpp -o bh_tsne To get an idea of its current performance, let’s test it on some real-world data, the GloVe word vectors:$ curl -LO http://nlp.stanford.edu/data/glove.6B.zip
% Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
0   308    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0
100  822M  100  822M    0     0  5661k      0  0:02:28  0:02:28 --:--:-- 3072k
$unzip glove.6B.zip Archive: glove.6B.zip inflating: glove.6B.50d.txt inflating: glove.6B.100d.txt inflating: glove.6B.200d.txt inflating: glove.6B.300d.txt We have to strip out the first column, which is the words themselves:$ cut -d' ' -f2- glove.6B.50d.txt >glove.6B.50d.cut.txt

And now we can use the Python wrapper to run the tool:

$./bhtsne.py -v -r1 -i glove.6B.50d.cut.txt -o glove.6B.50d.tsne.txt Read the 400000 x 50 data matrix successfully! Using random seed: 1 Using no_dims = 2, perplexity = 50.000000, and theta = 0.500000 Computing input similarities... Building tree... - point 0 of 400000 ... - point 390000 of 400000 Input similarities computed in 21409.84 seconds (sparsity = 0.000682)! Learning embedding... Iteration 50: error is 131.843420 (50 iterations in 729.03 seconds) ... Iteration 999: error is 6.373127 (50 iterations in 654.72 seconds) Fitting performed in 15039.99 seconds. Wrote the 400000 x 2 data matrix successfully! ./bhtsne.py -v -r1 -i glove.6B.50d.cut.txt -o glove.6B.50d.tsne.txt 36697.76s user 228.13s system 99% cpu 10:17:09.79 total Wow, that took a long time! Let’s see if we can speed it up a bit. ### Profiling First, because premature optimization is the root of all evil, we should profile to see where the bottlenecks are. We don’t have all day, so let’s make a smaller input file for benchmarking, with just the first 1,000 points:$ head -n1000 glove.6B.50d.cut.txt >glove.6B.50d.small.txt

And use the Linux perf tool to sample the behavior of the program:

$perf record --call-graph=dwarf ./bh_tsne ‘perf report’ shows where it’s spending most of its time: - 94.58% 0.00% bh_tsne bh_tsne [.] main - main - 94.57% TSNE::run - 87.76% TSNE::computeGradient - 49.31% SPTree::computeNonEdgeForces - 48.92% SPTree::computeNonEdgeForces + 47.42% SPTree::computeNonEdgeForces 21.25% SPTree::computeEdgeForces + 10.90% SPTree::SPTree + 6.11% SPTree::~SPTree + 4.02% TSNE::evaluateError + 2.10% TSNE::computeGaussianPerplexity So it looks like SPTree::computeNonEdgeForces(), which happens to be the octree traversal, is the main candidate for optimization. Diving into the annotated assembly view of perf report by pressing A, a couple hotspots are evident. ### Avoiding square roots Let’s start with this one: │ if(is_leaf || max_width / sqrt(D) < theta) { 5.24 │ sqrtsd %xmm2,%xmm4 │ ↓ ja 2115 <SPTree::computeNonEdgeForces 10.37 │18e: divsd %xmm4,%xmm1 0.45 │ comisd %xmm1,%xmm3 1.48 │ ↓ ja 20b0 <SPTree::computeNonEdgeForces(unsigned int, double, 1e0 The criterion $$\frac{w_{\max}}{\sqrt{D}} < \theta$$ determines whether the current node is a good enough approximation of its children, or if we have to recurse into it and handle its children separately. Square roots and division are both very expensive operations, but because $$x^2$$ is an increasing function on positive numbers, we can rewrite the inequality in the more efficient form $$w_{\max}^2 < \theta^2 D$$. With this patch, bhtsne is about 5% faster on this dataset. ### Hoisting invariants Now let’s move on to this hotspot: │ double cur_width; │ for(unsigned int d = 0; d < dimension; d++) { │ cur_width = boundary->getWidth(d); │ max_width = (max_width > cur_width) ? max_width : cur_width; 10.98 │ cc: maxsd (%r9),%xmm1 This loop is computing the maximum width of the node along any axis. But doesn’t that remain constant for the life of the node? We could compute that once for every node and store it in a field to save some work. Doing that makes it about another 5% faster. ### Shrinking the footprint Why exactly did the last patch help? ‘maxsd’ is not a slow operation (computing the maximum of two doubles) and it wasn’t a long loop (dimension == 2 in this case). In fact, it was slow for the same reason almost all code is slow these days: memory locality. The getWidth(d) call ultimately looks up the width in an array somewhere else in memory. As CPUs have gotten faster over the years, RAM has failed to keep pace. Levels of smaller and faster on-CPU caches have mitigated the issue, but the smaller your memory footprint is, and the fewer pointers you have to chase, the more of your working set will fit in lower levels of cache and the faster your program will be. So, if we can reduce the memory footprint of the tree, we’ll probably be able to keep more of it in the cache, resulting in faster code. After reworking the implementation to use standard C++ containers instead of manual memory management, this patch introduces a smaller node structure. These changes make it almost 40% faster yet on our benchmark. We can shrink the footprint a bit more by noticing that all the nodes at the same level as each other have the same sizes, just different positions, so they can share their width member. This patch gains us about another 5%. Finally, by eliminating the scratch space array buff and just recomputing the values as necessary, we get almost another 10% faster. Altogether, we’re now over 70% faster on this benchmark. ### Scaling up Now let’s check how this scales to larger datasets. Figure 1 – Run time before and after optimizations. Figure 1 shows that with 10,000 points, we’re over 100% faster, going from 3m09s to just 1m32s. But at 100,000 points, the benefit falls back down to 75%, going from 1h14m to 42m. The output hints that the bottleneck has started shifting to a different part of the program:$ ./bhtsne.py -v -r1 -i glove.6B.50d.large.txt -o glove.6B.50d.tsne.txt
...
Input similarities computed in 1388.20 seconds (sparsity = 0.002531)!
...
Fitting performed in 1081.47 seconds.
...
./bhtsne.py -v -r1 -i glove.6B.50d.large.txt -o glove.6B.50d.tsne.txt 2529.04s user 1.70s system 99% cpu 42:19.12 total

‘perf’ confirms that with more points, we’re spending more time in a different part of the program:

-   95.04%     0.26%  bh_tsne  bh_tsne              [.] TSNE::run
- 94.78% TSNE::run
+ 45.46% SPTree::computeNonEdgeForces
25.83% SPTree::computeEdgeForces
+ 4.71% SPTree::SPTree
+ 1.61% SPTree::~SPTree
- 11.67% TSNE::computeGaussianPerplexity
- 10.62% VpTree::search (inlined)
+ 10.23% VpTree::search (inlined)
0.79% expf32x
+ 4.56% TSNE::evaluateError
0.69% TSNE::symmetrizeMatrix

VpTree::search() is starting to show up in the profile, and its share of execution time only goes up as the number of points grows. This corresponds to the “Computing input similarities…” part of the program, which involves finding the nearest neighbors of every point in the dataset. A careful look at the profile shows that almost all that time is spent computing high-dimensional Euclidean distances:

+    9.59%     0.00%  bh_tsne  bh_tsne              [.] VpTree<DataPoint, &(euclidean_distance(DataPoint const&, DataPoint const&))>::search (inlined)
+    9.02%     9.00%  bh_tsne  bh_tsne              [.] euclidean_distance

Unfortunately, we can’t do the same thing we did the last time we saw a Euclidean distance computation. The code is using vantage-point trees to find nearest neighbors, which fundamentally rely on the metric properties of the distance function. Eliminating the square root makes the new “metric” violate the triangle inequality and would cause the VP tree to give incorrect results. Furthermore, the square root isn’t even the slowest part — most of the time is just spent iterating over the coordinates:

│            diff = (x1[d] - x2[d]);
6.08 │23:   movsd  (%rcx,%rax,8),%xmm1
21.08 │      subsd  (%rsi,%rax,8),%xmm1
│            dd += diff * diff;
0.15 │      lea    0x1(%rax),%rdx
10.72 │      mulsd  %xmm1,%xmm1
│        for(int d = 0; d < t1._D; d++) {
18.73 │      cmp    %rax,%rdi
0.00 │    ↑ jne    2e90 <euclidean_distance(DataPoint const&, 20
│        }
│        return sqrt(dd);
0.46 │      ucomisd %xmm0,%xmm2
11.74 │      sqrtsd %xmm0,%xmm1

High dimensional nearest neighbors is a hard problem at the best of times, because of the curse of dimensionality. The best choice here might be to use an approximate nearest neighbors algorithm, or at least an all-nearest-neighbors algorithm that avoids wasted work. But I’ll save that for a future post.

### Conclusion

So how did we do on the original, 400,000-point dataset?

\$ ./bhtsne.py -v -r1 -i glove.6B.50d.cut.txt -o glove.6B.50d.tsne.txt
Read the 400000 x 50 data matrix successfully!
Using random seed: 1
Using no_dims = 2, perplexity = 50.000000, and theta = 0.500000
Computing input similarities...
Building tree...
- point 0 of 400000
...
- point 390000 of 400000
Input similarities computed in 21370.59 seconds (sparsity = 0.000682)!
Learning embedding...
Iteration 50: error is 131.843420 (50 iterations in 273.23 seconds)
...
Iteration 999: error is 6.373127 (50 iterations in 239.38 seconds)
Fitting performed in 5394.53 seconds.
Wrote the 400000 x 2 data matrix successfully!
./bhtsne.py -v -r1 -i glove.6B.50d.cut.txt -o glove.6B.50d.tsne.txt  26989.04s user 59.58s system 99% cpu 7:32:10.40 total

With all the above optimizations, we’ve saved 2 hours and 45 minutes on this computation, making it 36% faster. The second stage of the computation, where the optimizations apply, is over 178% percent faster, starting at over four hours and ending up under one and a half. These optimizations have been submitted upstream as a pull request and I hope to make more improvements in the future!

### References

L.J.P. van der Maaten. Accelerating t-SNE using Tree-Based Algorithms. Journal of Machine Learning Research 15(Oct):3221-3245, 2014.
L.J.P. van der Maaten and G.E. Hinton. Visualizing High-Dimensional Data Using t-SNE. Journal of Machine Learning Research 9(Nov):2579-2605, 2008.
J. Barnes and P. Hut. A hierarchical O(N log N) force-calculation algorithm. Nature 324 (4): 446–449, 1986.

## Up Next

Programming languages and software engineering

### Summer of MakeCode for kids

Microsoft MakeCode is a web-based learning environment for kids and teens to create with technology. MakeCode takes a unique approach to computing education by combining the magic of making with the power of code as a way to make learning about computers and technology more accessible and interesting to a wider audience. We have created easy MakeCode projects that you can do with your kids at home over the summer using a different MakeCode product each week.

Jacqueline Russell

Principal Program Manager

Programming languages and software engineering

### The democratization of data science with Dr. Chris White

Episode 27, June 6, 2018 - Dr. White talks about his “problem-first” approach to research, explains the vital importance of making data understandable for everyone, and shares the story of how a one-week detour from academia turned into an extended tour in Afghanistan, a stint at DARPA, and, eventually, a career at Microsoft Research.

Microsoft blog editor

Programming languages and software engineering

### CodeTalk: Rethinking IDE accessibility

It is a bright afternoon in the Microsoft Research India lab. Research Fellow, Venkatesh Potluri, sits at his computer, frantically racing against the clock to fix one last bug before the end of the day. The computer blares into his headphones—not music, but a robotic rendition of the code, as Venkatesh uses a program called a […]

Suresh Parthasarathy

Senior Research Developer