xbe.at

Miscellany / Embeddings and vectors

Exploring Embeddings and Vector Search

9 min read, 15 sections

Contents
  1. Understanding Embeddings Fundamentals
  2. Vector Similarity Metrics
  3. Building a Basic Vector Store
  4. Implementing Approximate Nearest Neighbors
  5. Text Preprocessing for Embeddings
  6. Document Embedding Pipeline
  7. Vector Indexing Strategies
  8. Real-world Application - Semantic Search Engine
  9. Mathematical Foundations of Embeddings
  10. Dimensionality Reduction for Visualization
  11. Embedding Quality Assessment
  12. Production-Ready Vector Store Implementation
  13. Performance Optimization Techniques
  14. Results Analysis and Metrics
  15. Additional Resources

Understanding Embeddings Fundamentals

Embeddings are dense vector representations of data that capture semantic relationships and meaningful features. They transform high-dimensional categorical data into continuous vector spaces where similar items are positioned closer together, enabling efficient similarity computations.

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer

# Example text data
texts = [
    "Machine learning is fascinating",
    "Deep learning revolutionizes AI",
    "Neural networks process data"
]

# Create TF-IDF embeddings
vectorizer = TfidfVectorizer()
embeddings = vectorizer.fit_transform(texts)

# Convert to dense array for visualization
dense_embeddings = embeddings.toarray()
print(f"Embedding shape: {dense_embeddings.shape}")
print(f"First document embedding:\n{dense_embeddings[0]}")

Vector Similarity Metrics

Vector similarity computations form the foundation of embedding-based search systems. The choice of similarity metric significantly impacts search quality and performance. Common metrics include cosine similarity, euclidean distance, and dot product.

def cosine_similarity(v1, v2):
    # Compute cosine similarity between two vectors
    dot_product = np.dot(v1, v2)
    norm_v1 = np.linalg.norm(v1)
    norm_v2 = np.linalg.norm(v2)
    return dot_product / (norm_v1 * norm_v2)

# Calculate similarity between first two documents
sim_score = cosine_similarity(dense_embeddings[0], dense_embeddings[1])
print(f"Cosine similarity: {sim_score:.4f}")

Building a Basic Vector Store

A vector store is a specialized database designed to efficiently index and search high-dimensional vectors. The fundamental operations include vector insertion, indexing, and nearest neighbor search implementation using basic Python data structures.

class SimpleVectorStore:
    def __init__(self):
        self.vectors = []
        self.metadata = []
    
    def add_vector(self, vector, metadata=None):
        self.vectors.append(vector)
        self.metadata.append(metadata)
    
    def search(self, query_vector, k=1):
        similarities = [
            cosine_similarity(query_vector, vec)
            for vec in self.vectors
        ]
        top_k = np.argsort(similarities)[-k:][::-1]
        return [(self.metadata[i], similarities[i]) for i in top_k]

# Initialize and populate store
store = SimpleVectorStore()
for idx, vec in enumerate(dense_embeddings):
    store.add_vector(vec, texts[idx])

Implementing Approximate Nearest Neighbors

The k-nearest neighbors search becomes computationally expensive with large datasets. Approximate Nearest Neighbors (ANN) algorithms trade perfect accuracy for significant speed improvements using space partitioning techniques.

import faiss
import numpy as np

class ANNVectorStore:
    def __init__(self, dimension):
        self.dimension = dimension
        self.index = faiss.IndexFlatL2(dimension)
        self.metadata = []
    
    def add_vectors(self, vectors, metadata):
        vectors = np.asarray(vectors).astype('float32')
        self.index.add(vectors)
        self.metadata.extend(metadata)
    
    def search(self, query_vector, k=1):
        query_vector = np.asarray([query_vector]).astype('float32')
        distances, indices = self.index.search(query_vector, k)
        return [(self.metadata[idx], dist) for idx, dist in zip(indices[0], distances[0])]

# Convert embeddings to float32 for FAISS
vectors_float32 = dense_embeddings.astype('float32')

Text Preprocessing for Embeddings

Effective embedding generation requires careful text preprocessing to ensure quality vector representations. This includes tokenization, normalization, and handling of special characters while maintaining semantic meaning.

import re
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
import nltk

# Download required NLTK data
nltk.download('punkt')
nltk.download('stopwords')

def preprocess_text(text):
    # Convert to lowercase
    text = text.lower()
    
    # Remove special characters and digits
    text = re.sub(r'[^a-zA-Z\s]', '', text)
    
    # Tokenize
    tokens = word_tokenize(text)
    
    # Remove stopwords
    stop_words = set(stopwords.words('english'))
    tokens = [t for t in tokens if t not in stop_words]
    
    return ' '.join(tokens)

# Example usage
sample_text = "The quick brown fox jumps over the lazy dog!"
processed_text = preprocess_text(sample_text)
print(f"Original: {sample_text}")
print(f"Processed: {processed_text}")

Document Embedding Pipeline

Document embedding requires a systematic approach to convert raw text into meaningful vector representations. This pipeline handles document loading, preprocessing, and embedding generation using advanced natural language processing techniques.

from transformers import AutoTokenizer, AutoModel
import torch

class DocumentEmbeddingPipeline:
    def __init__(self, model_name='sentence-transformers/all-MiniLM-L6-v2'):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        
    def generate_embedding(self, text):
        # Preprocess text
        cleaned_text = preprocess_text(text)
        
        # Tokenize and encode
        inputs = self.tokenizer(cleaned_text, 
                              padding=True, 
                              truncation=True, 
                              return_tensors="pt")
        
        # Generate embeddings
        with torch.no_grad():
            outputs = self.model(**inputs)
            embeddings = outputs.last_hidden_state.mean(dim=1)
            
        return embeddings.numpy()

# Example usage
pipeline = DocumentEmbeddingPipeline()
doc_embedding = pipeline.generate_embedding("AI and machine learning advances.")
print(f"Embedding shape: {doc_embedding.shape}")

Vector Indexing Strategies

Vector indexing is crucial for efficient similarity search in high-dimensional spaces. Modern indexing techniques employ tree-based structures or clustering methods to partition the vector space and enable logarithmic-time queries.

class VectorIndex:
    def __init__(self, vectors, leaf_size=40):
        self.vectors = np.array(vectors)
        self.kdtree = self._build_kdtree()
        
    def _build_kdtree(self):
        from scipy.spatial import cKDTree
        return cKDTree(self.vectors)
    
    def query(self, vector, k=1):
        distances, indices = self.kdtree.query(vector, k=k)
        return [(idx, dist) for idx, dist in zip(indices, distances)]
    
    def batch_query(self, vectors, k=1):
        return self.kdtree.query(vectors, k=k)

# Example usage
index = VectorIndex(dense_embeddings)
query_vector = dense_embeddings[0]
results = index.query(query_vector, k=2)
print(f"Nearest neighbors: {results}")

Real-world Application - Semantic Search Engine

This implementation demonstrates a complete semantic search engine for document retrieval. It combines preprocessing, embedding generation, and efficient vector search to enable natural language queries over a document collection.

class SemanticSearchEngine:
    def __init__(self):
        self.embedding_pipeline = DocumentEmbeddingPipeline()
        self.vector_store = ANNVectorStore(384)  # MiniLM embedding dimension
        self.documents = []
        
    def add_documents(self, documents):
        self.documents.extend(documents)
        embeddings = [
            self.embedding_pipeline.generate_embedding(doc)[0]
            for doc in documents
        ]
        self.vector_store.add_vectors(embeddings, documents)
    
    def search(self, query, k=3):
        query_embedding = self.embedding_pipeline.generate_embedding(query)[0]
        results = self.vector_store.search(query_embedding, k=k)
        return results

# Example usage
search_engine = SemanticSearchEngine()
documents = [
    "Deep learning models achieve remarkable performance",
    "Natural language processing transforms text analysis",
    "Computer vision enables visual understanding"
]
search_engine.add_documents(documents)
results = search_engine.search("machine learning applications")

Mathematical Foundations of Embeddings

The mathematical principles underlying embeddings involve complex transformations and distance metrics. Understanding these foundations is crucial for optimizing embedding models and similarity computations.

# Mathematical formulas for embedding operations
"""
1. Cosine Similarity:
$$cos(θ) = \frac{A \cdot B}{||A|| ||B||}$$

2. Euclidean Distance:
$$d(p,q) = \sqrt{\sum_{i=1}^n (p_i - q_i)^2}$$

3. Dot Product Attention:
$$Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$$

4. Word2Vec Skip-gram Objective:
$$J(θ) = -\frac{1}{T}\sum_{t=1}^T\sum_{-c≤j≤c,j≠0} log p(w_{t+j}|w_t)$$
"""

# Implementation of distance metrics
def euclidean_distance(v1, v2):
    return np.sqrt(np.sum((v1 - v2) ** 2))

def manhattan_distance(v1, v2):
    return np.sum(np.abs(v1 - v2))

Dimensionality Reduction for Visualization

Dimensionality reduction techniques enable visualization of high-dimensional embeddings in lower-dimensional spaces while preserving relative distances between points. This is crucial for understanding embedding structure and quality assessment.

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

class EmbeddingVisualizer:
    def __init__(self, embeddings, labels=None):
        self.embeddings = embeddings
        self.labels = labels or range(len(embeddings))
        
    def reduce_dimensions(self, n_components=2):
        tsne = TSNE(n_components=n_components, random_state=42)
        return tsne.fit_transform(self.embeddings)
    
    def plot(self):
        reduced_embeddings = self.reduce_dimensions()
        plt.figure(figsize=(10, 8))
        plt.scatter(reduced_embeddings[:, 0], 
                   reduced_embeddings[:, 1])
        
        for i, label in enumerate(self.labels):
            plt.annotate(label, 
                        (reduced_embeddings[i, 0], 
                         reduced_embeddings[i, 1]))
        return plt

# Example usage
visualizer = EmbeddingVisualizer(dense_embeddings, texts)
plot = visualizer.plot()

Embedding Quality Assessment

Quantitative evaluation of embedding quality involves measuring clustering effectiveness, retrieval accuracy, and semantic preservation. These metrics guide model selection and optimization decisions.

from sklearn.metrics import silhouette_score
from sklearn.cluster import KMeans

class EmbeddingEvaluator:
    def __init__(self, embeddings):
        self.embeddings = embeddings
        
    def evaluate_clustering(self, n_clusters=3):
        kmeans = KMeans(n_clusters=n_clusters, random_state=42)
        labels = kmeans.fit_predict(self.embeddings)
        
        silhouette_avg = silhouette_score(self.embeddings, labels)
        inertia = kmeans.inertia_
        
        return {
            'silhouette_score': silhouette_avg,
            'inertia': inertia,
            'cluster_labels': labels
        }
    
    def evaluate_retrieval(self, query_idx, k=5):
        query_vector = self.embeddings[query_idx]
        distances = [
            cosine_similarity(query_vector, vec)
            for vec in self.embeddings
        ]
        top_k = np.argsort(distances)[-k-1:-1][::-1]
        return top_k, [distances[i] for i in top_k]

# Example usage
evaluator = EmbeddingEvaluator(dense_embeddings)
clustering_metrics = evaluator.evaluate_clustering()
print(f"Silhouette Score: {clustering_metrics['silhouette_score']:.4f}")

Production-Ready Vector Store Implementation

A production vector store implementation requires careful consideration of scalability, persistence, and concurrent access. This implementation includes disk-based storage and batch operations for efficiency.

import pickle
from pathlib import Path
import threading

class ProductionVectorStore:
    def __init__(self, dimension, index_path='vector_store'):
        self.dimension = dimension
        self.index_path = Path(index_path)
        self.index_path.mkdir(exist_ok=True)
        
        self.index = faiss.IndexFlatL2(dimension)
        self.metadata = {}
        self.lock = threading.Lock()
        
    def save_state(self):
        with self.lock:
            faiss.write_index(self.index, 
                str(self.index_path / 'vectors.idx'))
            with open(self.index_path / 'metadata.pkl', 'wb') as f:
                pickle.dump(self.metadata, f)
    
    def load_state(self):
        with self.lock:
            if (self.index_path / 'vectors.idx').exists():
                self.index = faiss.read_index(
                    str(self.index_path / 'vectors.idx'))
            if (self.index_path / 'metadata.pkl').exists():
                with open(self.index_path / 'metadata.pkl', 'rb') as f:
                    self.metadata = pickle.load(f)
    
    def batch_add(self, vectors, metadata_list):
        with self.lock:
            start_idx = self.index.ntotal
            vectors = np.asarray(vectors).astype('float32')
            self.index.add(vectors)
            
            for idx, metadata in enumerate(metadata_list):
                self.metadata[start_idx + idx] = metadata
            
        self.save_state()

# Example usage
store = ProductionVectorStore(dimension=dense_embeddings.shape[1])
store.batch_add(dense_embeddings, texts)

Performance Optimization Techniques

Vector search optimization involves careful balancing of accuracy and speed through techniques like quantization, clustering, and caching. These methods significantly reduce memory usage and search latency in large-scale deployments.

import faiss
import numpy as np
from collections import LRUCache

class OptimizedVectorStore:
    def __init__(self, dimension, n_lists=100, n_probes=10):
        # Initialize quantizer
        self.quantizer = faiss.IndexFlatL2(dimension)
        
        # Create IVF index with Product Quantization
        self.index = faiss.IndexIVFPQ(
            self.quantizer,
            dimension,
            n_lists,    # number of Voronoi cells
            8,          # number of sub-vectors
            8           # bits per code (PQ)
        )
        
        self.n_probes = n_probes
        self.cache = LRUCache(maxsize=1000)
        
    def train(self, vectors):
        vectors = np.asarray(vectors).astype('float32')
        self.index.train(vectors)
        
    def add(self, vectors):
        vectors = np.asarray(vectors).astype('float32')
        self.index.add(vectors)
        
    def search(self, query_vector, k=5):
        # Check cache first
        cache_key = hash(query_vector.tobytes())
        if cache_key in self.cache:
            return self.cache[cache_key]
        
        # Configure search parameters
        self.index.nprobe = self.n_probes
        
        # Perform search
        distances, indices = self.index.search(
            query_vector.reshape(1, -1).astype('float32'), 
            k
        )
        
        # Cache results
        results = list(zip(indices[0], distances[0]))
        self.cache[cache_key] = results
        
        return results

# Example usage
optimized_store = OptimizedVectorStore(dimension=dense_embeddings.shape[1])
optimized_store.train(dense_embeddings)
optimized_store.add(dense_embeddings)

Results Analysis and Metrics

A comprehensive evaluation of vector search performance requires analyzing multiple metrics including latency, recall, and precision. This implementation provides tools for measuring and visualizing these metrics.

class SearchEvaluator:
    def __init__(self, ground_truth, search_results):
        self.ground_truth = ground_truth
        self.search_results = search_results
        
    def calculate_precision_at_k(self, k):
        relevant = set(self.ground_truth[:k])
        retrieved = set(idx for idx, _ in self.search_results[:k])
        if not retrieved:
            return 0.0
        return len(relevant & retrieved) / len(retrieved)
    
    def calculate_recall_at_k(self, k):
        relevant = set(self.ground_truth[:k])
        retrieved = set(idx for idx, _ in self.search_results[:k])
        if not relevant:
            return 0.0
        return len(relevant & retrieved) / len(relevant)
    
    def calculate_average_precision(self):
        precisions = []
        relevant_count = 0
        
        for k, (idx, _) in enumerate(self.search_results, 1):
            if idx in self.ground_truth:
                relevant_count += 1
                precision = relevant_count / k
                precisions.append(precision)
                
        if not precisions:
            return 0.0
        return sum(precisions) / len(self.ground_truth)

# Example usage
evaluator = SearchEvaluator(
    ground_truth=[0, 1, 2],
    search_results=[(0, 0.9), (2, 0.8), (1, 0.7)]
)
print(f"Precision@3: {evaluator.calculate_precision_at_k(3):.4f}")
print(f"Recall@3: {evaluator.calculate_recall_at_k(3):.4f}")
print(f"MAP: {evaluator.calculate_average_precision():.4f}")

Additional Resources

More in Embeddings and vectors

  1. Exploring the Four Fundamental Subspaces of Matrices11 min

    The four fundamental subspaces of a matrix are essential concepts in linear algebra. These subspaces provide crucial insights into the structure and behavior...

  2. Embedding Layers Transforming Text into Machine-Readable Vectors8 min

    An embedding layer transforms discrete word tokens into continuous vector spaces, enabling mathematical operations on text data. The fundamental process...

  3. Improving Question-Answer Semantics with HyDE11 min

    The HyDE approach revolutionizes traditional RAG systems by addressing the semantic gap between questions and answers through a novel two-step process that...

  4. Describing Location with Latitude and Longitude Vectors11 min

    A vector space implementation focusing on basic operations like addition, scalar multiplication, and dot products. This forms the foundation for...

  5. Improving Semantic Similarity in Question Answering with HyDE8 min

    Traditional RAG systems face semantic similarity challenges between queries and answers. HyDE (Hypothetical Document Embeddings) addresses this by generating...

  6. Combinatorics Fundamentals6 min

    Combinatorics is a branch of mathematics that deals with counting, arrangement, and combination of objects. It forms the foundation for probability theory...