Optimizing Vector Embeddings for Better Search Results

Vector embeddings have revolutionized information retrieval by enabling semantic search capabilities that understand the meaning behind queries rather than just matching keywords. However, the effectiveness of embedding-based search systems depends heavily on how these embeddings are generated, processed, and indexed. In this technical guide, we'll explore advanced techniques for optimizing vector embeddings to achieve better search relevance, reduced latency, and improved overall system performance.

Understanding the Vector Embedding Pipeline

Before diving into optimization techniques, let's review the key components of a vector embedding pipeline:

1. Document Processing: Preparing source documents through cleaning, normalization, and chunking
2. Embedding Generation: Converting text chunks into vector representations using embedding models
3. Vector Indexing: Building efficient data structures for storing and retrieving vectors
4. Query Processing: Transforming user queries into vectors and retrieving relevant results
5. Ranking and Filtering: Post-processing to improve result quality

Optimizations can be applied at each stage of this pipeline, with improvements in earlier stages often cascading throughout the system.

Document Processing Optimization Techniques

1. Advanced Chunking Strategies

How you divide documents into chunks significantly impacts retrieval quality. Here are advanced chunking techniques that go beyond simple fixed-length segmentation:

from sentence_transformers import SentenceTransformer
from sklearn.cluster import AgglomerativeClustering
import numpy as np

def semantic_chunking(document, min_chunk_size=100, max_chunk_size=1000):
    # Step 1: Split into initial small segments (sentences or paragraphs)
    sentences = document.split('. ')

    # Step 2: Generate embeddings for each sentence
    model = SentenceTransformer('all-MiniLM-L6-v2')
    embeddings = model.encode(sentences)

    # Step 3: Cluster similar sentences
    clustering_model = AgglomerativeClustering(
        n_clusters=None,
        distance_threshold=0.25,  # Adjust based on desired granularity
        affinity='cosine',
        linkage='average'
    )
    clusters = clustering_model.fit_predict(embeddings)

    # Step 4: Form chunks based on clusters
    chunks = []
    current_chunk = []
    current_size = 0
    current_cluster = clusters[0]

    for i, sentence in enumerate(sentences):
        if (clusters[i] != current_cluster or
            current_size + len(sentence) > max_chunk_size) and current_size >= min_chunk_size:
            # Start a new chunk
            chunks.append('. '.join(current_chunk) + '.')
            current_chunk = [sentence]
            current_size = len(sentence)
            current_cluster = clusters[i]
        else:
            # Continue the current chunk
            current_chunk.append(sentence)
            current_size += len(sentence)

    # Add the last chunk
    if current_chunk:
        chunks.append('. '.join(current_chunk) + '.')

    return chunks

def sliding_window_chunking(document, chunk_size=500, overlap=100):
    """
    Split document into overlapping chunks using a sliding window approach.

    Args:
        document (str): The document text to chunk
        chunk_size (int): Target size of each chunk
        overlap (int): Number of characters to overlap between chunks

    Returns:
        list: A list of text chunks
    """
    if len(document) <= chunk_size:
        return [document]

    chunks = []
    start = 0

    while start < len(document):
        # Find the end position for this chunk
        end = start + chunk_size

        # Don't cut words - find the nearest space after the end position
        if end < len(document):
            # Look for the next paragraph break first
            next_para = document.find('\n\n', end - 50, end + 50)
            if next_para != -1 and next_para - end < 100:
                end = next_para
            else:
                # Fall back to finding the next space
                next_space = document.find(' ', end)
                if next_space != -1:
                    end = next_space
        else:
            end = len(document)

        # Extract the chunk and add to list
        chunks.append(document[start:end])

        # Move the start position, accounting for overlap
        start = end - overlap

        # Make sure we don't get stuck in a loop with small documents
        if start >= len(document) - overlap:
            break

    return chunks

def hierarchical_chunking(document):
    """
    Create a multi-level hierarchy of chunks.

    Returns:
        dict: A hierarchical structure of chunks
    """
    # Level 1: Document level
    doc_embedding = {
        "text": document,
        "level": "document",
        "children": []
    }

    # Level 2: Section level
    sections = split_into_sections(document)
    for i, section in enumerate(sections):
        section_chunk = {
            "text": section,
            "level": "section",
            "parent_idx": 0,  # Points to the document
            "children": []
        }
        doc_embedding["children"].append(section_chunk)

        # Level 3: Paragraph level
        paragraphs = section.split("\n\n")
        for j, para in enumerate(paragraphs):
            if len(para.strip()) > 50:  # Exclude very small paragraphs
                para_chunk = {
                    "text": para,
                    "level": "paragraph",
                    "parent_idx": i
                }
                section_chunk["children"].append(para_chunk)

    return doc_embedding

def split_into_sections(document):
    """Split document into sections based on headings."""
    import re

    # This pattern matches common heading patterns like '# Heading' or 'Section 1:'
    heading_pattern = r'(?:\n|^)(?:#{1,6}\s+[^\n]+|\d+\.\s+[^\n]+|[A-Z][A-Za-z\s]+:)'

    # Find all potential section boundaries
    matches = list(re.finditer(heading_pattern, document))

    sections = []
    for i, match in enumerate(matches):
        start = match.start()
        # If this is the last match, the section goes to the end of the document
        end = matches[i+1].start() if i < len(matches)-1 else len(document)

        # Extract the section text
        section = document[start:end]
        sections.append(section)

    # Handle the case where there are no clear section headings
    if not sections:
        sections = [document]

    return sections

2. Content-Aware Preprocessing

Applying domain-specific preprocessing can significantly improve embedding quality:

• Entity Normalization: Standardize entity mentions (e.g., "IBM" and "International Business Machines") to improve consistency.
• Domain-Specific Tokenization: Use specialized tokenizers for technical, legal, or medical content to better handle domain-specific terms.
• Structural Element Preservation: Retain important structural indicators like headings, lists, and tables with special tokens.

def preprocess_technical_document(text):
    """Specialized preprocessing for technical documentation."""
    import re

    # 1. Preserve code blocks with special tokens
    text = re.sub(r'```(?:\w+)?\n(.*?)\n```', r' [CODE] \1 [/CODE] ', text, flags=re.DOTALL)

    # 2. Highlight headings with special tokens
    text = re.sub(r'(#{1,6})\s+(.*?)(?:\n|$)', r' [HEADING] \2 [/HEADING] ', text)

    # 3. Normalize technical terms and acronyms
    tech_terms = {
        "javascript": "JavaScript",
        "js": "JavaScript",
        "py": "Python",
        "ML": "machine learning",
        "DL": "deep learning",
        "NLP": "natural language processing",
        # Add more domain-specific normalizations
    }

    for term, replacement in tech_terms.items():
        text = re.sub(r'\b' + re.escape(term) + r'\b', replacement, text, flags=re.IGNORECASE)

    # 4. Handle API references and function names
    # Identify and preserve function calls with special tokens
    text = re.sub(r'\b(\w+)\((.*?)\)', r' [FUNCTION] \1(\2) [/FUNCTION] ', text)

    return text

Embedding Generation Optimization

1. Model Selection and Fine-Tuning

The choice of embedding model dramatically impacts search quality. Here are approaches to optimize embedding generation:

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

def fine_tune_embedding_model(train_examples, base_model='all-MiniLM-L6-v2', epochs=10):
    """
    Fine-tune a sentence transformer model on domain-specific examples.

    Args:
        train_examples: List of tuples (sentence1, sentence2, similarity_score)
        base_model: Base model to fine-tune
        epochs: Number of training epochs

    Returns:
        Fine-tuned model
    """
    # Convert training examples to the format expected by sentence-transformers
    examples = [
        InputExample(texts=[s1, s2], label=score)
        for s1, s2, score in train_examples
    ]

    # Create dataloader
    train_dataloader = DataLoader(examples, shuffle=True, batch_size=16)

    # Load base model
    model = SentenceTransformer(base_model)

    # Define loss function - CosineSimilarityLoss for similarity scores
    train_loss = losses.CosineSimilarityLoss(model)

    # Train the model
    model.fit(
        train_objectives=[(train_dataloader, train_loss)],
        epochs=epochs,
        warmup_steps=100,
        output_path="fine-tuned-embeddings-model"
    )

    return model

# Example usage
# Generate pairs of similar texts from your domain
training_pairs = [
    ("How do I configure the API authentication?",
     "What's the process for setting up API auth credentials?",
     0.9),  # High similarity

    ("What programming languages are supported?",
     "Do you support Python integration?",
     0.7),  # Medium similarity

    ("How much does the enterprise plan cost?",
     "Can I deploy the model on my own hardware?",
     0.1),  # Low similarity

    # Add more domain-specific examples
]

# Fine-tune the model
domain_model = fine_tune_embedding_model(training_pairs)

class EnsembleEmbedder:
    """
    Combines multiple embedding models into an ensemble for improved performance.
    """
    def __init__(self, models, weights=None):
        """
        Initialize ensemble with multiple models and optional weights.

        Args:
            models: List of SentenceTransformer models
            weights: Optional list of weights for each model (defaults to equal weights)
        """
        self.models = models

        if weights is None:
            # Equal weighting by default
            self.weights = [1.0 / len(models)] * len(models)
        else:
            # Normalize weights to sum to 1
            total = sum(weights)
            self.weights = [w / total for w in weights]

    def encode(self, texts, normalize=True):
        """
        Encode texts using the ensemble of models.

        Args:
            texts: List of texts to encode
            normalize: Whether to L2-normalize individual embeddings

        Returns:
            Combined embeddings
        """
        import numpy as np

        # Get embeddings from each model
        all_embeddings = []
        for i, model in enumerate(self.models):
            emb = model.encode(texts, normalize_embeddings=normalize)
            all_embeddings.append(emb * self.weights[i])

        # Combine embeddings
        combined = np.zeros_like(all_embeddings[0])
        for emb in all_embeddings:
            combined += emb

        # Re-normalize if needed
        if normalize:
            norms = np.linalg.norm(combined, axis=1, keepdims=True)
            combined = combined / norms

        return combined

# Example usage
from sentence_transformers import SentenceTransformer

# Load different embedding models
general_model = SentenceTransformer('all-MiniLM-L6-v2')
mpnet_model = SentenceTransformer('paraphrase-mpnet-base-v2')
domain_model = SentenceTransformer('fine-tuned-embeddings-model')

# Create ensemble with custom weights
ensemble = EnsembleEmbedder(
    models=[general_model, mpnet_model, domain_model],
    weights=[0.3, 0.3, 0.4]  # More weight to domain-specific model
)

# Generate embeddings
query = "How do I integrate the API with my application?"
embedding = ensemble.encode([query])[0]

2. Dimensionality and Efficiency Techniques

def optimize_embeddings_with_pca(embeddings, target_dimensions=256):
    """
    Reduce embedding dimensions using PCA.

    Args:
        embeddings: Original high-dimensional embeddings
        target_dimensions: Target number of dimensions

    Returns:
        Reduced-dimension embeddings
    """
    from sklearn.decomposition import PCA
    import numpy as np

    # Fit PCA on embeddings
    pca = PCA(n_components=target_dimensions)
    pca.fit(embeddings)

    # Transform embeddings to lower dimension
    reduced_embeddings = pca.transform(embeddings)

    # Calculate how much variance is retained
    explained_variance = sum(pca.explained_variance_ratio_)
    print(f"Retained {explained_variance:.2%} of original variance with {target_dimensions} dimensions")

    return reduced_embeddings, pca

def quantize_embeddings(embeddings, bits=8):
    """
    Quantize embeddings to lower precision.

    Args:
        embeddings: Original embeddings
        bits: Target bits per value (8 or 16)

    Returns:
        Quantized embeddings and scale factors for reconstruction
    """
    import numpy as np

    # Find min and max for scaling
    mins = embeddings.min(axis=0)
    maxs = embeddings.max(axis=0)

    # Calculate scale to use full range of target precision
    scales = (maxs - mins) / (2**bits - 1)

    # Scale and convert to integers
    if bits == 8:
        dtype = np.uint8
    elif bits == 16:
        dtype = np.uint16
    else:
        raise ValueError("Only 8 or 16 bits supported")

    quantized = np.round((embeddings - mins) / scales).astype(dtype)

    return quantized, mins, scales

def dequantize_embeddings(quantized, mins, scales):
    """
    Restore quantized embeddings to floating point.

    Args:
        quantized: Quantized embeddings
        mins: Minimum values per dimension
        scales: Scale factors per dimension

    Returns:
        Approximation of original embeddings
    """
    return (quantized.astype(float) * scales) + mins

Vector Indexing and Retrieval Optimization

1. Index Structures and Algorithms

The choice of vector index dramatically impacts both search speed and accuracy:

class HybridVectorIndex:
    """
    Hybrid vector index combining exact search for high-priority documents
    and approximate search for the long tail.
    """
    def __init__(self, dimension, ann_algorithm='hnsw'):
        import faiss
        import numpy as np

        self.dimension = dimension
        self.ann_algorithm = ann_algorithm

        # Exact index for high-priority vectors
        self.exact_index = faiss.IndexFlatL2(dimension)

        # Approximate index for the rest
        if ann_algorithm == 'hnsw':
            # HNSW index for fast approximate search
            self.approx_index = faiss.IndexHNSWFlat(dimension, 32)  # 32 neighbors per layer
            self.approx_index.hnsw.efConstruction = 100  # Higher values = better quality but slower build
        elif ann_algorithm == 'ivf':
            # IVF index for memory-efficient search
            nlist = 100  # Number of Voronoi cells
            quantizer = faiss.IndexFlatL2(dimension)
            self.approx_index = faiss.IndexIVFFlat(quantizer, dimension, nlist)
            self.approx_index.nprobe = 10  # Number of cells to visit during search

        # Track which IDs are in which index
        self.exact_ids = []
        self.approx_ids = []

    def add_to_exact(self, vectors, ids):
        """Add high-priority vectors to exact index."""
        import numpy as np

        vectors = np.array(vectors).astype('float32')
        if ids is None:
            ids = np.arange(len(vectors)) + len(self.exact_ids)

        faiss.normalize_L2(vectors)  # Normalize for cosine similarity
        self.exact_index.add(vectors)
        self.exact_ids.extend(ids)

    def add_to_approx(self, vectors, ids=None):
        """Add regular vectors to approximate index."""
        import numpy as np

        vectors = np.array(vectors).astype('float32')
        if ids is None:
            ids = np.arange(len(vectors)) + len(self.approx_ids)

        faiss.normalize_L2(vectors)  # Normalize for cosine similarity

        # Train index if needed (for IVF)
        if self.ann_algorithm == 'ivf' and not self.approx_index.is_trained:
            self.approx_index.train(vectors)

        self.approx_index.add(vectors)
        self.approx_ids.extend(ids)

    def search(self, query_vector, top_k=10, exact_weight=0.7):
        """
        Search both indexes and combine results.

        Args:
            query_vector: Query vector
            top_k: Number of results to return
            exact_weight: Weight to give exact results vs approximate

        Returns:
            Combined search results
        """
        import numpy as np

        query_vector = np.array([query_vector]).astype('float32')
        faiss.normalize_L2(query_vector)

        # Number of results to get from each index
        exact_k = min(top_k, len(self.exact_ids))
        approx_k = min(top_k * 2, len(self.approx_ids))  # Get extra results from approx

        # Search exact index
        if exact_k > 0:
            exact_distances, exact_indices = self.exact_index.search(query_vector, exact_k)
            exact_results = [(self.exact_ids[idx], score * exact_weight)
                            for idx, score in zip(exact_indices[0], exact_distances[0])]
        else:
            exact_results = []

        # Search approximate index
        if approx_k > 0:
            approx_distances, approx_indices = self.approx_index.search(query_vector, approx_k)
            approx_results = [(self.approx_ids[idx], score * (1.0 - exact_weight))
                             for idx, score in zip(approx_indices[0], approx_distances[0])]
        else:
            approx_results = []

        # Combine and sort results
        all_results = exact_results + approx_results
        all_results.sort(key=lambda x: x[1])

        return all_results[:top_k]

class MetadataEnhancedVectorSearch:
    """
    Vector search with metadata filtering capabilities.
    """
    def __init__(self, dimension):
        import faiss

        self.dimension = dimension
        self.index = faiss.IndexFlatL2(dimension)
        self.metadata = []  # List to store metadata for each vector

    def add_vectors(self, vectors, metadata_list):
        """
        Add vectors with associated metadata.

        Args:
            vectors: Vectors to add
            metadata_list: List of metadata dictionaries for each vector
        """
        import numpy as np

        assert len(vectors) == len(metadata_list), "Length mismatch between vectors and metadata"

        vectors = np.array(vectors).astype('float32')
        faiss.normalize_L2(vectors)

        self.index.add(vectors)
        self.metadata.extend(metadata_list)

    def search(self, query_vector, top_k=100, filters=None):
        """
        Search vectors with optional metadata filtering.

        Args:
            query_vector: Query vector
            top_k: Number of initial candidates to retrieve
            filters: Dictionary of metadata filters ({field: value} or {field: [value1, value2]})

        Returns:
            Filtered search results with distances and metadata
        """
        import numpy as np

        query_vector = np.array([query_vector]).astype('float32')
        faiss.normalize_L2(query_vector)

        # Get initial candidates - retrieve extra to allow for filtering
        search_k = min(top_k * 10, self.index.ntotal) if filters else top_k
        distances, indices = self.index.search(query_vector, search_k)

        results = []
        for i, idx in enumerate(indices[0]):
            # Skip invalid indices that can occur with empty indices
            if idx < 0 or idx >= len(self.metadata):
                continue

            meta = self.metadata[idx]
            distance = distances[0][i]

            # Apply filters if specified
            if filters and not self._matches_filters(meta, filters):
                continue

            results.append({
                "id": idx,
                "distance": float(distance),
                "metadata": meta
            })

            # Stop once we have enough results after filtering
            if len(results) >= top_k:
                break

        return results

    def _matches_filters(self, metadata, filters):
        """Check if metadata matches all filters."""
        for field, value in filters.items():
            if field not in metadata:
                return False

            if isinstance(value, list):
                # Check if metadata value is in the list of acceptable values
                if metadata[field] not in value:
                    return False
            else:
                # Direct comparison
                if metadata[field] != value:
                    return False

        return True

# Example usage
import numpy as np

# Initialize search system
vector_search = MetadataEnhancedVectorSearch(dimension=384)

# Add vectors with metadata
vectors = [
    [0.1, 0.2, ..., 0.3],  # Vector representation 1
    [0.5, 0.1, ..., 0.9],  # Vector representation 2
]

metadata = [
    {"doctype": "article", "domain": "finance", "date": "2025-01-15"},
    {"doctype": "faq", "domain": "technical", "date": "2025-03-20"},
]

vector_search.add_vectors(vectors, metadata)

# Search with metadata filters
query = [0.2, 0.3, ..., 0.1]  # Query vector
results = vector_search.search(
    query_vector=query,
    top_k=5,
    filters={"domain": "finance", "doctype": "article"}
)

2. Query Optimization Techniques

Optimizing how queries are processed can substantially improve search relevance:

def generate_query_variations(query, num_variations=3):
    """
    Generate semantically similar variations of the query.

    This can help capture relevant documents that use different terminology.

    Args:
        query: Original query text
        num_variations: Number of variations to generate

    Returns:
        List of query variations
    """
    import tensorflow as tf
    import tensorflow_hub as hub

    # Load universal sentence encoder (or other appropriate model)
    model = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")

    # Define potential prefixes to create variations
    prefixes = [
        "In other words, ",
        "Similarly, ",
        "Another way to ask this is, ",
        "To rephrase, ",
        "Alternatively, "
    ]

    # Combine the query with prefixes to create a batch of potential paraphrases
    inputs = prefixes[:num_variations]

    # Use the language model to complete these prompts
    # This is a simplified example - in practice, you might use a more sophisticated
    # approach with a generative model like GPT or T5

    # For now, let's assume we have a function that generates paraphrases
    variations = generate_paraphrases(query, num_variations)

    # Include the original query
    all_queries = [query] + variations

    return all_queries

def query_with_expansion(search_system, query_text, top_k=10):
    """
    Perform search with query expansion.

    Args:
        search_system: Vector search system
        query_text: Original query text
        top_k: Number of results to return

    Returns:
        Merged search results
    """
    # Generate query variations
    queries = generate_query_variations(query_text)

    # Generate embeddings for all queries
    query_embeddings = embed_texts(queries)

    # Search with each query variation
    all_results = []
    for i, embedding in enumerate(query_embeddings):
        results = search_system.search(embedding, top_k=top_k)

        # Apply a weight based on similarity to original query
        query_weight = 1.0 if i == 0 else 0.8  # Original query gets full weight

        for result in results:
            result["score"] *= query_weight
            all_results.append(result)

    # Merge results by document ID, taking the highest score
    merged = {}
    for result in all_results:
        doc_id = result["id"]
        if doc_id not in merged or result["score"] > merged[doc_id]["score"]:
            merged[doc_id] = result

    # Sort by final score and return top results
    final_results = sorted(merged.values(), key=lambda x: x["score"], reverse=True)
    return final_results[:top_k]

class HybridSearchEngine:
    """
    Combines vector search with keyword search for better results.
    """
    def __init__(self, vector_search, keyword_search):
        self.vector_search = vector_search
        self.keyword_search = keyword_search

    def search(self, query, top_k=10, vector_weight=0.7):
        """
        Perform hybrid search combining vector and keyword approaches.

        Args:
            query: Search query
            top_k: Number of results to return
            vector_weight: Weight to give vector results (0-1)

        Returns:
            Combined search results
        """
        # Get more results than needed from each system to ensure good coverage
        vector_k = min(top_k * 2, 100)
        keyword_k = min(top_k * 2, 100)

        # Get vector search results
        vector_results = self.vector_search.search(query, top_k=vector_k)

        # Get keyword search results
        keyword_results = self.keyword_search.search(query, top_k=keyword_k)

        # Normalize scores - convert to 0-1 range
        self._normalize_scores(vector_results)
        self._normalize_scores(keyword_results)

        # Create lookup dictionaries
        vector_dict = {result["id"]: result for result in vector_results}
        keyword_dict = {result["id"]: result for result in keyword_results}

        # Find all unique document IDs
        all_ids = set(vector_dict.keys()) | set(keyword_dict.keys())

        # Combine scores
        combined_results = []
        for doc_id in all_ids:
            vector_score = vector_dict.get(doc_id, {"score": 0})["score"]
            keyword_score = keyword_dict.get(doc_id, {"score": 0})["score"]

            # Weighted combination
            combined_score = (vector_score * vector_weight) + (keyword_score * (1 - vector_weight))

            # Get the metadata from whichever result has it
            metadata = (vector_dict.get(doc_id) or keyword_dict.get(doc_id))["metadata"]

            combined_results.append({
                "id": doc_id,
                "score": combined_score,
                "vector_score": vector_score,
                "keyword_score": keyword_score,
                "metadata": metadata
            })

        # Sort by combined score
        combined_results.sort(key=lambda x: x["score"], reverse=True)

        return combined_results[:top_k]

    def _normalize_scores(self, results):
        """Normalize scores to 0-1 range."""
        if not results:
            return

        # Find max and min scores
        scores = [r["score"] for r in results]
        max_score = max(scores)
        min_score = min(scores)

        # Avoid division by zero
        score_range = max_score - min_score
        if score_range == 0:
            # All scores are the same
            for result in results:
                result["score"] = 1.0
            return

        # Normalize to 0-1
        for result in results:
            result["score"] = (result["score"] - min_score) / score_range

Benchmark Results and Trade-offs

We've benchmarked these optimization techniques across different datasets and use cases. Here are key findings:

Conclusion

Optimizing vector embeddings is a multi-faceted challenge that involves trade-offs between relevance, performance, and complexity. For most systems, a layered approach works best:

1. Start with intelligent document processing and chunking
2. Select and potentially fine-tune appropriate embedding models
3. Implement efficient indexing with metadata filtering capabilities
4. Add query optimization techniques where needed
5. Apply dimensionality reduction and quantization selectively based on scale requirements

By carefully applying these optimization techniques, you can significantly improve the relevance and performance of embedding-based search systems. The key is to focus optimization efforts on the most impactful areas for your specific use case and dataset characteristics.