Embeddings and Contrastive Learning: Making Similar Mean Close

Features and Overfitting’s router saw each question as a bag of words: one column per word, a 1 if the word appears. That representation has a blind spot. cache and caching are unrelated columns, and so are token and tokens. Nothing about the representation says two words, or two questions, mean similar things.

An embedding fixes that by giving each thing a short dense vector, arranged so that similar things are close together. Every later part of the series depends on it: a language model’s first layer is an embedding table, retrieval (Vector Search) searches by embedding distance, and multimodal models (Seeing with Transformers) put images and text in one shared space. The question this post answers is where the arrangement comes from. The answer is always the same: from a loss that rewards putting some things close and others far apart. An embedding is only as meaningful as that loss, and as the data it saw.

Table of contents

From one-hot to dense

A one-hot vector for word $i$ of a $V$-word vocabulary is all zeros except a 1 in position $i$. Every two different words are equally far apart, and orthogonal: their dot product is 0. There’s no notion of similar. A dense embedding is a learned vector of, say, 64 numbers per word, stored as a $V \times 64$ table. Looking up word $i$ is the same as multiplying its one-hot vector by the table, so an embedding layer is just a linear layer applied to one-hot inputs, done efficiently.

Two ways to compare dense vectors:

Cosine similarity only measures direction; the dot product also grows with the vectors’ lengths. When all vectors have unit length they’re the same thing, which is why embeddings are usually normalized before comparing.

Word2vec: meaning from neighbors

The idea behind word2vec (Mikolov et al., 2013) is an old one in linguistics: a word is known by the company it keeps. Its skip-gram with negative sampling turns that into a loss. Take a word and one of the words near it in the text, a real (center, context) pair, and also $k$ random words from the vocabulary, fake pairs. Then train two embedding tables so that the dot product of a real pair is high and of a fake pair low, with Learning Is Optimization’s logistic loss:

Here $u_c$ is the center word’s vector, $v_o$ the real context word’s, and $v_{n_j}$ the random ones’. (These symbols, and $k$, are local to this section.) Words that appear in similar contexts get pushed toward the same context vectors, and so end up near each other. Trained on this blog’s English posts, prose, code and some inline SVG included, with the sections held out for the test later in this post removed (split.py is shown there):

The training loop of word2vec.py, after it has read the training text and kept the words seen five or more times. It makes two simplifications from the original method: frequent words aren’t subsampled, and every distance in the window counts equally.

torch.manual_seed(seed)
D, window, k = 64, 4, 5
center_emb = torch.nn.Embedding(len(vocab), D)
context_emb = torch.nn.Embedding(len(vocab), D)
torch.nn.init.normal_(center_emb.weight, 0, 0.1); torch.nn.init.zeros_(context_emb.weight)
freq = torch.tensor([counts[w] for w in vocab], dtype=torch.float) ** 0.75   # word2vec's sampling rule
opt = torch.optim.Adam(list(center_emb.parameters()) + list(context_emb.parameters()), lr=3e-3)
for step in range(4000):
    pos = torch.randint(window, len(ids) - window, (1024,))
    offset = torch.randint(1, window + 1, (1024,)) * (torch.randint(0, 2, (1024,)) * 2 - 1)
    c, o = ids[pos], ids[pos + offset]                                # a real (center, context) pair
    neg = torch.multinomial(freq, 1024 * k, replacement=True).view(1024, k)  # k random "contexts"
    vc = center_emb(c)
    pos_score = (vc * context_emb(o)).sum(-1)
    neg_score = torch.einsum("bd,bkd->bk", vc, context_emb(neg))
    # Logistic loss ("Learning Is Optimization"): real pairs toward 1, random pairs toward 0.
    loss = -(torch.nn.functional.logsigmoid(pos_score).mean()
             + torch.nn.functional.logsigmoid(-neg_score).sum(-1).mean())
    opt.zero_grad(); loss.backward(); opt.step()

Full script: word2vec.py

The script then saves the vectors, raw and normalized to unit length, as word2vec.pt, prints each test word’s six nearest neighbors by cosine, then the cosine and rank of three word pairs.

60,355 words, 1,685 kept (5+ occurrences)
final loss 2.071 (at the start: (1 + k) log 2 = 4.16)
attention   -> sparse, multi-head, hybrid, transformer, flashattention, gqa
cache       -> kv, kv-cache-explained, linearly, per-sequence, quantization, iteration-level
gradient    -> descent, alive, mse, convex, curvature, failures
token       -> representation, accepted, dimensional, speculative, decoder, single
vietnamese  -> english, questions, route, reader, foundations, list
gpu         -> h100, datasheet, laptop, tflops, cores, bandwidth
cos(cache, caching) = 0.40, rank 71 of 1,685
cos(token, tokens) = 0.49, rank 14 of 1,685
cos(cache, cached) = 0.39, rank 89 of 1,685

The loss falls from 4.16, what a model that knows nothing pays, to about 2.07. Nobody told it anything about the topics, yet the neighbors are what a reader of this blog would guess. Which of them are real, and which are the luck of one random start? Train with three seeds (more_seeds.py runs word2vec.py again with seeds 1 and 2) and count:

from collections import Counter
import torch

# Which neighbours are stable? Top 6 by cosine in each of three word2vec seeds.
seeds = [torch.load(f"word2vec{'' if s == 0 else s}.pt") for s in (0, 1, 2)]
for w in ("attention", "cache", "gradient", "token", "vietnamese", "gpu"):
    c = Counter()
    for m in seeds:
        vocab, E = m["vocab"], m["E"]
        i = vocab.index(w)
        c.update(vocab[j] for j in (E @ E[i]).argsort(descending=True)[1:7].tolist())
    print(f"{w:11s} in all 3 seeds: {', '.join(x for x, n in c.items() if n == 3) or '-'};  "
          f"in 2: {', '.join(x for x, n in c.items() if n == 2) or '-'}")

Full script: stable.py

attention   in all 3 seeds: sparse, multi-head, hybrid;  in 2: transformer
cache       in all 3 seeds: kv, kv-cache-explained, per-sequence;  in 2: linearly
gradient    in all 3 seeds: descent, alive, mse;  in 2: convex, failures
token       in all 3 seeds: representation;  in 2: accepted, decoder, re-read
vietnamese  in all 3 seeds: english, questions;  in 2: -
gpu         in all 3 seeds: h100, datasheet;  in 2: laptop, tflops, bandwidth

The stable neighbors are the right ones: gpu next to h100 and datasheet in every seed, cache next to kv and the KV cache post’s own name, gradient next to descent and mse, attention next to multi-head and sparse. The rest change from seed to seed, which is what 60,000 words buys. And the bag of words’ blind spot is partly fixed: one-hot put cache and caching at cosine 0, and here they’re at 0.40, token and tokens at 0.49, among each other’s nearest few percent.

The analogy test. Word2vec became famous for arithmetic like king − man + woman ≈ queen: the difference between two related words is a direction, and adding it to a third lands near the fourth. On this blog:

import torch

# The famous test: does b - a + c land nearest to d? ("king - man + woman ~ queen")
m = torch.load("word2vec.pt")
vocab, E = m["vocab"], m["E"]
index = {w: i for i, w in enumerate(vocab)}
tests = [("encoder", "decoder", "input", "output"), ("keys", "values", "key", "value"),
         ("prefill", "decode", "prompt", "generation"),
         ("layer", "layers", "token", "tokens"), ("head", "heads", "weight", "weights"),
         ("key", "keys", "value", "values")]
for a, b, c, d in tests:
    if not all(w in index for w in (a, b, c, d)):
        print(f"{a}:{b} :: {c}:{d}  (a word is missing from the vocabulary)")
        continue
    q = E[index[b]] - E[index[a]] + E[index[c]]
    sims = torch.nn.functional.normalize(q, dim=0) @ E.T
    for w in (a, b, c):
        sims[index[w]] = -1                          # the usual rule: exclude the question words
    ranked = sims.argsort(descending=True)
    rank = (ranked == index[d]).nonzero().item() + 1
    print(f"{a} : {b} :: {c} : ?   top answer '{vocab[ranked[0]]}', the expected '{d}' is ranked {rank} "
          f"of {len(vocab) - 3:,}")

Full script: analogies.py

encoder : decoder :: input : ?   top answer 'trigger', the expected 'output' is ranked 121 of 1,682
keys : values :: key : ?   top answer 'self', the expected 'value' is ranked 36 of 1,682
prefill : decode :: prompt : ?   top answer 'character', the expected 'generation' is ranked 168 of 1,682
layer : layers :: token : ?   top answer 'speculative', the expected 'tokens' is ranked 91 of 1,682
head : heads :: weight : ?   top answer 'neurips', the expected 'weights' is ranked 15 of 1,682
key : keys :: value : ?   top answer 're-read', the expected 'values' is ranked 9 of 1,682

Every one fails, even the easy plurals. A random guess would rank the right answer around 840th of 1,682, so these ranks (9th to 168th) show some signal, but none reaches the top. That’s worth knowing. The original results came from vectors trained on billions of words, and even those analogies depend on the benchmark’s conventions, such as excluding the question words from the answers. With 60,000 words, the vectors capture “these words appear together” well and consistent directions barely at all. The neighbors are useful; the arithmetic isn’t, at this scale.

Cosine vs dot, measured. Using the raw, unnormalized vectors:

import torch

# Dot product vs cosine on the raw word2vec vectors. The dot product is the
# cosine times both lengths, so a long vector rises in any list it's already
# near.
m = torch.load("word2vec.pt")
vocab, raw, counts = m["vocab"], m["raw"], torch.tensor(m["counts"], dtype=torch.float)
index = {w: i for i, w in enumerate(vocab)}
norms = raw.norm(dim=1)
r = torch.corrcoef(torch.stack([norms, counts.log()]))[0, 1]
print(f"vector length vs log word frequency: correlation {r:.2f}")
q = raw[index["cache"]]
for name, scores in [("dot product", raw @ q), ("cosine", (raw @ q) / (norms * q.norm()))]:
    scores[index["cache"]] = -1e9
    top = scores.argsort(descending=True)[:6]
    print(f"nearest to 'cache' by {name:11s}: " + ", ".join(f"{vocab[j]} ({norms[j]:.1f})" for j in top))
print(f"(vector lengths in brackets; the median word's is {norms.median():.1f})")

Full script: cosine_vs_dot.py

vector length vs log word frequency: correlation -0.20
nearest to 'cache' by dot product: kv (4.9), xhat (7.1), kv-cache-explained (5.4), nh (5.7), color-background (7.9), fill (7.4)
nearest to 'cache' by cosine     : kv (4.9), kv-cache-explained (5.4), linearly (4.6), per-sequence (4.9), quantization (3.8), iteration-level (4.9)
(vector lengths in brackets; the median word's is 4.3)

The dot product’s list is pulled toward words with long vectors, like xhat, nh, color-background and fill (lengths 5.7–7.9 against a median of 4.3), which are variable names and SVG attributes from the posts’ code, not related words. Length here has little to do with how common a word is (the correlation is −0.20); the longest vectors belong to code tokens whose contexts are nearly fixed. A long vector multiplies whatever its cosine is, so it rises in any list it’s already near. Cosine removes the length and gives the sensible list. Normalize before comparing, unless the length is meant to carry information.

InfoNCE: a cross-entropy over one right answer

Word2vec’s loss treats each pair separately. The loss behind most modern embeddings instead makes a choice: given a query, pick its one true match out of a set of candidates. That’s InfoNCE (van den Oord, Li and Vinyals, 2018), and it is exactly the cross-entropy of Cross-Entropy and the Softmax Gradient, with the candidates as the classes. For a batch of $B$ pairs $(a_i, b_i)$, compute every similarity $s_{ij} = \cos(a_i, b_j)$ and, for each row, the loss for picking the diagonal ($a$, $b$ and $s$ are local to this section):

The other $B - 1$ candidates in the batch are the negatives: the same data used twice, no extra cost. (Many models, including CLIP, use both directions, each $a$ picking its $b$ and each $b$ its $a$, and average the two losses.) Worked on three pairs:

import math
import torch

def info_nce(a, b, tau):
    """Row i of a should match row i of b and no other row: cross-entropy over
    the similarity matrix, with the diagonal as the correct 'class' ("Cross-Entropy and the Softmax Gradient")."""
    a, b = torch.nn.functional.normalize(a, dim=1), torch.nn.functional.normalize(b, dim=1)
    logits = a @ b.T / tau
    return torch.nn.functional.cross_entropy(logits, torch.arange(len(a))), a @ b.T

# Three pairs of 2-D embeddings: each b is a slightly rotated copy of its a.
a = torch.tensor([[1.0, 0.0], [0.0, 1.0], [-1.0, 0.2]])
b = torch.tensor([[0.9, 0.3], [0.2, 1.0], [-1.0, -0.1]])
loss, sim = info_nce(a, b, 1.0)
print("cosine similarities (rows: a, columns: b):")
for row in sim:
    print("   " + "  ".join(f"{v:+.3f}" for v in row))
for tau in (1.0, 0.5, 0.1, 0.05):
    print(f"temperature {tau:<4}: loss {info_nce(a, b, tau)[0]:.4f}")

# Random, unrelated embeddings: the loss is about log(batch size), a uniform
# guess among B ("Cross-Entropy and the Softmax Gradient"), and the best possible loss is 0. The bound:
# mutual information >= log B - loss. (At a low temperature, random
# similarities spread out and the loss goes above log B.)
torch.manual_seed(0)
for B in (4, 64, 1024):
    x, y = torch.randn(B, 64), torch.randn(B, 64)
    print(f"batch {B:5d}: random embeddings give loss {info_nce(x, y, 1.0)[0]:.2f} at temperature 1, "
          f"{info_nce(x, y, 0.1)[0]:.2f} at 0.1; log B = {math.log(B):.2f}")

Full script: infonce.py

cosine similarities (rows: a, columns: b):
   +0.949  +0.196  -0.995
   +0.316  +0.981  -0.100
   -0.868  +0.000  +0.956
temperature 1.0 : loss 0.5106
temperature 0.5 : loss 0.2332
temperature 0.1 : loss 0.0006
temperature 0.05: loss 0.0000
batch     4: random embeddings give loss 1.43 at temperature 1, 2.11 at 0.1; log B = 1.39
batch    64: random embeddings give loss 4.17 at temperature 1, 4.93 at 0.1; log B = 4.16
batch  1024: random embeddings give loss 6.94 at temperature 1, 7.68 at 0.1; log B = 6.93

• The similarity matrix is the whole story. Each $a_i$ is most similar to its own $b_i$ (0.949, 0.981, 0.956 on the diagonal), so the loss is already low.
• The temperature $\tau$ divides the similarities before the softmax, as in Cross-Entropy and the Softmax Gradient. Cosines live between $-1$ and $1$, so at $\tau = 1$ the softmax can’t be confident: even these well-separated pairs pay 0.51. At $\tau = 0.1$ the same similarities give a loss of 0.0006. A low temperature makes the loss care about small differences in similarity; most contrastive models use somewhere around 0.05 to 0.1, or learn it.
• Random embeddings pay about $\log B$ at $\tau = 1$. With unrelated vectors each row is a uniform guess among $B$: 1.43 for 4 candidates, 6.94 for 1,024, against $\log B$ = 1.39 and 6.93. At a low temperature they pay more, because random similarities, spread out by the division, look like confident wrong answers.
• The loss bounds an information. The mutual information between two views measures how many nats knowing one saves about the other. Van den Oord et al. showed it is at least $\log B - \mathcal{L}$ (for the expected loss, with independent negatives; Poole et al., 2019, give the careful version). A larger batch raises that ceiling, one reason contrastive models train with thousands of in-batch negatives. The bound is loose, it can never certify more than $\log B$, and when the loss is above $\log B$ it says nothing.

Drag the points below. Each colored circle $a_i$ should be closest to its own square $b_i$; the table is the similarity matrix and the loss is InfoNCE:

A contrastive head on a tiny corpus

SimCLR (Chen et al., 2020) made the recipe general: take one thing, make two views of it by random changes that shouldn’t change what it is, and train an encoder with InfoNCE so the two views match and differ from everything else in the batch. For images the views are crops, color changes and flips. For text, the simplest view is a random crop, the recipe behind Contriever (Izacard et al., 2021): two 40-word windows from the same blog section are a positive pair.

The test uses sections nothing trains on, 30% of all sections held out (word2vec above was trained without them too): the first 40 words of each must find the same section’s last 40 words among all of them (recall@1). The encoder is deliberately small: word2vec vectors, frozen, averaged over the window, then one learned $64 \times 64$ linear map that starts as the identity. So before training it is the word2vec average, and training can only rotate, stretch and shift it. First, split.py holds the test sections out, once for every script in the post:

import re
import torch

# Split the blog into sections of raw words once: 30% of sections are held
# out for testing and never trained on, not even by word2vec. Everything
# else, sections and the text between them, is training text.
VIEW = 40                                                      # words per view
text = open("blog_train.txt").read() + open("blog_val.txt").read()
chunks = re.split(r"\n## ", text)
long_enough = [i for i, c in enumerate(chunks) if len(c.split()) >= 3 * VIEW]
order = torch.randperm(len(long_enough), generator=torch.Generator().manual_seed(0))
held_out = {long_enough[i] for i in order[:len(long_enough) * 3 // 10]}
test_secs = [chunks[i].split() for i in sorted(held_out)]
train_secs = [chunks[i].split() for i in long_enough if i not in held_out]
train_text = "\n## ".join(c for i, c in enumerate(chunks) if i not in held_out)

Full script: split.py

sections.py loads the word2vec vectors and holds two helpers that contrastive.py and pretrained.py share: to_ids turns a window of words into the ids of the words word2vec knows, and crop cuts a random 40-word window out of a section.

Full script: sections.py

Then the encoder, in contrastive.py:

class Encoder(torch.nn.Module):
    """Average the (frozen) word2vec vectors, then a learned linear map."""
    def __init__(self, init):
        super().__init__()
        self.emb = torch.nn.Embedding.from_pretrained(init.clone(), freeze=True)
        self.proj = torch.nn.Linear(init.shape[1], init.shape[1])
        torch.nn.init.eye_(self.proj.weight); torch.nn.init.zeros_(self.proj.bias)
    def forward(self, windows):                                # a list of word windows
        pooled = torch.stack([self.emb(to_ids(w)).mean(0) for w in windows])
        return torch.nn.functional.normalize(self.proj(pooled), dim=-1)

Each step draws a batch of sections, crops two views of each, and applies InfoNCE at temperature 0.1, three seeds per batch size:

for B in (4, 16, 64):
    scores, fit = [], []
    for seed in range(3):
        g = torch.Generator().manual_seed(seed); torch.manual_seed(seed)
        enc = Encoder(W2V)
        opt = torch.optim.Adam(enc.parameters(), lr=1e-3)
        for step in range(1500):
            pick = torch.randint(0, len(train_secs), (B,), generator=g)
            x1 = [crop(train_secs[i], g) for i in pick]
            x2 = [crop(train_secs[i], g) for i in pick]
            logits = enc(x1) @ enc(x2).T / 0.1                  # InfoNCE, temperature 0.1
            same = pick[:, None] == pick[None, :]               # repeats in a batch aren't negatives
            logits = logits.masked_fill(same & ~torch.eye(B, dtype=torch.bool), float("-inf"))
            loss = torch.nn.functional.cross_entropy(logits, torch.arange(B))
            opt.zero_grad(); loss.backward(); opt.step()
        scores.append(recall_at_1(enc))
        fit.append(recall_at_1(enc, train_secs[:len(test_secs)]))   # same number of training sections
    print(f"contrastive, batch {B:3d} ({B - 1:2d} negatives): recall@1 "
          + ", ".join(f"{s:.3f}" for s in scores) + " (3 seeds); on training sections "
          + ", ".join(f"{s:.3f}" for s in fit))

Full script: contrastive.py

The rest of the script scores the baselines: averaged random vectors and averaged word2vec vectors from each of the three seeds.

156 training sections, 66 held-out sections
before training, the encoder is exactly the word2vec average: recall@1 0.152
random word vectors, averaged:   recall@1 0.061
word2vec vectors, averaged:      recall@1 0.152, 0.136, 0.136 (3 word2vec seeds)
contrastive, batch   4 ( 3 negatives): recall@1 0.197, 0.152, 0.212 (3 seeds); on training sections 0.121, 0.121, 0.121
contrastive, batch  16 (15 negatives): recall@1 0.167, 0.182, 0.212 (3 seeds); on training sections 0.136, 0.152, 0.136
contrastive, batch  64 (63 negatives): recall@1 0.212, 0.227, 0.212 (3 seeds); on training sections 0.182, 0.182, 0.167

With 66 test sections, one extra hit is 1.5 points of recall, chance is 1.5%, and the binomial standard error near 20% is about ±5 points. Read with that in mind:

• Word vectors help over nothing. Averaged random vectors, which amount to a random projection of the bag of words, find the right section 6% of the time, four times chance. Averaged word2vec vectors find 14–15%, across three word2vec seeds.
• The contrastive head helps a little, and more with a bigger batch. With 63 negatives per row, all three seeds land at 21–23%, about six and a half points above the average it started from, roughly 1.3 standard errors. With 3 negatives it’s 15–21%. More negatives, a harder choice to learn from, as the bound above suggests.
• It isn’t memorizing. Its recall on the same number of training sections is no higher than on held-out ones.

That’s what 156 sections and a linear map can do. The encoders below were trained on vastly more pairs, with a full Transformer to adjust instead of one matrix.

Pretrained sentence encoders

Two pretrained encoders, on the same held-out task and on a second one: match each section of the posts that have a Vietnamese translation to its translation, among all the translated sections. One, all-MiniLM-L6-v2, was trained contrastively on about a billion English sentence pairs. The other got its languages by distillation (Reimers and Gurevych, 2020): an English teacher model, itself trained contrastively, and a multilingual student trained with a squared-error loss so that a sentence and its translation both land where the teacher puts the English sentence:

def recall(model, queries, docs):
    q = model.encode(queries, normalize_embeddings=True, convert_to_tensor=True)
    d = model.encode(docs, normalize_embeddings=True, convert_to_tensor=True)
    return ((q @ d.T).argmax(1) == torch.arange(len(q), device=q.device)).float().mean().item()

def spread(E):
    """Average cosine between *different* held-out English windows."""
    S = E @ E.T
    return S[~torch.eye(len(E), dtype=torch.bool)].mean().item()

w2v = torch.nn.functional.normalize(torch.stack([W2V[to_ids(s[:VIEW])].mean(0) for s in test_secs]), dim=-1)
print(f"averaged word2vec: average cosine between different held-out windows {spread(w2v):.2f}")
for name, path in models.items():
    model = SentenceTransformer(path, device="cpu", revision=REVISION[path])
    params = sum(p.numel() for p in model.parameters())
    same = recall(model, first, last)
    t0 = time.perf_counter()
    cross = recall(model, [e for e, _ in pairs], [v for _, v in pairs])
    secs = time.perf_counter() - t0
    E = model.encode(first, normalize_embeddings=True, convert_to_tensor=True)
    print(f"{name}: {params / 1e6:.0f}M parameters, reads up to {model.max_seq_length} tokens")
    print(f"    same-section recall@1 {same:.3f}   English -> Vietnamese recall@1 {cross:.3f} "
          f"({2 * len(pairs)} sections in {secs:.1f} s)   average cosine between different windows {spread(E):.2f}")

Full script: pretrained.py

The start of the script pins both models to a revision and pairs the English and Vietnamese sections of the blog snapshot that tar xzf blog-47b01b1.tar.gz unpacks into content/.

66 held-out English sections; 118 English/Vietnamese section pairs
averaged word2vec: average cosine between different held-out windows 0.81
all-MiniLM-L6-v2 (English): 23M parameters, reads up to 256 tokens
    same-section recall@1 0.318   English -> Vietnamese recall@1 0.576 (236 sections in 2.0 s)   average cosine between different windows 0.16
multilingual MiniLM-L12: 118M parameters, reads up to 128 tokens
    same-section recall@1 0.212   English -> Vietnamese recall@1 1.000 (236 sections in 2.0 s)   average cosine between different windows 0.29

(The table-of-contents and reference sections are left out: the first are only a heading, identical in every post, and the second are near-identical citation lists. Both models read only the first 128 to 256 tokens of each section, so the translation task is mostly “match the heading and first paragraph”.)

• On the same-section task, the English model is clearly best. It finds 32%, against 14–15% for the word2vec average, more than two standard errors apart, and the contrastive head’s 21–23% in between. The multilingual model, 21%, trades some English skill for its languages.
• Across languages, the training pairs decide everything. The English model finds the right Vietnamese section 58% of the time, probably helped by the English technical terms, headings and code that the Vietnamese posts keep. The multilingual model finds all 118. It’s also larger, 118 million parameters to 23 million, mostly for a vocabulary covering dozens of languages, but what matters is what it was trained to do: put a sentence and its translation at the same point.
• Embedding is cheap. 236 sections, truncated as above, took about 2 seconds on this laptop’s CPU.

Anisotropy. The last column is the average cosine between two different held-out windows. For averaged word2vec vectors it’s 0.81: every window points in nearly the same direction, and only small differences carry meaning. The two pretrained encoders, one trained contrastively and one distilled from a contrastive teacher, spread the same windows out, to 0.16 and 0.29. Representations learned without a contrastive loss tend to crowd into a narrow cone of the space (anisotropy; Ethayarajh, 2019, measured it in language models), and contrastive training is one way to spread them apart, pulling matches together while pushing everything else toward a uniform spread (Wang and Isola, 2020). Subtracting the mean embedding before comparing is a cheaper partial fix.

A practical note for retrieval: these are bi-encoders. Each text is embedded on its own, by averaging the Transformer’s token vectors, so a corpus can be embedded once and searched fast. A cross-encoder reads the query and a candidate together and scores the pair, much more accurately and much more slowly; that’s the reranker that Reranking and Late Interaction adds.

Where this leads

• What “Generative” Means puts the character model of Sequences Before Attention and these embeddings together: a language model’s first layer is an embedding table, trained by next-token prediction instead of a contrastive loss.
• Vector Search, Semantic and Hybrid Search and Reranking and Late Interaction build retrieval for the Notebook Assistant on sentence embeddings like these, add keyword search for what embeddings miss, and a reranker.
• Seeing with Transformers trains the same InfoNCE loss on image–caption pairs: that’s CLIP.

Takeaways

• A one-hot vector says nothing about similarity; a dense embedding is a learned table, and an embedding lookup is a linear layer on a one-hot input.
• Word2vec learns vectors from which words appear near each other, with a logistic loss on real vs random pairs. On 60,000 words the stable neighbors are sensible and the analogies all fail.
• Cosine measures direction only; the dot product also rewards long vectors, and here pulled in code and SVG tokens. Normalize first.
• InfoNCE is cross-entropy over one positive and $B - 1$ in-batch negatives. Random embeddings pay about $\log B$; the temperature decides how sharply small similarity differences count; more negatives give a higher ceiling on what the loss can measure.
• A contrastive head on frozen word vectors, trained on 156 sections, added about six and a half points of recall, most with the largest batch. A pretrained encoder, trained on a billion pairs, added more than twice that.
• For cross-lingual search, the training pairs are everything: an English encoder matched 58% of sections to their Vietnamese translations, a multilingual one all of them.
• Averaged word vectors crowd into a narrow cone (anisotropy); contrastive training spreads embeddings out.

References

Start with Jay Alammar’s The Illustrated Word2vec for intuition, then the SimCLR paper for the contrastive recipe and Wang and Isola for why it spreads embeddings out.

• Mikolov, Chen, Corrado and Dean, Efficient Estimation of Word Representations in Vector Space (2013): skip-gram and CBOW.
• Mikolov, Sutskever, Chen, Corrado and Dean, Distributed Representations of Words and Phrases and their Compositionality (NeurIPS, 2013): skip-gram with negative sampling.
• van den Oord, Li and Vinyals, Representation Learning with Contrastive Predictive Coding (2018): InfoNCE and its mutual-information bound.
• Poole, Ozair, van den Oord, Alemi and Tucker, On Variational Bounds of Mutual Information (ICML, 2019): the careful version of the InfoNCE bound.
• Chen, Kornblith, Norouzi and Hinton, A Simple Framework for Contrastive Learning of Visual Representations (ICML, 2020): SimCLR, with views, in-batch negatives and temperature.
• Izacard et al., Unsupervised Dense Information Retrieval with Contrastive Learning (TMLR, 2022): Contriever, with random crops as views for text.
• Wang and Isola, Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere (ICML, 2020): why contrastive training spreads embeddings out.
• Reimers and Gurevych, Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks (EMNLP, 2019): Sentence-BERT, the background for the sentence encoders used here.
• Reimers and Gurevych, Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation (EMNLP, 2020): the multilingual model used here.
• Sentence Transformers, all-MiniLM-L6-v2 (model card): contrastive training on about a billion sentence pairs.
• Ethayarajh, How Contextual are Contextualized Word Representations? (EMNLP, 2019): anisotropy of language-model embeddings.
• Radford et al., Learning Transferable Visual Models From Natural Language Supervision (ICML, 2021): CLIP, the same loss on images and text.