Coala Embeddings

A 20 million document index takes 80 GB of memory. Coala can fit it into under 1 GB—retaining over 94% search quality.

Coala Embeddings is a compression framework that shrinks embeddings by up to 99% by training a lightweight neural network along with quantization. It enables compression of any existing embedding model in seconds, without sacrificing retrieval accuracy.

The main features are

• Fast search: Even on low end CPUs, the speed can reach over 1000 queries per second.
• Excellent Quality: Our novel neural network and losses are up to 27% better quality than previous SOTA.
• Lightweight code: The repository only has around 1k lines of code and easy to install using pip.
• Extreme RAM savings: can compress embeddings up to 256x, saving significantly more RAM than just pure quantization (TurboQuant, RabitQ, etc.)

Usage

Simple 1 line installation

pip install coala_embeddings

Simple 3 line usage

from coala_embeddings.manager import ProposedIndex

# 1. Initialize with a sample of your dataset (4,000+ documents recommended)
index = ProposedIndex(train_embeddings=train_data, directory="document_embeddings")

# 2. Add all your document embeddings to the compressed index
index.add_documents(documents)

# 3. Search to get the top 10 most relevant document IDs
scores, doc_ids = index.search(query, fine_k=10)

Note: Train_data, documents, query should be formatted as a numpy array with shape like this (num_documents, dimension). Documents/Network are automatically saved to directory.

Loading documents

from coala_embeddings.manager import ProposedIndex
index = ProposedIndex(load=True, directory="document_embeddings")

## If you don't want to load documents but load the neural network compresser only, then do the below
## encoder_path = 'document_embeddings/encoder.pkl' ## Make sure to change the path if it's not correct.
## index = ProposedIndex(encoder_path=encoder_path, load=True, directory="new_doc_embeddings", scratch_doc=True)

Advanced Usage

Usage with description of inputs

import numpy as np
from coala_embeddings.manager import ProposedIndex

# 1. Initialize with GPU support and custom dimensions
index = ProposedIndex(
    train_embeddings=train_data,  # Pre-training sample (~4k+ vectors)
    original_dim=1024,            # Input vector dimension
    slow_dim=256,                 # Re-ranking stage dimension(stored on disk)
    fast_dim=32,                  # Fast IVF candidate generation dimension(stored on memory)
    device='cpu',                 # Acceleration ('cuda' or 'cpu')
    probes=256,                   # Centroids for IVF index (increase for 1M+ docs)
    directory="document_embeddings",       # Path to save/load index binaries
    num_threads=8                 # CPU threads for FAISS
)

# 2. Add full document corpus
index.add_documents(documents)

# 3. Fine-tune two-stage search performance
scores, doc_ids = index.search(
    query,
    coarse_k=200,  # Number of candidates retrieved in Stage 1 (Fast)
    fine_k=10,     # Final re-ranked results returned in Stage 2 (Slow)
    nprobe=16      # IVF centroids to visit during search. Higher = better quality but slower.
)

Load and search

from coala_embeddings.manager import ProposedIndex

# Reload existing index files directly from disk
index = ProposedIndex(
    load=True,
    directory="document_embeddings",
    device='cuda'
)

# Ready to search immediately
scores, doc_ids = index.search(query, fine_k=10, nprobe=16)

Tips

• RAM Management for Indexing: Build your index on a high-RAM machine. Indexing operations (like IVF clustering and quantization) require significantly more RAM during setup than during search. Once built, you can easily load and run it on a low-RAM device.
• Hardware Acceleration: Utilize a GPU for the initial training phase. For loading and searching, CPU execution is fine and often faster than transferring data back and forth to a GPU per query.
• Training Corpus Size: Make sure you provide at least 4,000 document vectors in your training corpus (train_embeddings) for optimal compression quality.
• Hands-on Examples: Check out the notebooks/ directory for step-by-step real world examples using popular models like mxbai-embed-large and nomic-embed-text-v1.5.

Frequently Asked Questions

Q: How does this work?

Modern embedding models produce high-dimensional vectors (768 to 2048+ dimensions). We use a lightweight neural network trained with novel loss functions to project these embeddings into two paths:

• Low-dim path (32-64 dim): Stored in memory to perform an ultra-fast coarse search (for example, fetching top 200 candidates). The loss used for training optimizes for coarse recall.
• High-dim path (~256 dim): Stored on disk to re-rank those candidates into the final top 10 results. The loss used for training optimizes fine retrieval.

Both paths are further compressed using 4-bit quantization.

Q: How does it achieve 64x memory savings?

By keeping only the ultra-compressed low-dimensional index in RAM while leaving the higher-dimensional re-ranking index on disk. Because disk reads are only triggered for the top candidate results, you get massive memory savings without a performance issue.

Q: How does it compare to SVD, MRL, or other techniques?

It significantly outperforms them. Traditional post-hoc methods like SVD degrade quickly at extreme compression levels (e.g., 32 dimensions), and MRL requires expensive model re-training. Coala delivers superior retrieval quality post-hoc on your existing models in seconds. You can see this in the next section.

Benchmarks

As seen in the image, our method (colored green) is consistently #1, infact up to 27% better then SVD. The method is very close to no-compression (colored black), reaching 95% of it's quality.

View full benchmarks

Roadmap

• Improve efficiency and memory.
• Huggingface space demo with real world model.
• Release paper for technique and more benchmarks.

Acknowledgments

• Faiss: Very useful to create optimized IVF indexes and quantize code.
• Beir: Simple and effective benchmarking code.
• Matryoshoka Paper: Useful paper with several ideas used here.

Final Notes

The code is licensed under the Apache-2.0 license. See LICENSE for details.

Stars would be appreciated, thank you.

Email: yatharthsharma3501@gmail.com