Published at Oct 30, 2024

Implementing Cosine Similarity in Python

AI development for all developers, not just AI experts. Build your AI app with Tiger Cloud today.

Embeddings map language to vectors, while cosine similarity measures relationships between them, enabling various artificial intelligence (AI) applications. Cosine similarity calculates the angle between two vectors—similar vectors score 1, while dissimilar ones score 0. We explored the topic in depth in a previous article and explained how cosine similarity is used in AI systems.

This time, we will focus on implementing cosine similarity from scratch using Python. You will be able to touch the insides of the computation, locate bottlenecks in the algorithm, and look for essential tool optimizations to improve its performance. Let’s jump right in.

Basic Cosine Similarity Implementation in Python

Cosine similarity measures the similarity between two non-zero vectors by calculating the cosine of the angle between them in vector space. Simply put, it tells us how close two vectors are, regardless of their magnitude.

Cosine similarity formula

The formula for cosine similarity between two vectors, X and Y, is:

Here, X and Y represent two vectors in an n-dimensional space, meaning they are numerical arrays of the same length, such as:

The cosine similarity is computed by taking the inner product of X and Y, and dividing it by the product of their magnitudes.

Inner product function

The first step is calculating the inner product (the dot product) between the two vectors. This is done by multiplying each pair of corresponding elements from the two vectors and summing the results.

Here’s the Python function to compute the inner product:

def inner_prod(X, Y): sum_prod = 0 for x_i, y_i in zip(X, Y): sum_prod += x_i * y_i return sum_prod

This function loops over the entries of X and Y, multiplies each pair 𝑥_i and 𝑦_i, and adds them to sum_prod, which is returned at the end.

Magnitude function

The magnitude (Euclidean norm) of a vector X is calculated by squaring each element, summing them, and taking the square root of the sum. The same calculation is applied to vector Y.

Here’s the Python function to compute the magnitude of a vector:

import math def magnitude(X): sum_squares = 0 for x_i in X: sum_squares += x_i ** 2 return math.sqrt(sum_squares)

In this function, sum_squares accumulates the sum of squares of all elements in X, and then the square root is taken to return the final magnitude.

Complete cosine similarity code

Now, we combine the inner product and magnitude functions to calculate cosine similarity:

def cosine_similarity(X, Y): return inner_prod(X, Y) / (magnitude(X) * magnitude(Y))

This function simply calls the inner product and magnitude functions and returns the cosine similarity by dividing the dot product by the product of the magnitudes of X and Y.

Algorithm analysis

For vectors with 𝑛 entries:

The inner product computation requires 2𝑛 operations (one multiplication and one addition per element).
The magnitude computation for each vector involves 2𝑛 + 1 operations (squaring, summing, and taking the square root).
The algorithm's total complexity is O(6n), which means it scales linearly with the vector length or “dimension.”

This makes cosine similarity computation highly efficient for moderate dimensions. However, AI systems often work with vector embeddings with dimensions in the millions or even billions of numerical values. As a result, the computational cost increases significantly, primarily when cosine similarity is used in pairwise searches across large datasets.

Search With Cosine Similarity

In practice, cosine similarity is often used to search a database of vectors to find those most similar to a given query vector. The goal is to rank the vectors by their cosine similarity, where the results closer to 1 are more similar to the query vector, and those closer to -1 are dissimilar.

For instance, suppose we have a database of vectors Vector_db = [Y_0, Y_1, ..., Y_(m-1)] and a query vector X. We can perform a pairwise search by calculating the cosine similarity between X and each vector 𝑦_i in the database. The vectors can then be sorted in descending order based on their similarity to the query vector.

Naive pairwise search code

This function performs a pairwise search by calculating cosine similarity between a query vector and each vector in a database, returning the results sorted by similarity in descending order. Here’s a naive implementation of a pairwise search with cosine similarity:

def search_db(X, vector_db): results = [] for Y_i in vector_db: similarity = cosine_similarity(X, Y_i) results.append((Y_i, similarity)) # Sort by similarity in descending order results.sort(key=lambda x: x[1], reverse=True) return results

In this code, the function search_db takes a query vector X and a database of vectors vector_db. For each vector 𝑦_i in the database, we calculate the cosine similarity using our previously defined cosine_similarity function. The results are appended to a list sorted by the similarity score in descending order.

Explanation of inefficiencies

While this method works for small datasets, it becomes inefficient as the vector database grows. Here are the main issues with naive pairwise search code:

Computational complexity: The naive approach requires calculating cosine similarity 𝑚 times, where 𝑚 is the number of vectors in the database. This alone makes it scale poorly with large datasets.
Dimensional scaling: For each similarity computation, we must consider the dimensionality of the vectors. High-dimensional vectors (common in AI applications) require more time to compute their cosine similarity.
Sorting overhead: After calculating the similarities, the function creates and sorts an array of length 𝑚 to produce ranked results. Sorting further adds computational overhead.

This naive approach becomes impractical for large databases with high-dimensional vectors. The sheer volume of computations makes it a severe bottleneck for performance.

Nearest neighbors search algorithms

The process of ranking vectors based on similarity is called nearest neighbors (NN) search. Various algorithms and techniques have been developed to improve the efficiency of similarity searches. These methods optimize the search by adding extra structures to the vector data and reducing the number of complete pairwise comparisons.

HNSW (hierarchical navigable small world)

HNSW is an algorithm designed to efficiently search large vector databases. It organizes the data into multiple layers, starting with a sparse layer containing only a few vectors. The algorithm searches through the layers, moving in the direction of the query vector, from lower to higher similarity.

How it works: Starting at the sparse top layer, the algorithm moves through each layer to get closer to the most similar vectors. The search progressively refines itself as it navigates through denser layers.
Performance: This method significantly speeds up the search process while maintaining high accuracy.

More information on HNSW can be found here.

Centroid indexing

Centroid indexing is another optimization technique for searching for nearest neighbors. In this approach, the vectors in the database are clustered, and each cluster's centroid (average vector) is calculated.

How it works: Instead of comparing the query vector against all database vectors, the algorithm first finds the most similar centroid. It then searches only within that cluster for the closest vectors.
Performance: This method improves efficiency for large datasets by reducing the number of vectors that need to be searched.

A detailed explanation of centroid indexing can be found here.

Approximate nearest neighbor (ANN) methods

ANN methods are designed to significantly speed up similarity computations by sacrificing a small degree of accuracy. These methods use statistical techniques to approximate the nearest neighbors, trading off perfect accuracy for faster search times.

How it works: Rather than doing an exhaustive search, ANN algorithms use approximations that may mis-rank some vectors but will return the most relevant results for most cases.
Example: Microsoft’s DiskANN, an ANN algorithm, provides substantial performance improvements while only losing about five percent in accuracy. It is optimized to run on standard hardware, like 64 GB of RAM and SSD storage, making it scalable for massive datasets.

You can learn more about DiskANN here.

Cosine Similarity Optimization

Optimization techniques can significantly improve the performance of cosine similarity calculations when dealing with high-dimensional vector spaces. One key area for optimization is using sparse vectors—vectors where many entries are zero. For vectors with billions of dimensions, most entries are often zero, and ignoring them in computations can drastically reduce the workload.

Sparse vector optimization

High-dimensional AI and machine learning vectors often contain many zero entries. These are known as sparse vectors. Zero entries contribute nothing to the inner product's final result and cosine similarity magnitude calculations. So, instead of performing computations on every entry in the vector, we can optimize by focusing only on the non-zero elements.

Non-zero entry optimization

The idea behind optimizing cosine similarity for sparse vectors is first to identify the vectors' non-zero indices. Then, rather than wasting resources computing with zeros, we compute the inner product and magnitude only on the index that matters.

Here’s a function that returns the indices of non-zero elements in a vector:

def non_zero(X): return [i for i, x in enumerate(X) if x != 0]

Using this non_zero function, we can optimize our inner product calculation by computing only the product of the non-zero elements from both vectors.

def optimized_inner_prod(X, Y): non_zero_X = non_zero(X) non_zero_Y = non_zero(Y) common_indices = set(non_zero_X).intersection(set(non_zero_Y)) inner_product = sum(X[i] * Y[i] for i in common_indices) return inner_product

In this code, optimized_inner_prod calculates the inner product of two vectors, but only for non-zero indices in both vectors. This calculation avoids unnecessary multiplications involving zeros, making it highly efficient for sparse data.

Storage optimization with sparse vectors

For extremely high-dimensional data systems, optimizing storage becomes crucial. Instead of storing the entire vector—mostly filled with zeros—we can only store the non-zero entries and their corresponding indices. This optimization saves memory and speeds up computation.

For example, a sparse vector can be stored like this:

sparse_vector = { "indices": [0, 3, 7], "values": [0.5, -1.2, 3.1] }We can expand the full vector by inserting zeros in the appropriate places when we need it. Storing the non-zero elements in this compressed format can significantly reduce the size of your vector database and improve search times.

Benefits of sparse vector optimization

Reduced computational complexity: We skip zero elements, reducing the number of calculations for the inner product and magnitude.
Efficient storage: Storing only the non-zero entries of sparse vectors can reduce memory requirements, making the storage of large datasets of high-dimensional vectors feasible.
Scalable search: Sparse vectors naturally arise in AI systems that deal with high-dimensional data. Hence, incorporating optimizations like non_zero computations into your data structure ensures scalability, even for billions of entries.

Focusing on the non-zero elements of sparse vectors makes the cosine similarity calculation and vector storage much more efficient, making it feasible to work with high-dimensional data at scale. This is a crucial optimization for AI applications where vectors often have sparse representations.

Cosine Distance Tools

The implementation above still requires further optimizations to be suitable for a production environment. Fortunately, numerous tools are available with built-in cosine similarity functions optimized explicitly for scalability and production use.

Pgvector: Cosine distance in PostgreSQL

Pgvector is an extension designed to store and search high-dimensional vector data within PostgreSQL. It includes a variety of vector distance metrics, including cosine distance, Euclidean distance, and L² distance. Here are some others:

Symbol	Distance Metric
<#>	Negative Inner Product
<~>	Hamming Distance
<%>	Jaccard Distance

In pgvector, cosine distance is defined as:

This metric ranks vector pairs based on how close they are in direction (rather than magnitude), with smaller values representing more similar results. In search queries, vectors are ranked from low to high, with the lowest cosine distance corresponding to the most similar items.

How to enhance cosine similarity in PostgreSQL

To turn PostgreSQL into a full-fledged platform for AI development that is accessible to all developers, Timescale created an open-source PostgreSQL stack with purpose-built extensions for AI that complement pgvector:

pgai brings more AI workflows to PostgreSQL, making it easier for developers to build search and retrieval-augmented generation (RAG) applications.
pgvectorscale enables developers to build more scalable AI applications with higher-performance embedding search and cost-efficient storage. It adds a StreamingDiskANN index to pgvector, as well as Statistical Binary Quantization, unlocking large-scale, high-performance AI use cases previously only achievable with specialized vector databases like Pinecone. Compared to Pinecone’s storage-optimized index (s1), PostgreSQL with pgvector and pgvectorscale achieves 28x lower p95 latency and 16x higher query throughput.
pgai Vectorizer enables embedding creation within PostgreSQL—no more complicated vector databases or clunky add-ons. Stale embeddings are a thing of the past, and you can test different embedding models and chunking/formatting strategies within PostgreSQL. Compare, tweak, and refine—all in one place.

Conclusion

Operators like cosine similarity serve as an essential foundation of semantic search. Optimized and efficient cosine similarity computation is critical in a vector-based AI system. Cosine similarity-based nearest-neighbor search opens up many possibilities for AI applications—even the most popular ones, like RAG.

With Timescale’s open-source stack for AI applications, you can leverage state-of-the-arnearest-neighboror algorithms, combining cosine similarity with other similarity metrics to power your AI tools.

Timescale made PostgreSQL the powerful database you trust—now it’s your AI platform, too. Pick your deployment option today: you can find self-hosting installation instructions on the pgai, pgvectorscale, and pgai Vectorizer GitHub repositories or quickly access all these extensions on any database service on Timescale’s cloud PostgreSQL platform.