Using k-d trees to efficiently calculate nearest neighbors in 3D vector space | Hacker News Reader

Using k-d trees to efficiently calculate nearest neighbors in 3D vector space

blog.krum.io

190 pointsnkurz40 comments

Extract from Andrew Moore's PhD Thesis: Ecient Memory-based Learning for Robot Control PhD. Thesis; Technical Report No. 209, Computer Laboratory, University of Cambridge. 1991.

[1991] https://www.ri.cmu.edu/pub_files/pub1/moore_andrew_1991_1/mo...

Table 6.4 describes my actual implementation of the nearest neighbour algorithm. It is called with four parameters: the kd-tree, the target domain vector, a representation of a hyperrectangle in Domain, and a value indicating the maximum distance from the target which is worth searching.

It’s been a really long time since I’ve read that paper. Thanks for the link and the call-out out 6.4. That “target distance worth searching” seems interesting when thinking about color-replacement.

A faster alternative, using a LUT (lookup table):

- At initialization time, from the given palette (e.g. the CSS palette), build a 16-bit (RGB565) to 8-bit LUT, and find for every 2^16 entry the nearest palette color. Time complexity: O(n^2), but done once, and can be avoided if having the LUT pre-calculated as a constant.

- At the image conversion, with per-pixel O(1) time complexity: for every RGB888 pixel, pack the pixel to RGB565, and use the LUT to find de palette index i.e. pal_index = lut[RGB888to565(pixel)]. With e.g. a 3 GHz modern 4-way OooE CPU you can convert pixels at > 1 Gpixel per second (> 10^9 pixels per second), because of the LUT high cache hit ratio.

In case of requiring more precision, you could use a 24-bit LUT (uint8_t[2²⁴] -> 16 MiB of memory, instead of the uint8_t[2¹⁶] -> 64 KiB of memory), or a intermediate case of willing to avoid running out of the L2/L3 cache e.g. a 21-bit LUT (RGB685, uint8_t[2^²¹] -> 2MiB). The 16-bit table would allow you run in the L1 data cache, being very fast (you could use even a smaller table, e.g. 15-bit RGB555 (32 KiB), or 14-bit RGB554 (16KiB), with similar results regarding color conversion, but with even lower L1 data cache requirements.

For those interested in this space, folks might like some of the work that went in to visualizing T Distributed Stochastic Neighbor Embedding and the successor to that LargeVis which use multiple forms of KNN trees each with their own trade offs:

https://arxiv.org/abs/1301.3342 (Barnes Hut)

https://arxiv.org/abs/1602.00370 (LargeVis)

One uses Vantage Point Trees which sudivides the space in to quadrants, and the other uses Random Projection Trees.

Both are interesting exercises in using KNN like techniques to visualize high dimensional spaces.

There's also FLANN which is an approximate nearest neighbour algorithm which is routinely used for things like feature/keypoint matching (e.g. SIFT keypoints are 128 dimensional).

https://www.cs.ubc.ca/research/flann/

Barnes Hut is a fun one to implement for N-body simulations.

I was toying with compqring similarity of texts just last night, is there anything that could handle 10k+ sized sparse vectors? Obviously there are tools like lucene, but I wanted to mess around in the interactive interpreter

Check out LargeVis then. The original github repo by the authors has a sparse mode: https://github.com/lferry007/LargeVis/

TSNE didn't scale well no matter what you did to it (hence the newer technique) - in general though you can't beat cosine similarity of TFIDF vectors as a baseline.

Fun Fact: I am an author of https://www.Photopea.com - advanced image editor, that has 50 000 lines of code.

In my code, I have three independent implementations of KD-trees, with different dimensionality for different purposes ... I did not realize it until now :D

This is a very good intro to K-D trees, big thank you to the author. However, I think that a three-dimensional lookup table would be more performant in this case. 16M (256^3) of entries are needed to cover RGB color space with 8 bits per channel, that's a lot, but lookup time is O(1) with a very small constant.

I think the problem there is that if you're only looking at a few images then generating the 16M lookup table takes more time than using the K-D tree itself.

Some kind of memoisation scheme might work though.

You can actually use the k-d table to compute the lookup table, as the lookup table is nothing more really than an image containing a "pixel" with each possible color exactly once. The author already ran the algorithm on a 12285x14550 image, which took 10x more time than running on a 256x256x256 "image".

Of course, if you want to compute the lookup table dynamically on each run, you might use memoization and hope the image is either small (and therefore uses comparatively few colors) or large with few actual colors in use.

I think if I was actually building this, I’d do that initial calculation, and then compress it into a linear decision tree. Potentially tiny data set, very fast lookups. There’s definitely going to be some inaccuracy there, but I suspect small

That might not have very good cache performance though?

Which reminds me that some time ago I build a simple [kd-tree library in haskell](https://github.com/fhaust/kdtree).

The interesting part is that it's easy to construct a (lazy) list of all points ordered based on the distance to the query point. Other queries (n-nearest neighbours, points in radius) are just derived from that function and traverse only the first some elements of the list.

Simpler:

- Compute z-value for each vector, store in binary tree (or sorted array, etc.), keyed by z-value.

- Compute z-value of search vector, search the data structure, find the k adjacent items in each direction (k is small, e.g. 1-5).

- Of the vectors found, compute minimum distance from the search vector, r.

- Do another search of your data structure, using a sphere of radius r, (easy z-order algorithm).

- The nearest neighbor is the vector resulting from this search whose distance from the search vector is minimum.

You’ve just described a single axis kd-tree roughly. Which means you’ll also perform catastrophically for patterns with say a line in the Z-axis. It’s not that hard to do a proper spatial tree. If you hate kd-trees, an octree or a grid are also more robust than a “single axis binary split tree”.

Not sure what you mean by a single-axis kd-tree.

Z-order + sorted array/tree/etc. is a generalization of an oct-tree. An octtree is just z-order combined with a degree 8 trie. And an octtree is simply an encoding of a 3-d grid.

So I really don't understand your objections.

Also, I haven't tried NN search using z-order, but range searches work very well. The partitioning induced by z-order adapts well to all sorts of distributions.

Oh! I see the problem.

I don't mean just a tree keyed on the z-axis. I mean z-order: https://en.wikipedia.org/wiki/Z-order

Ahh! I wasn’t sure, so sadly I chose “z as in z axis” not Z as in Z/space-filling/Hilbert/Morton order. That’s a grid in disguise! :).

That's only viable when you have a small data set, which is when you don't care much about indexing to begin with.

Sounds like GeoHash lookup which is actually pretty robust if you have a good data structure for indexing. For example: https://redis.io/commands/georadius

Huh? I've used z-order inside database systems, with many millions of records. Works incredibly well.

Building and balancing a k-d tree is not free, either. I suspect that the GP's approach is roughly the same in terms of theoretical complexity, but effectively trades memory for smaller constants and a chance to trip over edge cases like dense groups of points.

Semi-related: Anybody know what's the current state of the art for calculating nearest neighbors in high-dimensional spaces (at least 10D, maybe 100s of D)?

Various approaches based on LSH (locality sensitive hashing) and random projections.

But "the best" really depends on your data profile, as the generic problem formulation, for completely generic vectors, is too… generic. See also the curse of dimensionality [0].

In practical terms, if you're looking for an actual implementation to use, here's a benchmark of some state-of-the-art implementations for high-dimensional vector search: http://radimrehurek.com/2013/12/performance-shootout-of-near... and also the updated https://github.com/erikbern/ann-benchmarks.

In particular, for something good AND open source:

- Leonid Boytsov's NMSLIB [1]

- Erik Bernhardsson's Annoy [2]

- Facebook's FAISS [3]

For a commercial engine focused more on practical use (index management, transactions, versioning, support), see ScaleText [4] (disclaimer: our product).

[0] https://en.wikipedia.org/wiki/Curse_of_dimensionality [1] https://github.com/searchivarius/nmslib [2] https://github.com/spotify/annoy [3] https://github.com/facebookresearch/faiss [4] https://scaletext.com/

Very helpful, many thanks!

Discussion on reddit: https://www.reddit.com/r/programming/comments/7wg91g/using_k...

How easy is it to extend k-d trees to deal with objects with finite size (vs. points)? Are there any libraries that do this? I might have to do something similar to this in the future, but I need nearest neighbors for 'spheres' in higher dimensions.

AABB Trees (Axis Aligned Bounding Box) might be an option. Most realtime collision detection systems utilize them.

Or more generally bounding volume hierarchies, of which AABB trees are a subset.

They’re great. Good perf, while being one of the simpler ways of accelerating spatial queries (compared to other types of trees, at least)

Related to another thread on this discussion:

Z-order combined with sorted array/tree/etc. generalizes to non-point objects quite easily. I think the Z-order NN algorithm will "just work" with non-point data.

If by "just work" you mean testing each constituent vertex, it'll fail. Consider the case when the needle is at the midpoint of a line between two points far away.

No, that's not what I mean.

A point is represented by a z-value at the resolution of space.

A non-point is approximated by multiple z-values, at varying resolutions. You can trade the accuracy of the approximation against the number of z-values.

I explained the NN algorithm in another comment on this thread. It shouldn't matter if some of the adjacent z-values are actually from the approximation of the same object.

See r-trees, interval trees, and ball trees.

k-d trees can be used with things like lines and triangles. This was pretty common in computer graphics for a while. There is also a generalization called BSP tree that allows for arbitrary split planes.

The catch with both is that you need to be careful with primitives that intersect split planes. You either need to introduce either additional bookkeeping to avoid testing the same primitive multiple times, or you need to split the primitives at each plane. Both are not really desirable properties. Also, primitives that lie exactly on a split plane will lead to unstable results as the tests involving the split plane and the primitive are usually bound to to give slightly different results.

A bounding volume hierarchy is generally easier to implement and more reliable. And people have optimised the heck out of algorithms to construct them.

This is color quantization, right? https://en.wikipedia.org/wiki/Color_quantization

Which leads us to dithering

I can't find it right now but there was a lovely article on old ways and new ways to dither.