ElastiK Nearest Neighbors

Insight Data Engineering Project, Boston, April - June 2018

***, , ---, Image similarity search demo running searches on a cluster of 4 Elasticsearch nodes, ## Overview

I built Elasticsearch-Aknn (EsAknn), an Elasticsearch plugin which implements
approximate K-nearest-neighbors search for dense, floating-point vectors in
Elasticsearch. This allows data engineers to avoid rebuilding an infrastructure
for large-scale KNN and instead leverage Elasticsearch's proven distributed
infrastructure.

To demonstrate the plugin, I used it to implement an image similarity search
engine for a corpus of 6.7 million Twitter images. I transformed each image
into a 1000-dimensional floating-point feature vector using a convolutional
neural network. I used EsAknn to store the vectors and search for nearest
neighbors on an Elasticsearch cluster.

The repository is structured:

• demo directory: Twitter Image Similarity search pipeline and web-application
• elasticsearch-aknn directory: EsAknn implementation and benchmarks
• scratch directory: several smaller projects implemented while prototyping

Demo

• ~Web-application~ (Taken down at the end of the Insight program)
• Screencast demo on Youtube
• Presentation on Google Slides

Elasticsearch-Aknn

Usecase

EsAknn is useful for problems roughly characterized:

1. Have a large corpus of feature vectors with dimensionality ~50-1000.
2. Need to run similarity searches using K-Nearest-Neighbors.
3. Need to scale horizontally to support many concurrent similarity searches.
4. Need to support a growing corpus with near-real-time insertions. I.e.,
   when a new vector is created/ingested, it should be available for searching in
   less than 10 minutes.

How does it compare to other approximate-nearest-neighbors libraries?

Tldr: If you need to quickly run KNN on an extremely large corpus in an offline
job, use one of the libraries from Ann-Benchmarks.
If you need KNN in an online setting with support for horizontally-scalable searching and
indexing new vectors in near-real-time, consider EsAknn (especially if you already
use Elasticsearch).

There are about a dozen high-quality open-source approximate-nearest-neighbors libraries.
The Ann-Benchmarks project is a great place to
compare them. Most of them take a large corpus of vectors, build an index, and
expose an interface to run very fast nearest-neighbors search on that fixed corpus.

Unfortunately they offer very little infrastructure for deploying your
nearest-neighbors search in an online setting. Specifically, you still have to consider:

1. Where do you store millions of vectors and the index?
2. How do you handle many concurrent searches?
3. How do you handle a growing corpus? See this issue
   on the lack of support for adding to an index.
4. How do you distribute the index and make searches fault tolerant?
5. Who manages all the infrastructure you've built for such a simple algorithm?

Elasticsearch already solves the non-trivial infrastrcture problems, and
EsAknn implements approximate nearest-neighbors indexing and search atop
this proven infrastructure.

EsAknn's LSH implementation is very simplistic in the grand scheme of
approximate-nearest-neighbors approaches, but it maps well to Elasticsearch and still
yields high recall. EsAknn's speed for serial queries is much slower than other
approximate nearest-neighbor libraries, but it's also not designed for serial
querying. Instead it's designed to serve many concurrent searches over a convenient
HTTP endpoint, index new vectors in near-real-time, and scale horizontally
with Elasticsearch. For specific performance numbers, see the performance section
below and the slides linked in the demo section.

API

Create LSH Model

Given a sample of vectors, create a locality-sensitive-hashing (LSH) model
and store it as an Elasticsearch document.

POST <elasticsearch host>:9200/_aknn_create 

{
    "_index":   "aknn_models",
    "_type":    "aknn_model",
    "_id":      "twitter_image_search",
    "_source": {
        "_aknn_description": "LSH model for Twitter image similarity search",
        "_aknn_nb_tables": 64,
        "_aknn_nb_bits_per_table": 18,
        "_aknn_nb_dimensions": 1000
    },
    "_aknn_vector_sample": [
        # Provide a sample of 2 * _aknn_nb_tables * _aknn_nb_bits_per_table vectors
        [0.11, 0.22, ...],
        [0.22, 0.33, ...],
        ...
        [0.88, 0.99, ...]
    ]
}

This returns:

{ "took": <number of milliseconds> }

Index New Vectors

Given a batch of new vectors, hash each vector using a pre-defined LSH model
and store its raw and hashed values in an Elasticsearch document.

POST <elasticsearch host>:9200/_aknn_index

{
    "_index":       "twitter_images",
    "_type":        "twitter_image",
    "_aknn_uri":    "aknn_models/aknn_model/twitter_image_search"
    "_aknn_docs": [
        {
            "_id": 1,
            "_source": {
                "_aknn_vector": [0.12, 0.23, ...],

                # Any other fields you want...
            }
        }, ...
    ]
}

This returns:

{ "took": <number of milliseconds>, "size": <number of documents indexed> }

Similarity Search

Given a vector in the index, search for and return its nearest neighbors.

GET <elasticsearch host>:9200/twitter_images/twitter_image/1/_aknn_search?k1=1000&k2=10

This returns:

{
    "took": <number of milliseconds>,
    "timed_out": false,

    "hits": {
        "max_score": 0,
        "total": <number of hits returned, up to k2>,
        "hits": [
            {
                "_id": "...",
                '_index': "twitter_images",
                "_score": <euclidean distance from query vector to this vector>,
                '_source': {
                    # All of the document fields except for the potentially
                    # large fields containing the vector and hashes.
                }
            }, ...
        ]

    }

}

Implementation

The key things to know about the implementation are:

1. EsAknn runs entirely in an existing Elasticsearch cluster/node. It operates
   effectively as a set of HTTP endpoint handlers and talks to Elasticsearch via
   the Java Client API.
2. Searches can run in parallel. New vectors can be indexed on multiple nodes
   in parallel using a round-robin strategy. Parallel indexing on a single node has
   not been tested extensively.
3. EsAknn uses Locality Sensitive Hashing
   to convert a floating-point vector into a discrete representation which can be
   efficiently indexed and retrieved in Elasticsearch.
4. EsAknn stores the LSH models and the vectors as standard documents.
5. EsAknn uses a Bool Query
   to find k1 approximate nearest neighbors based on discrete hashes. It then
   computes the exact distance to each of these approximate neighbors and returns
   the k2 closest. For example, you might set k1 = 1000 and k2 = 10.
6. EsAknn currently only implements euclidean distance, but any distance
   function compatible with LSH can be added.

Performance

Speed

EsAknn's speed is generally characterized:

1. Create a new LSH model: < 1 minute.
2. Index new vectors: hundreds to low thousands per second.
3. Search for a vector's neighbors: < 500 milliseconds. Search time scales
   sub-linearly with the size of the corpus.

The corpus vs. search time generally follows a sub-linear pattern like this:

Beyond that, speed is a function of:

1. The vectors' dimensionality.
2. The number of tables (a.k.a. hash functions or trees) in the LSH model.
3. The number of bits in the LSH model's hashes.
4. The number of approximate neighbors retrieved, k1.
5. The number of exact neighbors returned, k2.

In the image similarity search engine, you can see that searches against an
index of 6.7 million 1000-dimensional vectors rarely exceed 200 milliseconds.

Recall

Recall is defined as the proportion of true nearest neighbors returned for
a search and can be evaluated at various values of k2. For example, if
you know your application needs to retrieve the top ten most similar items,
you should evaluate recall at k2 = 10.

Similar to speed, recall depends on the LSH configuration. Increasing k1
is typically the easiest way to increase recall, but the number of tables and
bits also play an important role. Finding a configuration to maximize
recall and minimize search time can be considered a form of hyper-parameter
optimization.

The figure below demonstrates that it is possible to find a configuration
with high-recall and low search-time at various corpus sizes. The points plotted
represent the "frontier" of recall/search-time. That is, I ran benchmarks on
many configurations and chose the configurations with the lowest median search
time for each median recall across three corpus sizes.

The table below shows the best configuration for each combination of corpus size,
median recall, median search time with a median recall >= 0.5., Corpus size, Med. recall, Med. search time, k1, _aknn_nb_tables, _aknn_nb_bits_per_table, ----, ---------------, ---------------, --------------------, ------, -------------------, ---------------------------, 1000000, 1, 191, 500, 200, 12, 1, 1000000, 0.9, 100, 500, 100, 14, 2, 1000000, 0.8, 62, 1000, 50, 16, 3, 1000000, 0.7, 49, 500, 50, 16, 4, 1000000, 0.6, 43, 250, 50, 16, 5, 1000000, 0.5, 50, 250, 50, 19, 6, 100000, 1, 26, 250, 100, 12, 7, 100000, 0.9, 21, 500, 50, 14, 8, 100000, 0.8, 14, 250, 50, 18, 9, 100000, 0.7, 11, 100, 50, 14, 10, 100000, 0.6, 11, 100, 50, 19, 11, 100000, 0.5, 14, 500, 10, 8, 12, 10000, 1, 8, 100, 100, 8, 13, 10000, 0.9, 5, 100, 50, 12, 14, 10000, 0.8, 5, 100, 50, 18, 15, 10000, 0.7, 6, 250, 10, 8, 16, 10000, 0.6, 6, 15, 100, 18, 17, 10000, 0.5, 3, 15, 50, 14, ## Image Processing Pipeline

Implementation

The image processing pipeline consists of the following components, shown in
pink and green above:

1. Python program ingests images from the Twitter public stream and stores in S3.
2. Python program publishes batches of references to images stored in S3 to a
   Kafka topic.
3. Python program consumes batches of image references, computes feature
   vectors from the images, stores them on S3, publishes references to Kafka.
   I use the conv_pred layer from
   Keras pre-trained MobileNet
   to compute the 1000-dimensional feature vectors.
4. Python program consumes image features from Kafka/S3 and indexes them in
   Elasticsearch via EsAknn.

Performance

Image feature extraction is the main bottleneck in this pipeline. It's
embarrassingly parallel but still requires thoughtful optimization. In the end
I was able to compute:

1. 40 images / node / second on EC2 P2.xlarge (K80 GPU, $0.3/hr spot instance).
2. 33 images / node / second on EC2 C5.9xlarge (36-core CPU, $0.6/hr spot instance).

My first-pass plateaued at about 2 images / node / second. I was able to improve
throughput with the following optimizations:

1. Produce image references to Kafka instead of full images. This allows
   many workers to download the images in parallel from S3. If you send the full
   images through Kafka, it quickly becomes a bottleneck.
2. Workers use thread pools to download images in parallel from S3.
3. Workers use process pools to crop and resize images for use with Keras.
4. Workers use the Lycon library for fast image resizing.
5. Workers use Keras/Tensorflow to compute feature vectors on large batches of
   images instead of single images. This is a standard deep learning optimization.

Elasticsearch Versions

• The plugin was developed on Elasticsearch version 6.2.4.
• User mingruimingrui has a fork fork for version 5.6.6.
• User mattiasarro has a fork for version 6.3

Helpful Resources

Here are a handful of resources I found particularly helpful for this project:

1. Locality Sensitive Hashing lectures by Victor Lavrenko
2. Elasticsearch Plugin Starter by Alexander Reelsen
3. Elasticsearch OpenNLP Plugin by Alexander Reelsen
4. Discussions about similarity search in Elasticsearch: one, two.