Back to home

AI Internals - Embeddings & Retrieval

12 min read
Cover Image for AI Internals - Embeddings & Retrieval
Lucas LemosLucas Lemos

Introduction

In AI Internals - Prompts & Context we treated retrieval as a slot in the assembly pipeline: inject small, relevant snippets instead of whole documents.

This part is about how those snippets get chosen. An embedding model maps text into a fixed-size vector. A retrieval stack stores those vectors, compares a query vector to them, and returns the nearest neighbors. That search problem is separate from generation — and most "RAG is broken" bugs are actually retrieval bugs that only show up after the LLM invents around missing evidence.

We stay on search here: vectors, similarity, indexes, hybrid retrieval, and reranking. Chunking strategies and grounding the answer in retrieved text land in RAG.

Two different "embeddings"

In How LLMs Work, "embedding" meant the model's input table: token ID → vector inside the transformer residual stream. That table is part of next-token prediction.

Here we mean a different artifact: a sentence or passage embedding produced by a model trained (or adapted) so that texts with related meaning end up close in vector space. You call an embeddings API or run a local encoder, get back something like 384 / 768 / 1536 floats, and never sample a completion from it.

Confusing the two leads to odd product choices — for example expecting gpt-4.1-mini chat completions to give you a reusable document index, or treating cosine distance as somehow "what the LLM thinks."

What the vector is encoding

An embedding model is still a neural net, but the training objective pushes related texts together and unrelated texts apart (contrastive losses, hard negatives, instruction-tuned retrieval pairs — the details vary by model family).

Operational consequences:

  • Same model for index and query. Mixing text-embedding-3-small vectors with an open nomic index is noise with a pretty UI. Rebuild the index when you change the embedding model.
  • Dimension is a contract. A 768-d index cannot store 1536-d vectors. Providers sometimes expose a dimensions parameter that truncates or projects — only use it if the docs say the shortened vectors stay comparable.
  • Normalization matters for the metric. Many APIs return L2-normalized vectors so cosine similarity collapses to a dot product. If you normalize twice, or skip normalization while assuming cosine, rankings shift silently.

Closeness is geometric, not magical:

query:  "refund policy for annual plans"
near:   "How do I cancel and get money back on a yearly subscription?"
far:    "How do I rotate API keys?"

Lexical overlap helps some models and metrics, but the point of dense retrieval is catching paraphrase. When it fails, it often fails on short queries, domain jargon the embedder never saw, or passages that answer the question without sharing surface words in a way the model was trained for.

Similarity: cosine, dot product, L2

Given two vectors a and b:

MetricIdeaTypical use
CosineAngle between vectors (direction)Default for many text embedding APIs
Dot productSame as cosine if both are unit-lengthFast path when vectors are pre-normalized
L2 / EuclideanStraight-line distanceSome indexes and vision embeddings

Pick one metric and keep it consistent from training assumptions through the index configuration. Most managed vector databases ask you to declare the distance function at collection create time; changing it later means reindexing.

Scores are only comparable inside the same query. A cosine of 0.72 on one corpus does not mean the same thing as 0.72 on another model or another collection. Thresholds need calibration on your data, not blog defaults.

The retrieval pipeline

Two phases, two clocks:

Index time (offline or async):

  1. Split source docs into retrieval units (passages). Chunking policy belongs to RAG; for this article treat each unit as a string with an ID and metadata.
  2. Embed each unit with the chosen model.
  3. Upsert (id, vector, metadata, raw text pointer) into a store.

Query time (per request):

  1. Embed the query (sometimes with a query-specific instruction prefix the model expects).
  2. Search top-k nearest neighbors, optionally filtered by metadata.
  3. Optionally rerank or fuse with lexical hits.
  4. Hand the surviving snippets to the prompt assembler from the previous article.

Latency budgets usually die on the embedding call plus the index round-trip. Caching query embeddings only helps when the same string repeats; product search and support bots often do not.

Exact search vs approximate indexes

Brute force: compute distance to every vector, sort, take top-k. Correct, and fine for tens of thousands of vectors in memory. Painful at millions under a tight p99.

Approximate nearest neighbor (ANN) indexes trade a little recall for speed. HNSW (Hierarchical Navigable Small World) is the graph-style structure you will see most often in pgvector, Qdrant, Weaviate, and friends: layered proximity graphs so search walks toward the query neighborhood instead of scanning everything. Other families (IVF, product quantization) push memory and disk harder with different recall curves.

What you need to remember as an app author:

  • ANN is approximate. Raising ef_search / probe counts usually improves recall and costs latency.
  • A bad index build (wrong metric, under-provisioned graph) looks like "the model is dumb" downstream.
  • Start exact or with high-recall settings while the corpus is small. Optimize when measurements say search is the bottleneck.

Metadata filters (tenant_id, product, lang) are not optional decoration. They cut the candidate set before or during search so you do not retrieve another customer's docs. Filter semantics differ by engine — some filter then search, some search then filter — and that changes recall when filters are selective.

A minimal embed-and-search loop

Below is the shape of an embeddings API call plus an in-memory top-k by cosine. Production code uses a real index; the arithmetic is the same.

package main

import (
  "bytes"
  "encoding/json"
  "fmt"
  "io"
  "math"
  "net/http"
  "os"
  "sort"
)

func embed(text string) ([]float64, error) {
  body, _ := json.Marshal(map[string]any{
    "model": "text-embedding-3-small",
    "input": text,
  })
  req, _ := http.NewRequest(
    "POST",
    "https://api.openai.com/v1/embeddings",
    bytes.NewReader(body),
  )
  req.Header.Set("Authorization", "Bearer "+os.Getenv("OPENAI_API_KEY"))
  req.Header.Set("Content-Type", "application/json")

  res, err := http.DefaultClient.Do(req)
  if err != nil {
    return nil, err
  }
  defer res.Body.Close()
  raw, _ := io.ReadAll(res.Body)

  var parsed map[string]any
  if err := json.Unmarshal(raw, &parsed); err != nil {
    return nil, err
  }
  first := parsed["data"].([]any)[0].(map[string]any)
  rawVec := first["embedding"].([]any)
  out := make([]float64, len(rawVec))
  for i, v := range rawVec {
    out[i] = v.(float64)
  }
  return out, nil
}

func cosine(a, b []float64) float64 {
  var dot, na, nb float64
  for i := range a {
    dot += a[i] * b[i]
    na += a[i] * a[i]
    nb += b[i] * b[i]
  }
  return dot / (math.Sqrt(na) * math.Sqrt(nb))
}

func main() {
  passages := []string{
    "Annual plans can be refunded within 14 days of purchase.",
    "API keys are rotated from the developer settings page.",
    "Incident severity is assigned by customer impact.",
  }

  type row struct {
    text  string
    score float64
  }

  q, err := embed("How do I get money back on a yearly subscription?")
  if err != nil {
    panic(err)
  }

  var ranked []row
  for _, p := range passages {
    v, err := embed(p)
    if err != nil {
      panic(err)
    }
    ranked = append(ranked, row{text: p, score: cosine(q, v)})
  }
  sort.Slice(ranked, func(i, j int) bool {
    return ranked[i].score > ranked[j].score
  })

  for _, r := range ranked {
    fmt.Printf("%.3f %s\n", r.score, r.text)
  }
}

Run that once and you will see the refund passage win. The interesting failures appear when you add near-miss passages, change the query to a product ID, or embed with a different model than the one that built a persisted index.

Dense, sparse, and hybrid

Dense retrieval (vectors) is strong on paraphrase and weak on exact tokens: error codes, SKUs, rare proper nouns, legal clause numbers. Sparse / lexical retrieval (BM25, Lucene-style inverted indexes) is the opposite: great when the right answer shares rare keywords, weak when the user paraphrases.

Hybrid search runs both and fuses the ranked lists — reciprocal rank fusion (RRF) is a common, parameter-light merge: score by position in each list, not by raw similarity magnitudes.

If your corpus is full of identifiers and quoted error strings, pure vector search will keep embarrassing you. If every question is a paraphrase of policy prose, BM25 alone will miss. Measure both before declaring a winner.

Reranking as a second stage

Top-k from an ANN index is a candidate generator, not a final judgment. A cross-encoder reranker (or a provider rerank API) scores (query, passage) pairs jointly and reshuffles the shortlist. It is slower per pair than a vector lookup, so you only run it on 20–100 candidates, not the whole corpus.

Typical pattern:

  1. Retrieve 50 with hybrid / dense.
  2. Rerank to 5–10.
  3. Pack those into the prompt under the token budget from Prompts & Context.

Reranking fixes a lot of "almost right chunk in position 7" pain. It does not fix an empty index, a wrong tenant filter, or an embedding model that cannot represent your domain.

Failure modes at the retrieval layer

Debug here before blaming the chat model:

  • Wrong neighbors, fluent answer: retrieval returned irrelevant snippets; the LLM filled gaps. Log id, score, and text of every injected chunk.
  • Right doc never appears: query embedding drift (instruction prefix missing), ANN recall too low, or filters excluding the row.
  • IDs and error codes miss: dense-only stack; add sparse / hybrid.
  • Scores look "good" but users hate results: uncalibrated thresholds; evaluate recall@k on a labeled query set instead of eyeballing cosine.
  • Index stale after doc updates: upsert/delete pipeline broken; generation looks current because the model still "knows" old training fluff.
  • Mixed embedding models in one collection: silent garbage. Version the model name next to the index.

A cheap habit: for every bad answer, print the retrieved set with scores. If the evidence was never there, no prompt rewrite will save you.

Comparison: what to reach for

ApproachStrengthsWeaknessesReach for it when
Dense onlyParaphrase; simple pipelineWeak on rare tokens; needs good embedderPolicy / prose corpora, natural-language questions
Sparse onlyExact keywords, cheap, explainableMisses paraphraseLogs, SKUs, error strings, legal citations
Hybrid + RRFCovers both failure modesMore moving partsMost production knowledge bases
+ RerankBest ordering of a shortlistExtra latency and costQuality-sensitive answers with a small final context

Most serious systems land on hybrid retrieval plus a rerank step once quality matters more than the demo.

Conclusion

Retrieval is a search system sitting in front of the prompt assembler: embed passages, index them, embed the query, fetch neighbors, optionally fuse and rerank. The generative model only sees what you inject. Keep embedding model, metric, and index version aligned, and treat top-k scores as candidates until something measures recall on your queries.

Next up is RAG — chunking those passages, grounding generation in what was retrieved, and the failure modes that appear only after search and generation are wired together.