In-Memory Vector Retrieval for Small RAG Systems

Embeddings & Vector Search intermediate Building with AI ~30 min read

Overview

Many retrieval-augmented generation systems start by adopting a managed vector database before the corpus actually needs one. A dedicated database becomes valuable when embeddings must survive restarts, support high write concurrency, filter by rich metadata, shard across machines, or serve low-latency queries over data that does not fit in memory. Before those requirements appear, the retrieval step can be implemented as a compact in-process matrix search.

The core observation is that dense vector retrieval is linear algebra. If document embeddings are normalized to unit length, cosine similarity reduces to a dot product between a query vector and every row in an embedding matrix. NumPy, BLAS, and scikit-learn can compute those operations efficiently on ordinary CPUs. For internal tools, prototypes, evaluation harnesses, personal knowledge bases, and small-to-medium documentation bots, this can remove network hops, service credentials, ingestion jobs, and operational surface area.

This pattern does not claim that vector databases are unnecessary. It gives engineers a decision boundary. Use an in-memory retriever when the corpus fits comfortably in RAM, rebuild time is acceptable, metadata filtering is simple, and the application can tolerate process-local state. Move to pgvector, Chroma, Milvus, Qdrant, Pinecone, or another database when persistence, concurrent updates, indexing at scale, replication, access control, or operational isolation become real requirements.

A good small-RAG implementation treats the in-memory store as an explicit component, not a toy script hidden inside the prompt path. It owns document chunks, normalized embeddings, metadata, top-k scoring, persistence snapshots if needed, and measurement hooks. That makes it easy to replace later because the rest of the RAG pipeline talks to a retrieval interface rather than to NumPy arrays directly.

Core Concepts

Embeddings as a Dense Matrix

After chunking a corpus, each chunk is embedded into a fixed-width vector. Store those vectors in a two-dimensional matrix with shape num_chunks x dimensions. A query embedding is a one-row vector with the same dimensionality. Retrieval computes a score for every chunk, ranks the scores, and returns the top matches plus their source text and metadata.

For 384-dimensional float32 embeddings, one million vectors require roughly 1.5 GB for the raw matrix. Larger models increase that footprint, but the math remains simple. The memory budget is often more forgiving than teams expect, especially for prototypes and internal corpora measured in thousands or low millions of chunks.

Cosine Similarity and Normalization

Cosine similarity compares vector direction rather than magnitude. If both document and query embeddings are normalized to length one, the cosine formula collapses to a dot product. That matters because matrix multiplication is one of the most optimized operations in the Python numerical ecosystem.

Normalization should happen at ingestion for document embeddings and at query time for the query embedding. Normalizing once prevents repeated denominator work during search and makes scoring behavior explicit. It also keeps the retriever compatible with backends that expect normalized dense vectors.

Top-k Selection Without a Database

A brute-force search computes every score, then selects the highest k values. Sorting every score works but is wasteful for large matrices because the application only needs the top few candidates. A production in-memory retriever should use partial selection, such as argpartition, then sort only the candidate slice.

Approximate or indexed retrieval can still exist without a standalone vector database. scikit-learn exposes nearest-neighbor implementations, and FAISS can provide local ANN indexes. The same decision rule applies: add indexing when measured latency requires it, not because the architecture diagram expects a vector database box.

Key Concept: Start With the Retrieval Contract

The important abstraction is not the storage product. It is the contract: ingest chunks, store embeddings and metadata, retrieve top-k candidates for a query, and expose enough metrics to know when the implementation no longer fits.

How It Works

Build a Minimal Store

An in-memory vector store keeps two synchronized collections: an embedding matrix for scoring and a list of records for returning text and metadata. The source article demonstrates this shape with Sentence Transformers, NumPy, and scikit-learn. The portable version below shows the same mechanism while keeping the embedding function injectable.

import numpy as np

class InMemoryVectorStore:
    def __init__(self, embed_fn):
        self.embed_fn = embed_fn
        self.records = []
        self.embeddings = None

    def _normalize(self, matrix):
        norms = np.linalg.norm(matrix, axis=1, keepdims=True)
        return matrix / np.maximum(norms, 1e-12)

    def add_documents(self, docs):
        texts = [doc["text"] for doc in docs]
        vectors = np.asarray(self.embed_fn(texts), dtype=np.float32)
        vectors = self._normalize(vectors)

        self.embeddings = vectors if self.embeddings is None else np.vstack([
            self.embeddings,
            vectors,
        ])
        self.records.extend(docs)

The store does not need a network service, schema migration, or background index builder. It only needs a stable embedding model and a consistent document record format. For a RAG application, each record should include the chunk text, source identifier, chunk index, and any simple metadata needed for filtering or attribution.

Search by Matrix Multiplication

At query time, embed and normalize the query, multiply the document matrix by the query vector, and select the top matches. Because each row is already normalized, the dot product produces cosine similarity scores directly.

    def search(self, query, k=5, filter_fn=None):
        if self.embeddings is None or not self.records:
            return []

        query_vec = np.asarray(self.embed_fn([query]), dtype=np.float32)
        query_vec = self._normalize(query_vec)
        scores = self.embeddings @ query_vec[0]

        if filter_fn is not None:
            mask = np.array([filter_fn(record) for record in self.records])
            scores = np.where(mask, scores, -np.inf)

        k = min(k, len(self.records))
        candidate_idx = np.argpartition(scores, -k)[-k:]
        ranked_idx = candidate_idx[np.argsort(scores[candidate_idx])[::-1]]

        return [
            {
                "score": float(scores[i]),
                "text": self.records[i]["text"],
                "metadata": self.records[i].get("metadata", {}),
            }
            for i in ranked_idx
            if np.isfinite(scores[i])
        ]

This is the retrieval core. The rest of the RAG pipeline can rerank candidates, pack context into a prompt, generate an answer, and display citations. If this function is fast enough and operationally simple enough, a database would add complexity before adding value.

Chunk, Embed, Retrieve, Generate

The retriever sits between document processing and generation. Documents are parsed and chunked, chunks are embedded and stored, and user queries are embedded against the same vector space. Retrieved chunks become grounding material for the LLM.

def chunk_text(text, chunk_size=1000, overlap=200):
    step = chunk_size - overlap
    for start in range(0, len(text), step):
        chunk = text[start:start + chunk_size]
        if len(chunk.strip()) >= 50:
            yield chunk

docs = []
for path in corpus_paths:
    text = path.read_text(encoding="utf-8")
    for i, chunk in enumerate(chunk_text(text)):
        docs.append({
            "text": chunk,
            "metadata": {"source": path.name, "chunk": i},
        })

store.add_documents(docs)
matches = store.search("How does the billing retry policy work?", k=6)

Chunking quality still matters. An in-memory store does not fix bad parsing, missing overlap, malformed metadata, or an embedding model that does not match the domain. It only simplifies the retrieval infrastructure while those upstream decisions are being tested.

Practical Applications

RAG Prototypes and Evaluation Harnesses

Early RAG projects need fast iteration on chunking, prompts, reranking, citation rendering, and evaluation datasets. An in-memory store lets a team rebuild the corpus repeatedly without provisioning a database, writing migration scripts, or debugging ingestion workers. This is especially useful when testing whether semantic retrieval is even the right retrieval strategy for the task.

Evaluation harnesses also benefit from process-local retrieval. A test runner can load a frozen embedding snapshot, run a suite of queries, score retrieval quality, and compare changes deterministically. That reduces the number of moving parts in CI and makes regressions easier to attribute.

Small Internal Knowledge Bases

Team handbooks, API docs, runbooks, design notes, and support macros often fit into memory after chunking. If updates are infrequent and the application can rebuild on deploy, a simple retriever may be enough. The system can snapshot embeddings to disk with np.save or a compressed archive and reload them at process start.

The tradeoff is operational: a process restart may need to reload or rebuild the matrix, and multiple application instances may each carry their own copy. That is acceptable for small corpora but becomes wasteful as the matrix grows or updates become frequent.

Edge, Local, and Offline Retrieval

Local developer tools, desktop applications, and edge deployments may not have reliable access to a managed vector service. In-memory retrieval works well when the corpus is bundled with the application, generated locally, or scoped to one user. It also keeps sensitive documents out of additional infrastructure.

Offline retrieval is not automatically secure or private, but it reduces exposure. The application still needs local encryption, access controls, and model choices appropriate for the data. The retriever simply avoids another remote system in the data path.

Implementation Patterns

Backend-Neutral Retriever Interface

Do not let application code depend directly on NumPy arrays. Define a retriever interface that can later be backed by an in-memory store, pgvector, Chroma, or a managed vector database. The RAG chain should not know which backend is serving the top-k candidates.

from typing import Protocol

class Retriever(Protocol):
    def add_documents(self, docs: list[dict]) -> None:
        ...

    def search(self, query: str, k: int = 5) -> list[dict]:
        ...

def answer_question(retriever: Retriever, llm, question: str):
    chunks = retriever.search(question, k=6)
    context = "\n\n".join(chunk["text"] for chunk in chunks)
    return llm.generate(
        "Answer using only this context:\n"
        f"{context}\n\nQuestion: {question}"
    )

This interface preserves the migration path. When the corpus outgrows RAM or needs concurrent updates, replace the backend without rewriting prompt assembly, answer generation, or evaluation code.

Snapshot Persistence

A pure in-memory matrix disappears when the process exits. For small systems, snapshot persistence is often enough: store the embedding matrix and document records together whenever ingestion completes, then load them at startup. Use version metadata so stale snapshots are rejected when the embedding model or chunking policy changes.

import json

def save_snapshot(store, path):
    np.save(path / "embeddings.npy", store.embeddings)
    (path / "records.json").write_text(json.dumps(store.records))
    (path / "manifest.json").write_text(json.dumps({
        "embedding_model": "all-MiniLM-L6-v2",
        "dimensions": int(store.embeddings.shape[1]),
        "chunk_policy": "chars-1000-overlap-200",
    }))

def load_snapshot(store, path):
    store.embeddings = np.load(path / "embeddings.npy")
    store.records = json.loads((path / "records.json").read_text())

Snapshots are not a replacement for database durability. They are a pragmatic bridge for applications whose corpus changes on deploy, on a schedule, or by explicit rebuild. Once the system needs row-level updates and transactional guarantees, move persistence into a database.

Measured Migration Thresholds

Set thresholds before the in-memory retriever becomes a liability. Track matrix size, startup load time, rebuild time, search latency, memory headroom, update frequency, and filter complexity. Those metrics turn “we should probably use a vector database” into a concrete engineering decision.

A useful rule is to migrate when one requirement is both real and painful: the matrix no longer fits comfortably in memory, rebuilds are too slow, filters need indexed metadata, updates and deletes are frequent, multiple services need shared retrieval, or operators need managed replication and backups.

Performance Considerations

Common Pitfalls

Further Reading