Embeddings and Vector Search from Scratch

01From one-hot to dense

The obvious encoding of a word is a one-hot vector, and it is a bad one: every pair of words is equally far apart, so nothing about similarity can be read off the representation.^[1] An embedding replaces it with a short dense vector, learned so that the dot product or cosine between two vectors tracks a relation the task cares about. The transformer's first layer is a table of such vectors indexed by token id; a search engine's index is a table of such vectors indexed by document.

Two choices recur. Cosine similarity ignores vector length and compares direction only; the raw dot product also rewards length. Faiss, for one, has no separate cosine metric: cosine is obtained by normalizing every vector to unit length before indexing and before querying, after which inner product and cosine coincide.^[2]

02word2vec: learn vectors from company kept

The skip-gram model gives every word two vectors, one for use as a center word and one for use as a context word, and models the probability of a context word given a center word as a softmax over dot products across the whole vocabulary.^[1] Maximizing that likelihood over a corpus pulls together words that appear in the same contexts. The softmax is the problem: every update touches every vocabulary vector.

Negative sampling replaces the softmax with K + 1 binary decisions: the observed pair should score high, K sampled noise words should score low. The cost per pair is linear in K instead of the vocabulary size.^[3] The noise words are drawn in proportion to their count raised to the power 0.75, which flattens the distribution so rare words are drawn more often than their frequency alone would give.^[4][5]

The reference implementation's defaults are a window of 5, shrunk at each position to a random width between 1 and 5, 5 negatives, a 1e-3 subsampling threshold that randomly discards very frequent words, and a starting learning rate of 0.025 for skip-gram (0.05 for CBOW, its default architecture) that decays linearly over training.^[5] The noise table is built once from counts raised to 0.75 and sampled by index.

03Contrastive learning: pairs as supervision

Modern sentence and image embeddings are trained from pairs that should match: a question and its answer, a caption and its image, two translations. The in-batch contrastive loss takes a batch of N pairs, computes the N×N matrix of similarities, and asks each anchor to pick its own positive out of the N candidates with a softmax; every other positive in the batch is a free negative.^[6] A temperature sets how sharply the softmax discriminates. CLIP uses cosine similarities multiplied by a learned scale whose initial value is 1/0.07, which is a temperature of 0.07.^[7]

04Searching millions of vectors

Finding the nearest stored vector to a query by brute force costs one distance per stored vector. It is exact, and for a workload of a few thousand queries the Faiss guidelines call it the most efficient option, since an index would not repay its build time; Faiss calls this a flat index and uses it as the baseline every approximate index is measured against.^[8] At higher query rates an approximate index trades a little recall for a large speed-up, and when memory is the constraint, product quantization compresses each vector to a code of M bytes, typically 64 or fewer, at a further cost in accuracy.^[8]

HNSW is the index Faiss recommends when memory allows.^[8] Each vector is linked to near neighbours in a bottom layer, up to 2M of them in hnswlib, and up to M in each layer above; a random subset is promoted to a sparser layer above, a subset of those to the layer above that, and so on, each vector's top level drawn as ⌊−ln(U)/ln(M)⌋ for a uniform U.^[9] A query enters at the top, greedily steps to the closest linked node until no neighbour is closer, drops a layer, and repeats; at the bottom it widens into a beam of ef candidates. Higher ef means more distance evaluations and higher recall; higher M means better graphs at roughly M × 8 to 10 bytes per element of memory, with 12 to 48 covering most workloads.^[10]

05What high dimensions do to distance

Nearest-neighbour search gets harder with dimension for a reason that has nothing to do with the index. For random points, the distance from a query to its nearest neighbour and to its farthest grow closer together as coordinates are added, until almost every point is about the same distance away. One reason graph-based indexes still work on real embeddings is that those usually have far lower intrinsic dimension than the coordinate count, but the effect sets the limits.

06Retrieval for language models

A language model knows what was in its training data as of its cutoff. To answer about anything else, the standard pattern is retrieval: split reference documents into chunks, embed each chunk, store the vectors; at query time embed the question, find the most similar chunks, and put them into the prompt alongside the question.^[12] Embedding search suits questions because a question and its answer often share few words, which defeats lexical matching, while a trained pair embedding places them near each other.^[12]

The pieces of this page are the pieces of that system. The embedding model is trained with the §3 loss on (question, passage) pairs. The store is a flat or HNSW index over normalized vectors. The quality knobs are the chunk size, the number of chunks retrieved, the recall of the index at the chosen ef, and whether a lexical index is searched alongside and the results merged, which the cookbook lists among the ways to improve on embeddings-only search.^[12] The fine-tuning page covers the other way to change what a model produces: further training of its weights.

07Skip-gram training, checked

These complete programs generate the same kind of two-topic corpus, build a noise table from counts raised to 0.75, train skip-gram vectors with negative sampling for a few epochs, and assert that the geometry came out right: "cat" is closer by cosine to "dog" than to "car", "car" is closer to "truck" than to "sheep", and the training loss fell. The random number generator is a fixed linear congruential sequence so both languages run the same experiment.

#include <cassert>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>
struct Rng { uint64_t state = 12345; double next() { state = state * 6364136223846793005ULL + 1442695040888963407ULL; return double(state >> 11) / 9007199254740992.0; } };
const std::vector<std::string> vocabulary{"the", "in", "on", "cat", "dog", "cow", "sheep", "eats", "sleeps", "barn", "field", "car", "truck", "bus", "bike", "drives", "stops", "road", "garage"};
int id(const std::string& word) { for (std::size_t index = 0; index < vocabulary.size(); ++index) if (vocabulary[index] == word) return int(index); return -1; }
// Templated sentences: "the <subject> <verb> in|on the <place>", half farm, half traffic.
std::vector<int> corpus(Rng& rng, int sentences) {
    const std::vector<std::vector<std::string>> farm{{"cat", "dog", "cow", "sheep"}, {"eats", "sleeps"}, {"barn", "field"}};
    const std::vector<std::vector<std::string>> traffic{{"car", "truck", "bus", "bike"}, {"drives", "stops"}, {"road", "garage"}};
    std::vector<int> tokens;
    for (int sentence = 0; sentence < sentences; ++sentence) {
        const auto& topic = rng.next() < 0.5 ? farm : traffic;
        auto pick = [&](const std::vector<std::string>& options) { return id(options[std::size_t(rng.next() * options.size())]); };
        for (int token : {id("the"), pick(topic[0]), pick(topic[1]), rng.next() < 0.5 ? id("in") : id("on"), id("the"), pick(topic[2])}) tokens.push_back(token);
    }
    return tokens;
}
double sigmoid(double value) { return 1 / (1 + std::exp(-value)); }
double cosine(const std::vector<double>& left, const std::vector<double>& right) {
    double product = 0, leftNorm = 0, rightNorm = 0;
    for (std::size_t index = 0; index < left.size(); ++index) { product += left[index] * right[index]; leftNorm += left[index] * left[index]; rightNorm += right[index] * right[index]; }
    return product / std::sqrt(leftNorm * rightNorm);
}
int main() {
    Rng rng;
    const std::vector<int> tokens = corpus(rng, 600);
    const std::size_t dimension = 8, words = vocabulary.size();
    std::vector<double> counts(words, 0);
    for (int token : tokens) counts[token] += 1;
    std::vector<double> noiseWeight(words); double noiseTotal = 0;
    for (std::size_t word = 0; word < words; ++word) { noiseWeight[word] = std::pow(counts[word], 0.75); noiseTotal += noiseWeight[word]; } // word2vec's noise distribution.
    auto sampleNoise = [&]() { double threshold = rng.next() * noiseTotal; for (std::size_t word = 0; word < words; ++word) { threshold -= noiseWeight[word]; if (threshold <= 0) return int(word); } return int(words) - 1; };
    std::vector<std::vector<double>> center(words, std::vector<double>(dimension)), context(words, std::vector<double>(dimension, 0));
    for (auto& vector : center) for (double& value : vector) value = (rng.next() - 0.5) / dimension; // Small random centers, zero contexts, as in word2vec.
    const int window = 2, negatives = 5; const double rate = 0.05;
    double firstLoss = 0, lastLoss = 0;
    for (int epoch = 0; epoch < 8; ++epoch) {
        double loss = 0; int pairs = 0;
        for (int position = 0; position < int(tokens.size()); ++position)
            for (int offset = -window; offset <= window; ++offset) {
                if (offset == 0 || position + offset < 0 || position + offset >= int(tokens.size())) continue;
                const int centerWord = tokens[position];
                std::vector<double> gradientCenter(dimension, 0);
                for (int draw = 0; draw <= negatives; ++draw) {
                    const int target = draw == 0 ? tokens[position + offset] : sampleNoise(); // One observed pair, then K noise words.
                    const double label = draw == 0 ? 1 : 0;
                    double score = 0;
                    for (std::size_t coordinate = 0; coordinate < dimension; ++coordinate) score += center[centerWord][coordinate] * context[target][coordinate];
                    const double probability = sigmoid(score), error = probability - label; // ∂loss/∂score = σ − label.
                    loss -= label ? std::log(probability) : std::log(1 - probability);
                    for (std::size_t coordinate = 0; coordinate < dimension; ++coordinate) {
                        gradientCenter[coordinate] += error * context[target][coordinate];
                        context[target][coordinate] -= rate * error * center[centerWord][coordinate];
                    }
                }
                for (std::size_t coordinate = 0; coordinate < dimension; ++coordinate) center[centerWord][coordinate] -= rate * gradientCenter[coordinate];
                ++pairs;
            }
        if (epoch == 0) firstLoss = loss / pairs;
        lastLoss = loss / pairs;
    }
    assert(lastLoss < firstLoss);
    assert(cosine(center[id("cat")], center[id("dog")]) > cosine(center[id("cat")], center[id("car")]));     // Same topic beats other topic.
    assert(cosine(center[id("car")], center[id("truck")]) > cosine(center[id("car")], center[id("sheep")]));
}

struct Rng { state: u64 }
impl Rng { fn next(&mut self) -> f64 { self.state = self.state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); (self.state >> 11) as f64 / 9007199254740992.0 } }
const VOCABULARY: [&str; 19] = ["the", "in", "on", "cat", "dog", "cow", "sheep", "eats", "sleeps", "barn", "field", "car", "truck", "bus", "bike", "drives", "stops", "road", "garage"];
fn id(word: &str) -> usize { VOCABULARY.iter().position(|candidate| *candidate == word).unwrap() }
// Templated sentences: "the <subject> <verb> in|on the <place>", half farm, half traffic.
fn corpus(rng: &mut Rng, sentences: usize) -> Vec<usize> {
    let farm: [&[&str]; 3] = [&["cat", "dog", "cow", "sheep"], &["eats", "sleeps"], &["barn", "field"]];
    let traffic: [&[&str]; 3] = [&["car", "truck", "bus", "bike"], &["drives", "stops"], &["road", "garage"]];
    let mut tokens = Vec::new();
    for _ in 0..sentences {
        let topic = if rng.next() < 0.5 { &farm } else { &traffic };
        let subject = id(topic[0][(rng.next() * topic[0].len() as f64) as usize]);
        let verb = id(topic[1][(rng.next() * topic[1].len() as f64) as usize]);
        let preposition = if rng.next() < 0.5 { id("in") } else { id("on") };
        let place = id(topic[2][(rng.next() * topic[2].len() as f64) as usize]);
        tokens.extend_from_slice(&[id("the"), subject, verb, preposition, id("the"), place]);
    }
    tokens
}
fn sigmoid(value: f64) -> f64 { 1.0 / (1.0 + (-value).exp()) }
fn cosine(left: &[f64], right: &[f64]) -> f64 {
    let product: f64 = left.iter().zip(right).map(|(a, b)| a * b).sum();
    product / (left.iter().map(|a| a * a).sum::<f64>() * right.iter().map(|b| b * b).sum::<f64>()).sqrt()
}
fn main() {
    let mut rng = Rng { state: 12345 };
    let tokens = corpus(&mut rng, 600);
    let (dimension, words) = (8, VOCABULARY.len());
    let mut counts = vec![0.0_f64; words];
    for &token in &tokens { counts[token] += 1.0; }
    let noise_weight: Vec<f64> = counts.iter().map(|count| count.powf(0.75)).collect(); // word2vec's noise distribution.
    let noise_total: f64 = noise_weight.iter().sum();
    let mut center: Vec<Vec<f64>> = (0..words).map(|_| (0..dimension).map(|_| (rng.next() - 0.5) / dimension as f64).collect()).collect(); // Small random centers.
    let mut context = vec![vec![0.0; dimension]; words]; // Zero contexts, as in word2vec.
    let (window, negatives, rate) = (2i64, 5, 0.05);
    let (mut first_loss, mut last_loss) = (0.0, 0.0);
    for epoch in 0..8 {
        let (mut loss, mut pairs) = (0.0, 0);
        for position in 0..tokens.len() as i64 {
            for offset in -window..=window {
                if offset == 0 || position + offset < 0 || position + offset >= tokens.len() as i64 { continue; }
                let center_word = tokens[position as usize];
                let mut gradient_center = vec![0.0; dimension];
                for draw in 0..=negatives {
                    let target = if draw == 0 { tokens[(position + offset) as usize] } else { // One observed pair, then K noise words.
                        let mut threshold = rng.next() * noise_total;
                        let mut chosen = words - 1;
                        for (word, weight) in noise_weight.iter().enumerate() { threshold -= weight; if threshold <= 0.0 { chosen = word; break; } }
                        chosen
                    };
                    let label = if draw == 0 { 1.0 } else { 0.0 };
                    let score: f64 = center[center_word].iter().zip(&context[target]).map(|(a, b)| a * b).sum();
                    let probability = sigmoid(score);
                    let error = probability - label; // ∂loss/∂score = σ − label.
                    loss -= if label == 1.0 { probability.ln() } else { (1.0 - probability).ln() };
                    for coordinate in 0..dimension {
                        gradient_center[coordinate] += error * context[target][coordinate];
                        context[target][coordinate] -= rate * error * center[center_word][coordinate];
                    }
                }
                for coordinate in 0..dimension { center[center_word][coordinate] -= rate * gradient_center[coordinate]; }
                pairs += 1;
            }
        }
        if epoch == 0 { first_loss = loss / pairs as f64; }
        last_loss = loss / pairs as f64;
    }
    assert!(last_loss < first_loss);
    assert!(cosine(&center[id("cat")], &center[id("dog")]) > cosine(&center[id("cat")], &center[id("car")]));     // Same topic beats other topic.
    assert!(cosine(&center[id("car")], &center[id("truck")]) > cosine(&center[id("car")], &center[id("sheep")]));
}

08Where embeddings mislead

"Similar" is whatever the training pairs made it. Co-occurrence embeddings put interchangeable words together, which includes antonyms and the barn–garage pair in the first widget; pair-trained retrieval embeddings put questions near their answers. Pick the model by what its training objective rewarded, and measure it with recall on a labelled retrieval set rather than by eyeballing cosines.

An approximate index trades recall for speed through ef. Report the recall at the chosen ef alongside any retrieval result, measured against brute force on a sample of queries as the sixteen-dimensional widget does; a system tuned for throughput can quietly drop the right passage one time in ten.

Vectors from two different models, or two versions of one model, share no coordinate system. Re-embed everything when the model changes, and never mix vectors from different models in one index.

09What's next

The embedding table is the first layer of a language model. Transformers and Attention builds what comes after it: tokenization, attention that mixes the vectors across positions, position encoding, normalization and the key-value cache that makes generation tractable.

10Sources

The cited texts and reference implementations define the objectives, the index parameters and the retrieval pattern. Widget readouts report values computed in the page from seeded corpora, vectors and indexes.

1. Aston Zhang, Zachary C. Lipton, Mu Li, Alexander J. Smola, 2023. Dive into Deep Learning, Section 15.1: Word Embedding (word2vec), Cambridge University Press. One-hot vectors carrying no similarity, the skip-gram model with separate center and context vectors, and its softmax over the vocabulary.
2. Faiss contributors. MetricType and distances. Inner product versus cosine, and normalizing vectors before indexing and searching to obtain cosine similarity.
3. Aston Zhang et al., 2023. Dive into Deep Learning, Section 15.2: Approximate Training. The negative-sampling loss with K noise words and its cost linear in K.
4. Aston Zhang et al., 2023. Dive into Deep Learning, Section 15.3: The Dataset for Pretraining Word Embeddings. Noise words sampled in proportion to counts raised to 0.75, citing Mikolov et al. (2013).
5. Tomas Mikolov et al. word2vec.c, reference implementation. The unigram table built with exponent 0.75, defaults of 5 negatives, window 5 with a random per-position shrink, subsampling threshold 1e-3, and starting learning rate 0.025 for skip-gram, 0.05 for CBOW.
6. Sentence Transformers contributors. Loss overview. MultipleNegativesRankingLoss, also called InfoNCE or in-batch negatives, as the loss commonly used to train top-performing embedding models, and larger batches via caching.
7. OpenAI. CLIP, repository. Logits as cosine similarities times a learned scale initialized to 1/0.07.
8. Faiss contributors. Guidelines to choose an index. Flat indexes as the exact baseline, HNSW when memory allows with M links and efSearch as the speed-accuracy dial, and product quantization when memory is tight.
9. hnswlib contributors. hnswalg.h. The level multiplier 1/ln(M) and the random level −ln(U) × that multiplier.
10. hnswlib contributors. HNSW algorithm parameters. ef, M and ef_construction, memory of roughly M × 8 to 10 bytes per element, and M = 12 to 48 for most use cases.
11. Erik Bernhardsson et al. ann-benchmarks. Benchmarking approximate nearest-neighbour implementations by recall against speed.
12. OpenAI. Question answering using embeddings-based search. The chunk, embed, store, search, ask procedure; questions rarely sharing words with their answers; and combining search methods.