Writing useful notes for myself 6 months ago

Feature engineering is often the most important part of a supervised learning project. It's how we translate our human-readable data into a machine-readable form. If we translate faithfully, the input features for our model should represent everything that a human could understand from each observation in the dataset. However, any mistranslations or wonky projection could hobble our model's performance before it's even started training.

While over-compression or -interpretation can be dangerous, some deliberate massaging of the input data can be useful. For example, if we know that the raw representation of a variable is too granular, we might compress the representation by bucketing its values, reducing the number of input dimensions. Incorporating these inductive biases into our models usually allows them to train faster and raises the performance ceiling. 

Even at a glance, intuition can tell us a lot about a timestamp. When we look at a date or time, we sense that:

While each one is distinct, dates are not totally independent, and similar dates share similar properties. For example, the 1st June is more similar to the 2nd June than it is to the 1st December.

The same applies to times - 3am is more similar to 4am than it is to 3pm.

Dates repeat annually, and the weather on 2023-06-01 is likely to be similar to the weather on 2022-06-01.

Days repeat weekly, and most people are likely to have different habits on weekdays vs weekends.

Times repeat daily, and peoples habits are also correlated with these cycles, eg most people are probably asleep at 3am on any given day.

Many of these patterns are based on the cyclical nature of dates and times, and the correlated repetition of patterns in nature and human behaviour.

However, these cyclical features are poorly expressed by the raw representations alone. Although a day starts at 00:00 and ends at 23:59, we know from experience that 00:01 is as similar to 23:59 as 16:29 is to 16:31. Although a year starts on 01-01 and ends on 12-31, in many respects the 1st January is as similar to 31st December as 1st June is to 2nd June.

We have to rely on our experience to understand these features. An untrained model doesn't have that experience, and the raw representations make them hard to learn! Our own intuitions and experiences can (and should!) guide how we translate the raw data into a machine-readable form.

As explored above, we know that a timestamp is comprised of many component parts (minutes, hours, days, weeks, months, years, etc). Those components can be extracted and used as individual features, in the hope that some of them might be correlated with our model's target variable.

Pandas' DatetimeIndex has a set of attributes which are commonly used for this purpose. If we start with a datetime column, we can define a new set of component columns based on the original:

df['year'] = df['datetime'].dt.year
df['month'] = df['datetime'].dt.month
df['day'] = df['datetime'].dt.day
df['day_of_week'] = df['datetime'].dt.day_of_week
df['day_of_year'] = df['datetime'].dt.day_of_year
df['minute_of_day'] = df['datetime'].dt.hour * 60 + df['datetime'].dt.minute
df['quarter'] = df['datetime'].dt.quarter
df['is_month_end'] = df['datetime'].dt.is_month_end
df['is_leap_year'] = df['datetime'].dt.is_leap_year
date_features = df[[
    'year', 'month', 'day', 'day_of_week', 'day_of_year', 'minute_of_day', 
    'quarter', 'is_month_end', 'is_leap_year'
]].values

Great! We've gone from one value to many, each of which represents a different meaningful aspect of the original timestamp. If we're lucky, the model will be able to use some of these features to make predictions.

However, these numbers aren't ideally formatted yet - we still need to translate these features into a more machine-readable form, ideally set in the unit interval (ie values between 0-1).

The simplest approach is to use an ordinal encoding, where each value is represented by a number. Ideally, the numbers should then be normalised to take values between 0-1.

For example, pandas days of the week are given the following numbers:

These can then be normalised by dividing by the maximum value:

df['day_of_week'] = df['datetime'].dt.day_of_week / 6

By applying the same idea to our other values, we can create a feature vector which represents the component parts of our timestamp in a form which is suitable for an ML model.

These values capture the fact that, for example, a week proceeds linearly from Monday to Sunday, and that a year proceeds linearly from January to December, with Tuesday being closer to Monday than it is to Sunday, and February being closer to January than it is to December.

However, this approach has a major drawback. We've imposed a boundary between the maximum and minimum values, losing any of the cyclical features of the original timestamp!

For example, our minute_of_day implies that 23:59 is as far as it is possible to be from 00:00, when in reality they're adjacent!

To shake off this boundary, people often reach for a one-hot encoding. This is a binary representation of each value, where each value is represented by a vector of length n, where n is the number of possible values. For example, the days of the week might be represented as follows:

By separating the possible value into linearly independent components, we've opened up the possibility for the model to learn about the circular nature of the data. However, this encoding scheme comes with a new set of drawbacks.

This encoding is very inefficient, especially for higher-cardinality data. While our 7-day week introduces 7 new columns to our dataset, a day_of_year feature becomes 365 new features. minute_of_day adds 1,440 new columns. This seems like a poor trade-off for just the potential for the model to learn about some latent behaviour.

Our model is also not imbued with a natural sense of the relationships between the columns - any relationships it develops must be learned from scratch. The classes' linear independence mean that if examples are rare or missing in the training data (eg we have no training observations from day_of_year 129), test-performance is likely to miss their meanings (our model can't intuit that day_of_year 129 will be similar to 128 or 130).

We can improve the efficiency of these one-hot encoded features by grouping our observations into a limited number of buckets. This lessens the impact of sparsity in training data. For example, we might choose to group our minute_of_day observations according to 24 buckets, one for each hour, instead of our original 1,440. We could use 48 buckets, each representing half an hour.

Choosing a suitable number of buckets is a tricky problem, and relies on a strong sense of what patterns might exist in the data, or a lot of time to spend on hyperparameter tuning.

However, this encoding scheme still lacks an explicit inductive bias towards the cyclical relationship between these features. Though that's easier to learn in a compressed space, it would be nice if we could include it more intentionally.

To overcome the linear independence of features and introduce an explicit relationship between neighbouring values/buckets, we can introduce fuzziness to our one-hot encoding. Instead of encoding a single 1 in each vector and filling the rest of the values with 0s, we can instead apply a gaussian distribution (or any other appropriate distribution) to each vector with its center where we had previously placed our 1.

For example, our days of the week might be encoded as:

By allowing the model to observe (weak) correlations between target features and neighbouring cyclical feature indexes, we make it easier for the model to learn that Sunday and Monday share some similarities/proximity, or that 23:59 is adjacent to 00:00.

All of the previous hacks, workarounds and adaptations get us close to a solution, but they're either indirect, inelegant, or inefficient. Fundamentally, they all rely on the model to learn relationships that we already know and understand well, and should be able to encode explicitly.

It is possible to encode all of our intuition about the cyclical nature of dates and times into a single feature vector, without losing any of the information about the component parts of the timestamp. It's all made possible by some simple trigonometry. Let's consider what we already have

A single normalised value between 0-1, representing the point of the observation in the cyclical feature space.

For example, a Tuesday in a 7-day week is encoded as 0.166.

A projection of our input variable to a new space, with values also bounded between 0-1. Distances between pairs of points in our input and output space should be correlated, while allowing for distances to be calculated across the unit interval as if it were continuously connected.

In short, we should be projecting the points on our 1D input onto a 2D circle!

We can achieve this by simply calculating the sine and cosine of our input variable.

Let's create a set of random input values between 0-1, representing observations of our cyclical feature.

Now, let's calculate sin(2πx) and cos(2πx) for our input values. The following plot has been scaled to place outputs back in the unit interval using the following code:

import numpy as np

def encode_cyclical(a: float):
    """
    returns the sine and cosine of the input value, scaled to the unit interval
    """
    x = (np.sin(2 * np.pi * a) + 1) / 2
    y = (np.cos(2 * np.pi * a) + 1) / 2
    return x, y

Plotting the outputs on independent dimensions shows us what we want!

We can see instinctively that the encoding we've created fulfils our requirements. Points are evenly spaced along our circular number line, and the boundaries of the unit interval have been joined to create a fully cyclical system. 

We can go further than a visual proof though. Let's verify our assumptions with some more data.

We can test whether our requirement for distances is met by our cyclical encoding scheme by calculating distances for random pairs of points in both spaces, and comparing the results.

import numpy as np
import pandas as pd

# first, create a table with two columns of random values between 0-1
df = pd.DataFrame({"a": np.random.random(100), "b": np.random.random(100)})

def min_distance_in_1D(a: float, b: float):
    """
    returns the shortest distance between two points in the unit interval,
    explicitly correcting for the assumption that the interval is cyclical
    """
    distance_within_unit_boundary = abs(a - b)
    distance_across_unit_boundary = abs(1 - distance_within_unit_boundary)
    return min(distance_within_unit_boundary, distance_across_unit_boundary)


def encode_cyclical(a: float):
    x = (np.sin(2 * np.pi * a) + 1) / 2
    y = (np.cos(2 * np.pi * a) + 1) / 2
    return x, y


def distance_on_circle(a: float, b: float):
    """
    returns the shortest euclidean distance between two points in the unit
    interval projected onto angles around a unit circle
    """
    xa, ya = encode_cyclical(a)
    xb, yb = encode_cyclical(b)
    return np.sqrt((xa - xb) ** 2 + (ya - yb) ** 2)

# find the distances
df["min_distance_in_1D"] = df.apply(
    lambda row: min_distance_in_1D(row.a, row.b), axis=1
)
df["distance_on_circle"] = df.apply(
    lambda row: distance_on_circle(row.a, row.b), axis=1
)

df.plot.scatter(x="min_distance_in_1D", y="distance_on_circle")

Distances between random pairs of points in the unit interval as measured by their (corrected) distance in one dimension, and the corresponding euclidean distance in our new cyclical space.

That line tells us that the scheme is natively encoding the cyclical nature of the data, without us having to explicitly correct for it in our distance calculation. It's not perfectly linear, but that's okay! The important thing is that the function is monotonically increasing. Now any learning algorithm should be able to resolve patterns in the data which cross the domain boundary.

Nothing beats the efficiency and simplicity of sine-cosine embeddings for representing cyclical features. By directly encoding information about the fundamental nature of the values into the constructed space, they offer unparalleled flexibility for models to interpret the data efficiently.

More than anything though, this post should serve as a reminder of the vast array of options available for feature engineering tasks, and highlight the importance of selecting the right approach!

A little bit of trigonometry is all you need to get the most from your datetime features

The best way to encode dates, times, and other cyclical features

A basis is the set of building-block vectors which we can use to represent any other vector in a particular space.

Basis vectors should be independent of one another, ie we should not be able to express any of the vectors in the set using a combination of the others. 

For example, in a 2D vector space, the vectors [1, 0] and [0, 1] form a basis, because any other vector in the space can be expressed as a linear combination of them, and it's not possible to express one with the other.

For example, the vector [2, 3] can be represented as:

The coefficients 2 and 3 represent how much of each basis vector is needed to form the original vector.

We could also use the vectors [1, 1] and [-1, 1] form an alternative 2D basis, allowing us to express the original vector differently:

Bases are important because they allow us to represent complex data, like the features of images or text, as collections of simpler components. The choice of basis can have a significant impact on the ability of an algorithm to learn patterns and relationships within the data. By choosing an appropriate basis, we can simplify the representation of the data and make it easier for an algorithm to extract meaningful features.

In certain problem settings, we can make architectural choices which deliberately bias a model to organise its internal representation of certain features along its basis dimensions. These are often called inductive biases. In these cases, those dimensions (or groups of dimensions) will develop specific importance or meaning through the training process. We call these bases privileged.

For example, in a convolutional neural network, convolutional layers are explicitly designed to learn local spatial patterns in images such as edges, corners, and textures. The neurons of these layers are constrained to the nearby regions of the input, which encourages the learning of spatially localised features that align with the basis dimensions.

For example, suppose that a CNN is trained to classify images of cats and dogs. If the learned features do not align with the basis dimensions that detect features relevant to cats and dogs, such as whiskers or ears, then the network may not be able to effectively distinguish between the two classes of images. In contrast, if the learned features align with the basis dimensions, then the network can effectively capture the desired patterns in the input and achieve high classification accuracy.

Similarly, the multi-head attention mechanism in a transformer model is designed such that different heads will attend to different aspects of the input sequence, promoting feature alignment.

Regularisation methods like weight decay, dropout, or batch normalisation can also add constraints to the learning process, controlling the model complexity or reducing overfitting. These constraints impose yet more priors on the model, implicitly pushing it to learn features that align with the basis dimensions.

In the creation of a non-privileged basis, the model is given no access to privileged information, and the optimiser is given no incentive to align its basis dimensions with a particular type of feature. Instead, the available dimensions are treated as equally important, and the model is free to distribute information equally across them.

For example, in an embedding space generated by a process like word2vec, each word is represented as a high-dimensional vector, with each dimension corresponding to a learned feature of the word. Rather than being deliberately engineered to represent specific linguistic features, the basis vectors are distributed based solely on the relationships of words to one other in the space, regardless of how that vector space itself is arranged.

The distribution of basis vectors in the space is driven wholly by the optimiser, rather than by engineering specific dimensions to capture specific aspects of the data. 

In cases where the space is larger than the information being encoded, this can result in information being represented redundantly and non-orthogonally across many dimensions. This can make model with non-privileged bases more robust to noise and variations, at the expense of interpretability.

Apart from considering the priors in an algorithm's architecture or its training data, the easiest way to determine whether a basis is privileged is to consider the effect of a rotation on algorithmic performance.

In a non-privileged basis, rotation should have no impact on the effectiveness of the vectors in a model, because their features are not aligned with any specific basis dimension.

In contrast, in a privileged basis, such as in a CNN or transformer MLP, the basis dimensions are designed to capture specific types of features, and there is an incentive for the learned features to align with these dimensions. Rotating the basis would change the orientation of the basis dimensions relative to the learned features. In almost all cases, this misalignment will degrade the model's performance.

For example, a tiger-or-not-tiger model might contain a filter which is trained to capture orange stripes in the input data. Rotating this filter to detect blue stripes instead would make this feature much less useful!

It's useful to consider the kind of basis you're working with when training or using a new machine learning model.

If you're using a pre-trained network to extract features from data, knowing whether it will produce vectors in a privileged / non-privileged basis should make its results much more useful and interpretable.

When training models for tasks where interpretability is key, or where specific features are expected to be strong predictors, a privileged basis with strong architectural inductive biases may be appropriate. In contrast, for tasks where the features are less well-defined or where robustness to noise and variations is more important, leaving models open to define a non-privileged basis themselves might be the better choice.

What's do we mean when we talk about privileged and non-privileged bases in neural networks? How is a basis chosen or trained, and what can we say about their features?

Privileged vs non-privileged bases in machine learning

I really enjoy working with Jupyter notebooks. I used to exclusively use the JupyterLab UI, but lately I've been doing more of my notebook writing in VSCode. However, I've been missing some of the keybindings from JupyterLab which sped up my workflow, like the ability to merge cells with ⇧m. By default, VSCode users can only do this by awkwardly right-clicking the selected cells and choosing "Join selected cells".

However, VSCode is super configurable and makes it very easy to add that shortcut back to your workflow. Here's how:

Search for Preferences: Open Keyboard Shortcuts (JSON)

Paste the following into keybindings.json

[
    {
        // merge selected cells in a jupyter notebook with ⇧m
        "key": "shift+m",
        "command": "notebook.cell.joinSelected",
        "when": "notebookEditorFocused && !notebookEditorHasMultipleKernels && !inputFocus"
    }
]

A custom keybinding to merge jupyter notebook cells in vscode

For state-of-the-art language models, context is important.

Polysemy, homography, and other complexities of real-world language mean that a static definition of a word is often not enough to understand its meaning; The ability to interpret a word according to it's context is crucial for any high performing language model.

For example, the meaning of the word "get" changes radically between the phrases "I’ll get a book" (where get means fetch), "I'll get smarter" (where get means become), and "I get it" (where get means understand). 

The meaning of words and phrases isn't just affected by the bag of words around them, but also by where they are positioned in relation to each other. For example, the individual words in these sentences are the same, but modifying their order changes the meaning of the sentence dramatically:

"She only told him that she loved him."

Without the additional information about the position of each token, a model is limited to a bag-of-words representation of the target text. It's possible to learn something about a sentence from the words alone, but understanding all of the subtleties of language with a bag-of-words model is extremely difficult.

Encoding context in prior language models

In a generation of network architectures before transformers, language models like recurrent neural networks (RNNs) and long short-term memory networks (LSTMs) used a state-machine approach to encode positional context into each token. The models stored a representation of the previous tokens’ aggregate meaning, and used a combination of a fixed word embedding and its context to represent each word, updating the context vector(s) as it progressed through the text. Bidirectional variants of these models allowed context to be passed to a token from before and after its position in the text.

However, this state-machine approach is tricky to parallelise, and the length of the context window is limited as a result. State-machine approaches also exhibit a strong bias towards recent tokens and struggle to retain information over a long distance, even within the available context window. This meant that while state-of-the-art models were able to generate locally coherent text, the meaning of longer passages would quickly fall apart.

The mechanism at the heart of transformer models like BERT and GPT is intended to solve the parallelism problem, and increase models' coherence over longer sequences.

Instead of trying to hold and update the meaning of a sequence for each token, these models split the representation of words and their position into two distinct inputs. In addition to providing a word embedding for each token, we also provide a positional embedding, which represents the position of each token in the sequence. Transformers join those inputs together with an attention matrix, which describes how dependent or related each token is to its neighbours. Removing the sequential dependence and making the dependence of each token on its neighbours explicit makes the linear algebra significantly easier to parallelise, and allows transformers to scale much further than LSTMs or RNNs.

In this context, word embeddings work the same way as they have in many generations of language models, so I won't characterise them in much more detail in this post. Positional embeddings are a newer idea though, and warrant some more elaboration and intuition-building.

Unlike word embeddings, the characteristics of positional embeddings aren't emergent from data - they can be deliberately constructed to have certain properties. Concisely, positional embeddings should:

Be reflective of their distances from one another in the sequence - we should be able to use a distance metric to compare vectors, with positionally similar vectors producing small distances, and positionally distant vectors producing large distances in the vector space we construct.

Have values bounded between 0 and 1 - we’re using these as inputs to a neural network, so well-bounded values are useful.

Be invariant to sequence length - the values assigned to each position should be the same, whether the text is long or short.

Be deterministic - they shouldn't change from one initialisation of a network to another.

Our goal is to design a process which delivers embeddings with these characteristics. It’s worth pausing at this point to imagine a few candidate approaches yourself, and checking whether they fulfil the criteria above.

As a super-simple first stab, let’s just fill a vector with the position’s value, and call it a positional embedding. Position 1 gets an array of 1s, position 2 gets an array of 2s, etc:

def positional_embedding(position, dimensionality):
    return np.full(
        shape=dimensionality,
        fill_value=position
    )

We can certainly find the distances between positional embeddings which are generated this way; all we need to do is subtract the vectors from one another to determine their relative positions. However, the values aren’t bound between 0-1. 

Let’s try dividing the values by the length of the sequence to produce each token’s fractional position within it.

def positional_embedding(position, dimensionality, sequence_length):
    return np.full(
        shape=dimensionality,
        fill_value=position/sequence_length
    )

Hm, still no good. Dividing the values by the sequence length gives us a set of values which are easily comparable and bound between 0-1, but by their nature, they’re not invariant to sequence length. The first token in a sequence of 30 words will have a different set of values to one where the sequence is only 5 words long, and there’s no way to know the lengths of the sequences we’ll be modelling ahead of time. Ideally, we'd like this to scale to sequences of any length.

The approach in the original transformers paper (“Attention is all you need”) is really fun and smart, and it fulfils all of the criteria we defined.

First, we specify a dimensionality for the embedding. Each dimension of our embeddings will be assigned a corresponding sinusoidal wave, each with a different frequency. Early dimensions of the embedding will be tied to low-frequency waves, and the frequencies will increase as the we move towards the end of the embedding. We also alternate between sine and cosine waves, depending on the element’s parity. To get a positional embedding, we plug our position value into that series of waves, and read off the resulting values.

def positional_embedding(position: int, d_model: int = 1024):
    i = np.arange(d_model)
    angles = position / np.power(10000, (2 * (i // 2) / d_model))
    angles[0::2] = np.sin(angles[0::2])
    angles[1::2] = np.cos(angles[1::2])
    return angles

This is the confusing bit - why does that process produce embeddings which fulfil our criteria??

First, using sine and cosine functions means that we’re bounding our values between 0-1, no matter how long our sequence length. We can also see that the process is deterministic, and doesn't depend on any stochastic/external parameters. But what about the other criteria? How do these rules produce an embedding space with values which can neatly represent a continuous number line?

Let’s illustrate the simplest possible example of the method. Here are two cosine waves, one with a low frequency, and one with a high frequency. Let’s imagine that each one is tied to a dimension in a two-dimensional positional embedding.

Now let’s select two nearby positions, and mark them in our space with a pair of vertical lines.

We can use our position values (aka vertical lines) to get the first elements in our positional embeddings. The height at which our low-frequency curve meets our vertical lines is what we'll use as the first element in our embedding. I've marked those heights with horizontal lines on the graph below.

The resulting values are very close together, because we're at a slowly changing point on our low-frequency curve, with a very shallow gradient. 

We can do the same thing with our second, higher-frequency curve to get the second elements of our 2D embeddings.

We're at a point on the curve with a much steeper gradient, which means that the values are much further apart! The higher frequency wave also changes gradient much more rapidly, so we can expect more variability in the values we read from it. At higher frequencies, this would be even more true.

Now let’s consider a new line, representing a more distant position:

Both values here are very different to our first two positions. Next, try to visualise the values for a position close to our new, third line. You should find that they're much more similar to the values for our third position than they are to our first two positions.

Now, imagine extending our example to include many more curves, with many different frequencies, representing a positional embedding with a much higher dimensionality. As we accumulate more values, correlations between similar positions (whose values are close at slowly-changing points on our curves) will become more apparent, as will the lack of correlation between positions which are positionally further apart (which are unlikely to share slowly-changing regions of curves). 

With increasing dimensionality, these correlated and uncorrelated dimensions start resembling a pretty neat vector space. We can use the cosine distance between points to get a measure of the distance between them, achieving our final criteria!

It's worth considering what might happen if we were to use sin or cosine waves alone, and the strange results that it might give us. At low-frequencies, we would see regions where neighbouring waves would all be at their steepest, with the corresponding values changing a lot, and other regions where their values would all be very similar and hard to distinguish. By alternating between sin and cosine curves across the elements of our embedding, we ensure that in regions where the sine curves are changing a lot, the cosines won’t be, and vice-versa, thus stabilising our embeddings.

Examining distances between positional embeddings

Given a large enough embedding, we can use the cosine distance to measure the distance between positions in our vector space.

from scipy.spatial.distance import cosine

cosine(positional_embedding(1), positional_embedding(2))
>>> 0.026488616022189992

Positions which are further apart should also generate larger distances between their embeddings

cosine(positional_embedding(1), positional_embedding(3))
>>> 0.09339161307513

cosine(positional_embedding(1), positional_embedding(30))
>>> 0.4323030365719962

We should also see positions which are far away from position 1 but close to each other at similarly small distances from one another

cosine(positional_embedding(30), positional_embedding(31))
>>> 0.02648861602218988

We can also visualise the distribution of these distances. Here are all of the positions up to 250, plotted against their distance from position 1 in embedding space.

Here’s the distance-matrix version of that chart, showing the distances from all of the first 250 positions to one another. The data in the first row matches the values we've scattered on the plot above. 

We can also squeeze the dimensionality of our positional embeddings to plot the distances between all of them on a 2D plane.

from umap import UMAP

embeddings = np.array([positional_embedding(i) for i in range(1000)])
embedder = UMAP(n_components=2, metric="cosine")
embedder.fit(embeddings)
embeddings_2d = embedder.transform(embeddings)

We'd expect them to fall in a continuous line, with neighbouring positions closest to one another. Here they are, coloured according to their position:

The original paper uses the same dimensionality for both word and positional embeddings, and uses their sum as the input to the network. To be honest, I don't understand that decision. I would expect that concatenating the embeddings would inject more coherent information into the network. If you're interested in that experiment, let me know - I might write a follow-up post with the results!

Many networks use the positional embedding scheme described above, while some models like the RoFormer expand on the idea, more effectively leveraging the information that positional embeddings encode.

The authors of “Attention is all you need” note that they also experimented with fully learned embeddings, as well as the sinusoidally-generated ones we’ve explored here, with both approaches producing similar results. Many transformer implementations choose to forego implementing the positional embeddings, using a standard learned embedding layer instead.

Personally, I still love the idea of an embedding layer which efficiently encodes information without having to rely on learning, and I’ll be looking for areas in my own work where I can make use of those neat trigonometric tricks too!

All of the code used to generate the results and visualisations in this post is openly licensed and available in the corresponding github repo. If you got something from this post, please give it a star! And if I've made any silly mistakes, open an issue and let me know! 😅

Positional embeddings are key to the success of transformer models like BERT and GPT, but the way they work is often left unexplored. In this deep-dive, I want to break down the problem they're intended to solve and establish an intuitive feel for how they achieve it.

Understanding positional embeddings in transformer models

Approximate nearest neighbour (ANN) search is a foundational task in data science and machine learning, where the goal is to efficiently find data points in a dataset that are most similar to a given query point. It's used all over the place, from music recommendation, to plagiarism detection, to solving massive jigsaw puzzles. 

While there are many algorithms for ANN search, one of the most effective is locality sensitive hashing (LSH), which allows for fast and approximate search by mapping similar data points to the same "buckets".

In this blog post, I'll describe how to to implement a super fast, production-ready ANN search engine by training an LSH model, inferring hashes for new data, and indexing and querying those hashes with elasticsearch. Whether you're a seasoned data scientist or a newcomer to the field, I hope you'll learn something new from this post!

An ANN search should produce a set of data points in a dataset that are closest to a given query point. These methods are called "approximate" because they use clever tricks to find an approximate answer to the problem, rather than an exact answer.

Allowing the calculated distances to be imprecise (particularly for matches which are far away from the query point) allows approximate methods to work much faster than traditional, exact methods for finding the nearest neighbours, which can be very slow for large datasets.

Elasticsearch is a mature, open-source platform for building search engines, with a powerful, flexible query API. It uses some smart algorithmic magic behind the scenes to quickly and accurately find relevant results based on a user's search query.

Term frequency-inverse document frequency, or TF-IDF, is the core method used by Elasticsearch to assign scores to matching documents by determining how often each words overlaps with the user's search terms, and the significance of those overlaps based on the rest of the documents in the index. The more unique a matched word is to a particular document, the higher its score will be. This allows Elasticsearch to prioritize and rank search results based on the relevance of specific words to the user's query.

Locality sensitive hashing (LSH) is a way of finding similar items in large datasets by creating a "hash" or unique identifier for each item. Typically, hashing algorithms are designed to minimise collisions; that is, different input data should always result in a different hash, no matter how similar the inputs are. In contrast, LSH methods are designed to maximise collisions; items that are similar to each other should result in a similar set of strings, making it easy to find similar items without having to compare every single item in the dataset with an exact distance calculation.

LSH models work by using a group of independent hash functions to map similar data points in a high-dimensional feature space into the same buckets, so that nearby points are more likely to be hashed to the same bucket. This allows for fast and approximate search, since we only need to compare the query point to the points in the same set of buckets, rather than the entire dataset.

To use LSH, we first need to choose a set of hash functions that have the property of being "locality sensitive", meaning that they are more likely to hash similar points to the same bucket. There are many ways to construct such hash functions, but one of the most common is to use random projections. This involves projecting the data points onto a random vector, and then dividing the projection values by a chosen threshold to get the hash values. Alternatively, we could divide the complete feature space into a set of subspaces by randomly choosing groups of features, and then apply a clustering model (eg k-means or HDBSCAN) to each one, defining regions which share similar features.

We can combine the ideas we've defined above to establish a new approximate nearest neighbour method: If we use an LSH model to produce a set of tokens which describe our data (in a way that maximises the chance of collisions for similar sets of features), we can then use TF-IDF to search and compare those hashes, returning the most relevant ones! In particular, elasticsearch's more_like_this query type will allow us to use a field from an existing document as the set of search terms for a new query, allowing us to make the ANN request without knowing the query document's features or hashes in advance.

Let's get started by creating an LSH model to hash our features. The inputs to our model should be set of arrays which precisely describe the source data - If we could use an exact NN method, these would be the features we'd use to calculate our distances. 
For example, if we were building an image similarity service, these feature vectors might be the penultimate layer of of activations from a VGG16 model. Alternatively, working with text to produce article recommendations, they might be BERT sentence-embeddings. For simplicity, I'll just define a set of random 256-dimensional vectors in this example.

import numpy as np

# an array of 10,000 256-dimensional vectors
features = np.random.rand(10_000, 256)

We're going to randomly select a set of 1,000 feature vectors to train with. We only have 10,000 total data points here, but in production settings we might need to train on a subset of the complete data. As long as the distribution of the training set matches that of the test set, this sampling method should be fine.

After that, we split each of our training examples into 32 chunks. For each original 256 dimensional array, we'll end up with 32 8-dimensional arrays, each representing a different set of features and acting as its own subspace. Importantly, these smaller subspaces can be compared just like the original feature set.

# a subset of 1000 of those vectors
n_training_vectors = 1000
random_indices = np.random.choice(
    features.shape[0], n_training_vectors, replace=False
)
training_features = features[random_indices]

def split_features(feature_vectors, n_groups):
    return np.split(feature_vectors, indices_or_sections=n_groups, axis=1)

n_groups = 32
feature_groups = split_features(training_features, n_groups)

Next we'll fit a k-means model to each of those subspaces, dividing each one into 16 regions based on the clusters of features that it finds. 

from sklearn.cluster import KMeans

def train_clusters(feature_group, m):
    clustering_alg = KMeans(n_clusters=m, n_jobs=-1).fit(feature_group)
    return clustering_alg

m = 16
models = [train_clusters(feature_group, m) for feature_group in feature_groups]


Taken in combination, this set of models is our complete LSH model! If we chunk every one of our original feature vectors the same way as we did for the training set, we can predict which of our buckets each chunk falls into, and thereby generate a meaningful hash!

def hash_features(feature_vector, models):
    feature_groups = split_features(feature_vector.reshape(1, -1), len(models))
    clusters = [
        model.predict(feature_group)
        for model, feature_group in zip(models, feature_groups)
    ]
    return [f"{i}-{j}" for i, j in enumerate(clusters)]

hashes = [
    hash_features(feature_vector, models)
    for feature_vector in features
]

Here, each token (or 'word') in the hash is comprised of two parts: the index of the feature group / subspace we're considering, and the index of the cluster which that datapoint has been classified into. This gives us a set of n x m possible tokens. In this case, we have 32 subspaces, each divided into 16 clusters, so each document will be given a set of 32 tokens from of 512 possible tokens. Feature vectors which are close in the original feature space should share more of those tokens with one another than feature vectors which are far apart.

In production settings, the dataset is likely to be higher-dimensional, and we're likely to have many more than 10,000 documents. At the scale of ~10,000 docs, computing the distances exactly might be sufficiently fast, without any need for LSH. A much more information-rich feature space like those found in the real world will probably require higher values of both n and m than the values we've used here.

Now all we need to do is index our hashes into an elasticsearch index, ready to be queried. We'll set the field type of the lsh field to keyword, and set the analyzer to whitespace so that our hashes remain intact and are not split apart by the hyphen in {i}-{j}.

from elasticsearch import Elasticsearch

# Create a client instance
es = Elasticsearch()

# Define the index settings and mapping
index_settings = {
    "settings": {
        "number_of_shards": 1,
        "number_of_replicas": 0
    },
    "mappings": {
        "properties": {
            "lsh": {
                "type": "keyword",
                "analyzer": "whitespace"
            },
        }
    }
}

# Create the index
es.indices.create(index="ann-search", **index_settings)


We can now index each of our documents into that newly defined mapping:

for i, document_hashes in enumerate(hashes):
    es.index(
        index="ann-search",
        id=i,
        document={"lsh": document_hashes}
    )

And finally, run a more_like_this query to see similar results! We use a random ID as the query, but we could use a specific one based on user requests if this were running in a production API accepting user queries, for example.

random_id = np.random.randint(0, len(hashes))
results = es.search(
    index="ann-search",
    query={
        "more_like_this": {
            "fields": ["lsh_features"],
            "like": {
                "_index": "ann-search",
                "_id": random_id
            },
            "min_term_freq": 1,
            "max_query_terms": 32
        }
    }
)

approximate_nearest_neighbours = [
    hit["_id"] for hit in results["hits"]["hits"]
]

Because of elasticsearch's powerful, flexible query API, we can also run more interesting searches! For our images example, if we also had additional text fields like titles or descriptions for each image, we could combine our nearest-neighbour queries with traditional searches in a single API call, eg "give me images whose features look like this one, but whose captions contain the word 'dog'". We could even work in a multi-modal way, and apply the same LSH process to features from the text as well as the image, querying both fields for documents which share text and image features.

This fundamental pattern is one I use all the time: working with a set of large, interesting, highly optimisable ML models behind the scenes with a simple, interpretable format for inferred data, subsequently stored in an index with powerful querying capabilities.

In a follow-up post, I'll walk through applying these ideas to a large image dataset, and architecting a system for deploying the resulting models and APIs in the real world for other people to use!

This novel approach to approximate nearest neighbour search achieves super-fast query times by using elasticsearch to index and query the outputs of an LSH model. The cheap, scalable architecture for both hashing and querying make it an ideal approach for huge datasets in production environments.

Super fast approximate nearest neighbour search with locality sensitive hashing and elasticsearch

If you're getting started with data science work and python, you might have come across the yield keyword and wondered how or why it's used. Once you've got the hang of it, using yield can be a powerful technique for working efficiently with large datasets, without using up all your machine's memory.

Let's say we have a huge dataset that we need to work through, line by line. To start, we might read the entire dataset into memory, and then loop through each line of the dataset to process it.

with open("dataset.txt", "r") as f:
    dataset = f.readlines()

for line in dataset:
    # Do something with each line of the dataset

This works well for small datasets, but it's not the most efficient approach. We're storing a complete copy of the dataset, despite only processing one line at a time! If our dataset is large, loading the whole thing ahead of time can quickly use up all of the available memory and crash our program.

That's where yield comes in. yield allows us to create a generator object which only produces the next item in the dataset when it's needed, instead of reading the entire thing into memory at once.

Here's an example of how we might adapt the code above to use yield:

def iterate_dataset():
    with open("dataset.txt", "r") as dataset:
        for line in dataset:
            yield line

for line in iterate_dataset():
    # Do something with each line of the dataset

Instead of loading the whole dataset into memory at once, we're loading each line only when we're ready to process it, leaving us much more memory available for processing. 

Sometimes you might want to access a single item in a generator, without looping through the whole thing; to access the next item, we can use python's built-in next and iter functions:

data_generator = iterate_dataset()
next_item = next(iter(data_generator))

Being able to combine these techniques is super powerful, and makes lots of otherwise tricky work very simple!

However, yield might not be the best tool in every situation. For example, it's worth remembering that using yield will create a one-time iterator. That is, because it generates the next item in the dataset only when it's needed, you can only iterate over your dataset once. Once you've iterated over the entire dataset, the iterator is exhausted and you won't be able to do so again unless you recreate it. For the same reason, you can't index into a generator in the same way as you would a regular list.

I hope you've found this post helpful, and you feel ready to experiment with yield in your own code!

Generators come to the rescue when your dataset is too big for your machine

Iterating over large datasets with "yield" in python

Coming from a background in physics, I used to write my CV in LaTeX, the typesetting tool of choice for dorks everywhere. But in the years since, I'd grown frustrated with the ugly language, the fiddly and never-ending typesetting tweaks, and the archaic versioning, compiling, and publishing patterns.

I decided to scrap the LaTeX version, and moved the content into a framework which matched the rest of this website. HTML, CSS, and typescript might have their own problems, but at least I had some use for them in other areas of my life.

You can see the result here, which I was very happy with!
It's clean, concise, fast, and accessible. The content is all written and stored in prismic, and is then structured and styled for my site by next.js and tailwind CSS respectively. The design is responsive, so it's readable on any device. Updating the CV is now just as easy as editing or publishing a new blog post, and any improvements I make to my site also carry over to the CV. Making it available online as rich, indexable HTML also makes it much easier to find by people searching for me or my skillset!

However, for one reason or another, CVs are one of the last documents in my life which people insist must conform to a standard printable paper size. Application forms almost always requires a PDF copy of a CV, so without a PDF option, a pretty page of HTML isn't enough.

Rather than maintaining two versions of the same content, I wanted to generate a PDF version of the CV using the HTML on the site.

The simplest change I could make was to add a link at the bottom of the page, which, when clicked, would instruct the user's browser to open the print dialog, from which they would be able to save the page as a pdf.

<a href="javascript:window.print()">Download this as a PDF</a>


The best way to encode dates, times, and other cyclical features

Privileged vs non-privileged bases in machine learning

A custom keybinding to merge jupyter notebook cells in vscode

Understanding positional embeddings in transformer models

Super fast approximate nearest neighbour search with locality sensitive hashing and elasticsearch

Iterating over large datasets with "yield" in python

Creating a downloadable .pdf copy of a page using next.js and puppeteer

You should commit your .env files to version control (carefully)

A CLI alias for quick and easy passwords

Better terminal tab titles

Automated testing with Vercel previews and GitHub actions

Using emoji as favicons

Anti-racism

def from above