Embeddings & Representation Learning

An embedding is a learned mapping from a discrete or high-dimensional input (a word token, an image, a user ID) to a dense vector in ℝᵈ, typically d ∈ {64, 256, 768, 1536, ...}.

Why dense vectors? Contrast with one-hot encoding. A vocabulary of 50,000 words yields 50,000-dimensional one-hot vectors. These are orthogonal — every word is equally far from every other word — which destroys all semantic information. A dense embedding in ℝ³⁰⁰ can place semantically similar words nearby in Euclidean space, encoding similarity as cosine distance.

The distributional hypothesis (Harris 1954, Firth 1957): words that appear in similar contexts have similar meanings. This is the empirical basis for all word embeddings. If 'cat' and 'dog' both co-occur frequently with 'vet,' 'leash,' 'fur,' and 'pet,' their embeddings should be nearby.

Geometric structure: After training, embedding spaces exhibit surprising linear structure:

• king − man + woman ≈ queen (analogy arithmetic)
• Paris − France + Italy ≈ Rome (relational knowledge)

This linearity is not guaranteed by theory — it emerges from optimization on the right objective. The geometry encodes the statistical structure of the training corpus.

Embedding lookup as a matrix multiply: The embedding layer is just a matrix E ∈ ℝ^{|V| × d}. Looking up token i returns row E[i]. This is equivalent to the one-hot vector eᵢ times E: eᵢ · E. During training, gradients flow back to update exactly the rows for tokens seen in the batch.

import torch
import torch.nn as nn

# Embedding layer: vocabulary size 10000, dimension 256
emb = nn.Embedding(num_embeddings=10000, embedding_dim=256)

# Forward pass: token ids → dense vectors
tokens = torch.tensor([42, 7, 1337])   # shape (3,)
vecs = emb(tokens)                      # shape (3, 256)

# Cosine similarity between two embeddings
def cosine_sim(a, b):
    return (a @ b) / (a.norm() * b.norm())

print(cosine_sim(vecs[0], vecs[1]).item())

Word2Vec (Mikolov et al., 2013) is the algorithm that made representation learning mainstream. It trains a shallow two-layer network on a self-supervised prediction task, then discards the network and keeps the learned weight matrix as embeddings.

Two architectures:

1. CBOW (Continuous Bag-of-Words): Given the surrounding context words, predict the center word. Fast, works well for frequent words.
2. Skip-Gram: Given the center word, predict each surrounding context word. Slower but produces better embeddings for rare words — the default choice in practice.

Skip-Gram objective: For each center word w and each context word c within a window of size k:

maximize Σ log P(c | w)

where P(c | w) = softmax(u_c · v_w) = exp(u_c · v_w) / Σ_{c'} exp(u_c' · v_w)

Note two weight matrices: V (input/center embeddings) and U (output/context embeddings). At the end you typically use V, or average V and U.

The softmax bottleneck: Computing the denominator over the full vocabulary is O(|V|) per step — far too slow for |V| = 50,000.

Negative Sampling solves this. Instead of normalizing over all words, treat the problem as binary classification: given pair (w, c), predict whether c is a true context word (positive) or a random noise word (negative). Sample k noise words per positive. The loss becomes:

L = -log σ(u_c · v_w) - Σᵢ₌₁ᵏ log σ(-u_{nᵢ} · v_w)

where σ is the sigmoid. This reduces each step to O(k) instead of O(|V|). With k=5–20, training on billions of tokens becomes tractable on a laptop.

Noise distribution: Words are sampled with probability proportional to f(w)^(3/4), where f(w) is unigram frequency. The 3/4 power smooths out the distribution, preventing very common words from dominating the negative samples.

import numpy as np

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))

def skipgram_neg_sample_step(center, context, negatives, V, U, lr=0.01):
    """
    One gradient step for skip-gram with negative sampling.
    V: input embeddings  [vocab, dim]
    U: output embeddings [vocab, dim]
    """
    v_c = V[center]          # center vector
    u_pos = U[context]       # positive context vector

    # Positive pair: maximize u_pos · v_c
    sig_pos = sigmoid(u_pos @ v_c)
    grad_vc  = (sig_pos - 1) * u_pos
    grad_upos = (sig_pos - 1) * v_c

    U[context] -= lr * grad_upos

    # Negative pairs: maximize -u_neg · v_c (i.e. push apart)
    for neg in negatives:
        u_neg = U[neg]
        sig_neg = sigmoid(u_neg @ v_c)
        grad_vc  += sig_neg * u_neg
        U[neg]   -= lr * sig_neg * v_c

    V[center] -= lr * grad_vc

GloVe (Pennington et al., 2014 — Stanford NLP) takes a different route to the same destination. Instead of scanning through a corpus word by word, it first computes the global co-occurrence matrix X, where X_{ij} = number of times word j appears in the context of word i. Then it factors this matrix into embeddings.

Key insight: The ratio of co-occurrence probabilities is more informative than the raw probabilities. For probe word k:

• P('solid' | 'ice') / P('solid' | 'steam') ≫ 1 → 'solid' discriminates ice vs steam
• P('water' | 'ice') / P('water' | 'steam') ≈ 1 → 'water' is uninformative

GloVe models: log X_{ij} ≈ wᵢ · w̃ⱼ + bᵢ + b̃ⱼ

The weighted least-squares objective:

J = Σᵢⱼ f(X_{ij}) (wᵢ · w̃ⱼ + bᵢ + b̃ⱼ − log X_{ij})²

where f(x) = (x/x_max)^α for x < x_max, else 1 (a weighting function that caps the influence of very frequent pairs).

Connection to SVD: It was later shown (Levy & Goldberg, 2014) that Word2Vec with negative sampling implicitly factorizes the Pointwise Mutual Information (PMI) matrix. PMI(w,c) = log[ P(w,c) / P(w)P(c) ]. This unified the apparently different approaches: both GloVe and Word2Vec are doing matrix factorization on statistics of word co-occurrence.

PMI defined: PMI(w, c) = log [P(w,c) / (P(w) · P(c))]. High PMI = words co-occur much more than by chance. The embedding dot product wᵢ · w̃ⱼ approximates PMI(i,j) after training.

# Conceptual: build co-occurrence matrix then factorize with SVD
import numpy as np
from scipy.sparse.linalg import svds

def build_cooccurrence(corpus_tokens, vocab, window=5):
    V = len(vocab)
    w2i = {w: i for i, w in enumerate(vocab)}
    X = np.zeros((V, V))
    for t, token in enumerate(corpus_tokens):
        if token not in w2i:
            continue
        i = w2i[token]
        ctx_start = max(0, t - window)
        ctx_end   = min(len(corpus_tokens), t + window + 1)
        for s in range(ctx_start, ctx_end):
            if s != t and corpus_tokens[s] in w2i:
                j = w2i[corpus_tokens[s]]
                X[i, j] += 1.0
    return X

# PPMI (Positive PMI) transformation
def ppmi(X):
    total = X.sum()
    row_sums = X.sum(axis=1, keepdims=True)
    col_sums = X.sum(axis=0, keepdims=True)
    with np.errstate(divide='ignore', invalid='ignore'):
        M = np.log2((X * total) / (row_sums * col_sums))
    M[~np.isfinite(M)] = 0.0
    M = np.maximum(M, 0.0)   # positive PMI
    return M

# SVD gives dense embeddings
# U, s, Vt = svds(ppmi(X), k=300)
# embeddings = U * np.sqrt(s)   # shape (vocab_size, 300)

Contrastive learning is a paradigm for learning representations without labels. The idea: define what it means for two views of data to be 'similar,' then train a network to pull similar pairs together and push dissimilar pairs apart in embedding space.

Intuition: If you crop two overlapping patches from the same photo, the high-level content (a dog, a sunset) is the same even though the pixels differ. A good encoder should map both crops to nearby vectors. Random crops from different images should be far apart.

SimCLR framework (Chen et al., 2020):

1. For each image x in a batch, create two augmented views x̃₁, x̃₂ (random crop, color jitter, blur)
2. Encode both: zᵢ = g(f(x̃ᵢ)) where f is the backbone (ResNet) and g is a small projection MLP
3. Maximize agreement between positive pairs, minimize it for negative pairs (all other images in the batch)

InfoNCE loss (van den Oord et al., 2018, CPC paper):

L = -log [ exp(z₁ · z₂ / τ) / Σⱼ exp(z₁ · zⱼ / τ) ]

This is just cross-entropy over a 1-of-N classification problem: given z₁, identify its true positive z₂ among all N vectors in the batch. τ (temperature) controls sharpness — low τ makes the distribution peakier.

Why does it work? InfoNCE is a lower bound on the mutual information between the two views: I(Z₁; Z₂) ≥ log(N) - L. Minimizing the loss maximizes a bound on mutual information between views. Intuitively: the encoder learns to preserve only the information shared across augmentations — invariant, high-level features — while discarding augmentation-specific noise.

Collapse problem: A trivial solution is to map all inputs to the same constant vector. Contrastive methods avoid this via large batches of negatives (SimCLR uses batch size 4096+). Non-contrastive methods (BYOL, SimSiam) use momentum encoders or stop-gradient tricks instead.

import torch
import torch.nn.functional as F

def infonce_loss(z1, z2, temperature=0.07):
    """
    z1, z2: (N, D) normalized embeddings from two views.
    Returns scalar InfoNCE loss.
    """
    N = z1.size(0)
    # Normalize onto unit sphere
    z1 = F.normalize(z1, dim=1)   # (N, D)
    z2 = F.normalize(z2, dim=1)   # (N, D)

    # Similarity matrix: (N, N), entry [i,j] = z1[i]·z2[j] / τ
    logits = (z1 @ z2.T) / temperature   # (N, N)

    # Positive pairs are on the diagonal (index i matches index i)
    labels = torch.arange(N, device=z1.device)

    # Cross-entropy: for each row i, classify which column j is the positive
    loss = F.cross_entropy(logits, labels)
    return loss

# Symmetric version (used in SimCLR): average both directions
def simclr_loss(z1, z2, temperature=0.07):
    l1 = infonce_loss(z1, z2, temperature)
    l2 = infonce_loss(z2, z1, temperature)
    return (l1 + l2) / 2

Transfer learning is the practice of taking a model trained on one large task and adapting it to a different, usually smaller task. It is the dominant paradigm in modern ML: almost nothing of value is trained from scratch.

Why does transfer work? Deep networks learn hierarchical features. Early layers in a vision network learn edges and textures; middle layers learn shapes and parts; late layers learn object-level semantics. These low and mid-level features are useful for almost any visual task. Starting from pretrained weights means you already have these features — fine-tuning just adapts the top layers to the new task.

The regime spectrum:

Feature extraction: Freeze the pretrained backbone entirely. Pass all target examples through it, collect the final-layer activations as fixed feature vectors, then train any classifier (SVM, logistic regression, MLP) on those features. Fast, no GPU needed for the downstream task.

Fine-tuning: Unfreeze some or all layers, continue training with a small learning rate (1e-5 vs 1e-3 typical) to avoid catastrophic forgetting — overwriting pretrained knowledge with noisy gradient updates from a small dataset.

Catastrophic forgetting: Empirically, fine-tuning with too large a learning rate on too few examples destroys pretrained performance. Mitigations: gradual unfreezing (ULMFiT: Howard & Ruder 2018), layer-wise learning rate decay (lower LR for early layers), elastic weight consolidation (EWC).

ImageNet as the universal backbone: From 2012–2020, AlexNet/VGG/ResNet pretrained on ImageNet-1k (1M images, 1000 classes) were the universal starting point for almost all computer vision research. This pretrain-then-finetune paradigm directly drove most computer vision benchmarks down.

import torch
import torchvision.models as models
import torch.nn as nn

# Load pretrained ResNet-50 (ImageNet weights)
model = models.resnet50(weights='IMAGENET1K_V1')

# --- Option A: Feature extraction (freeze everything) ---
for param in model.parameters():
    param.requires_grad = False

# Replace final FC layer for a 10-class target task
model.fc = nn.Linear(model.fc.in_features, 10)
# Only model.fc parameters will be updated

# --- Option B: Fine-tuning with layer-wise LR decay ---
backbone_params = [p for n, p in model.named_parameters()
                   if 'fc' not in n]
head_params     = list(model.fc.parameters())

optimizer = torch.optim.AdamW([
    {'params': backbone_params, 'lr': 1e-5},  # small LR for backbone
    {'params': head_params,     'lr': 1e-3},  # larger LR for head
])

Self-supervised learning (SSL) generates free supervision from the structure of unlabeled data itself, without any human annotation. The model is trained on an artificially constructed pretext task, and the representations learned in service of that task generalize broadly.

Classic pretext tasks:

• Colorization (Zhang et al., 2016): predict the color of a grayscale image patch. Forces the network to understand object categories (grass is green, sky is blue).
• Rotation prediction (Gidaris et al., 2018): predict which of {0°, 90°, 180°, 270°} a patch was rotated by.
• Jigsaw puzzle (Noroozi & Favaro, 2016): given shuffled 3×3 image patches, predict the permutation index.
• Next-sentence prediction: given two sentences, predict if the second follows the first (used in BERT).

Masked Language Modeling (MLM) — BERT's core pretext task:

• Randomly mask 15% of tokens (replace with [MASK], a random token, or leave unchanged — 80/10/10 split)
• Train the encoder to predict the original token at each masked position
• Loss: cross-entropy over the vocabulary at masked positions only

This forces the model to use bidirectional context — unlike language models, which only look left. The 80/10/10 trick prevents the model from learning [MASK] is a special signal that only appears at test time.

Masked Autoencoders (MAE) — extended to vision (He et al., 2021):

• Mask a very high fraction (75%) of image patches
• Encode only visible patches with ViT
• Decode the masked patches from the encoded representations
• The high masking ratio makes the task non-trivial (no blurring/interpolation), requiring genuine scene understanding

Causal Language Modeling (CLM) — GPT's pretext task:

• Predict the next token given all previous tokens (left-to-right)
• Loss: cross-entropy over vocabulary at each position
• Simple, scales beautifully, directly useful for generation

# BERT-style MLM data preparation (pseudocode, clear and exact)
import random
import torch

MASK_TOKEN = 103  # [MASK] token id in BERT vocabulary

def mask_tokens(input_ids, vocab_size, mask_prob=0.15):
    """
    Returns (masked_input, labels) where labels[i] = original
    token if position i was selected for masking, else -100
    (PyTorch ignores index -100 in cross-entropy).
    """
    labels = input_ids.clone()
    prob_matrix = torch.full(labels.shape, mask_prob)
    masked_positions = torch.bernoulli(prob_matrix).bool()

    labels[~masked_positions] = -100  # only compute loss on masked

    # 80%: replace with [MASK]
    replace_with_mask = torch.bernoulli(torch.full(labels.shape, 0.8)).bool()
    replace_with_mask &= masked_positions
    input_ids[replace_with_mask] = MASK_TOKEN

    # 10%: replace with random token
    replace_with_rand = torch.bernoulli(torch.full(labels.shape, 0.5)).bool()
    replace_with_rand &= masked_positions & ~replace_with_mask
    rand_tokens = torch.randint(vocab_size, labels.shape)
    input_ids[replace_with_rand] = rand_tokens[replace_with_rand]

    # remaining 10%: keep original (already unchanged)
    return input_ids, labels

The transformer (Vaswani et al., 2017) fundamentally changed what 'embedding' means. In Word2Vec, each word has exactly one static vector. In transformers, each token gets a different vector depending on its context — a contextualized embedding.

Token + Positional Encoding: Transformers have no built-in notion of sequence order (unlike RNNs). To inject position information, a positional encoding PE is added to the token embedding before the first layer:

input = TokenEmbedding(token_id) + PositionalEncoding(position)

The original sinusoidal encoding:

PE(pos, 2i)   = sin(pos / 10000^(2i/d))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

Different frequencies for different dimensions let the model attend to relative positions via linear combinations (since sin(a+b) is a linear function of sin(a) and sin(b)).

Modern models often use Rotary Position Embedding (RoPE) (Su et al., 2021) or ALiBi instead. RoPE encodes position by rotating the query/key vectors before the dot-product attention — this generalizes better to longer sequences than absolute additive encodings.

Contextualized representations: After passing through L transformer layers, the vector at position i has been updated by attending to all other positions. The same token 'bank' will have very different final-layer vectors in 'river bank' vs 'savings bank' — the model has disambiguated via context. This is why BERT representations transfer so powerfully: they encode lexical meaning, syntax, and local context all at once.

[CLS] token: BERT prepends a special [CLS] token. Its final-layer embedding is used as a fixed-size representation of the entire sequence for classification tasks — a learned pooling operation.

import torch
import math

class SinusoidalPositionalEncoding(torch.nn.Module):
    """
    Adds sinusoidal positional encodings to token embeddings.
    Matches the original 'Attention is All You Need' formula.
    """
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)   # (T, D)
        pos = torch.arange(max_len).unsqueeze(1).float()  # (T, 1)
        div = torch.exp(
            torch.arange(0, d_model, 2).float() *
            (-math.log(10000.0) / d_model)
        )  # (D/2,)
        pe[:, 0::2] = torch.sin(pos * div)   # even dims
        pe[:, 1::2] = torch.cos(pos * div)   # odd dims
        self.register_buffer('pe', pe.unsqueeze(0))  # (1, T, D)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        return x + self.pe[:, :x.size(1)]

Word-level embeddings are powerful, but many applications need a single vector for an entire sentence or document: semantic search, clustering, duplicate detection, retrieval-augmented generation.

Naive pooling: Average or max-pool the token embeddings from BERT. Surprisingly poor in practice — BERT was trained with MLM, not to produce good sentence vectors via pooling.

Sentence-BERT (SBERT) (Reimers & Gurevych, 2019): Fine-tune BERT with a siamese network on NLI (Natural Language Inference) data. Two sentences are passed through the same BERT; their mean-pooled embeddings are compared. Training objective: contrastive loss or softmax over (entailment, contradiction, neutral). The resulting embeddings have much better sentence-level cosine similarity properties.

Bi-encoder vs. Cross-encoder:

• Cross-encoder: Concatenate both sentences, feed to BERT, output a scalar score. Highly accurate (the model sees both texts together) but O(N) inference at query time — must re-encode every document for each query.
• Bi-encoder: Encode query and document independently, compare with dot product. O(1) per document after offline indexing. Used in all large-scale retrieval systems (dense passage retrieval, DPR). The bi-encoder is always worse than a cross-encoder in isolation, but enables efficient search.

OpenAI text-embedding- and similar*: Large-scale contrastive training on pairs (web search query, clicked document) produces strong general-purpose embeddings. The training data and scale dwarf academic benchmarks — many are competitive with fine-tuned models.

MTEB (Massive Text Embedding Benchmark) evaluates embeddings across 56 tasks including retrieval, clustering, classification, and STS. Use it to compare models objectively.

# Sentence-BERT style: mean pooling of BERT outputs
import torch
from transformers import AutoTokenizer, AutoModel

model_name = 'sentence-transformers/all-MiniLM-L6-v2'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
model.eval()

def mean_pool(model_output, attention_mask):
    """Average token embeddings, ignoring padding tokens."""
    token_embeddings = model_output.last_hidden_state  # (B, T, D)
    # Expand mask to (B, T, D) for broadcasting
    mask = attention_mask.unsqueeze(-1).float()
    sum_embeddings = (token_embeddings * mask).sum(dim=1)
    count = mask.sum(dim=1).clamp(min=1e-9)
    return sum_embeddings / count   # (B, D)

def encode(sentences):
    enc = tokenizer(sentences, padding=True, truncation=True,
                    return_tensors='pt', max_length=128)
    with torch.no_grad():
        out = model(**enc)
    return mean_pool(out, enc['attention_mask'])

vecs = encode(["The cat sat on the mat.", "A kitten rested on a rug."])
print(torch.cosine_similarity(vecs[0:1], vecs[1:2]))

CLIP (Radford et al., OpenAI, 2021) demonstrates one of the most elegant applications of contrastive learning: aligning representations across two completely different modalities — images and text — using nothing but naturally occurring (image, alt-text) pairs from the web.

Architecture: Two independent encoders:

• Image encoder: ViT-B/32 or ResNet-50/101
• Text encoder: 12-layer transformer Both produce a fixed-size embedding (e.g., 512-d). A projection layer maps both into a shared multimodal space.

Training objective: For a batch of N (image, text) pairs, build the N×N similarity matrix. Train to maximize the diagonal (matching pairs) and minimize off-diagonal (mismatches) using InfoNCE symmetrically over both rows and columns. Trained on 400M (image, text) pairs from the internet.

Zero-shot transfer: At inference, to classify an image into K categories:

1. Construct K text prompts: 'a photo of a {class}'
2. Encode all K prompts → text embeddings
3. Encode the image → image embedding
4. Classify as the class with the highest cosine similarity

This achieves 76.2% top-1 accuracy on ImageNet — without seeing any ImageNet training examples, purely by matching image embeddings to text embeddings learned from internet captions.

Why it works: The internet's alt-text is rich with semantic content. By learning to align photos with their captions, CLIP's image encoder builds representations that capture the semantic content useful for description — which turns out to be exactly what object classification needs.

Limitations: CLIP struggles with precise counting, spatial relations ('to the left of'), and abstract reasoning. Its embedding space reflects statistical co-occurrence on web data, including its biases.

# CLIP-style contrastive loss (both directions)
import torch
import torch.nn.functional as F

def clip_loss(image_emb, text_emb, temperature_logit_scale):
    """
    image_emb: (N, D) normalized
    text_emb:  (N, D) normalized
    temperature_logit_scale: learnable log(1/tau), exponentiated
    """
    image_emb = F.normalize(image_emb, dim=-1)
    text_emb  = F.normalize(text_emb,  dim=-1)

    # Scale factor (CLIP learns this; starts at log(1/0.07))
    scale = temperature_logit_scale.exp()

    # N x N similarity matrix
    logits = scale * image_emb @ text_emb.T   # (N, N)

    # Targets: diagonal is the positive pair
    labels = torch.arange(len(image_emb), device=image_emb.device)

    # Symmetric cross-entropy
    loss_i = F.cross_entropy(logits,   labels)  # image->text
    loss_t = F.cross_entropy(logits.T, labels)  # text->image
    return (loss_i + loss_t) / 2

A foundation model (Bommasani et al., Stanford 2021) is a large model trained on broad, general data at scale, which can be adapted (fine-tuned, prompted, or used zero-shot) for a wide range of downstream tasks. GPT-3, BERT, T5, CLIP, LLaMA, and Gemini are all foundation models.

The scaling hypothesis: Kaplan et al. (2020) showed that for autoregressive language models, both loss and downstream performance follow smooth power laws in model size N, dataset size D, and compute C:

L(N) ∝ N^(-0.076)    L(D) ∝ D^(-0.095)

Doubling compute → predictable improvement in loss. This motivated training GPT-3 (175B params) and later models. The Chinchilla paper (Hoffmann et al., 2022) refined this: for a fixed compute budget, the optimal allocation is roughly equal scaling of N and D (not over-training small models).

Emergent abilities: At certain scale thresholds, capabilities appear that are absent in smaller models — few-shot reasoning, arithmetic, chain-of-thought. These are not smooth extrapolations; they are phase transitions. Whether they are truly emergent or just below the measurement threshold at small scale remains debated.

Representation quality vs. model size: Larger pretraining → richer internal representations. GPT-3 embeddings of intermediate layers encode syntactic structure, semantic roles, and world knowledge implicitly. Probing classifiers (train a linear probe on frozen activations) reveal that different layers encode different linguistic phenomena.

The representation bottleneck: The embedding dimension is the information bottleneck between input and output. Increasing embedding dimension allows the model to represent finer distinctions. In BERT-base, d=768; in GPT-3, d=12288. The model's ability to separate near-synonyms, encode subtle context, and generalize to rare words all improve with d.

Instruction tuning and RLHF: Foundation models are pretrained for next-token prediction, not instruction-following. Fine-tuning on instruction-response pairs (InstructGPT, Alpaca) aligns the model to user intent. The embeddings remain; only the final-layer behavior changes. Representation quality from pretraining is what makes post-training work at all.

# Compute-optimal training (Chinchilla): N_opt and D_opt
# From Hoffmann et al. 2022: for a fixed FLOP budget C,
# optimal N and D satisfy: N_opt ≈ D_opt (roughly)
# Total FLOPs ≈ 6 * N * D (approximate rule of thumb)

def chinchilla_optimal(compute_budget_flops):
    """
    Given total FLOPs, return approx optimal model size
    and token count according to Chinchilla scaling laws.
    """
    # C ≈ 6 * N * D and N ≈ D at optimum
    # => C ≈ 6 * N^2  => N = sqrt(C / 6)
    import math
    N_opt = math.sqrt(compute_budget_flops / 6)
    D_opt = compute_budget_flops / (6 * N_opt)
    return N_opt, D_opt

# Example: GPT-3 compute budget ≈ 3.14e23 FLOPs
N, D = chinchilla_optimal(3.14e23)
print(f"Chinchilla-optimal: N={N:.2e} params, D={D:.2e} tokens")
# Note: GPT-3 actual: 175B params, 300B tokens (over-parameterized
# by Chinchilla standards; Chinchilla-optimal would use ~70B params
# trained on ~1.4T tokens for same compute)

Embeddings encode statistical regularities from training data — including human biases. Evaluating and debugging embeddings is as important as training them.

Intrinsic evaluation:

• Word similarity: Correlate embedding cosine similarity with human judgments on datasets like WordSim-353 or SimLex-999. Spearman ρ is the standard metric.
• Analogy benchmarks: Google Analogy Test Set (Mikolov et al.) — 19,544 analogy questions. Solve via: v(king) - v(man) + v(woman) = ? Nearest neighbor should be v(queen).
• Probing tasks: Train a linear classifier on frozen embeddings to predict POS tags, syntactic depth, etc. High accuracy with a linear probe = the information is linearly encoded.

Extrinsic evaluation: Use the embeddings on a real downstream task (NER, sentiment, STS). Ultimately the only evaluation that matters.

Bias in embeddings: Word2Vec trained on news corpora exhibits:

• Gender bias: 'programmer' nearest to 'man'; 'nurse' nearest to 'woman'
• Racial bias in sentiment associations

This is captured by WEAT (Word Embedding Association Test), an analog of the Implicit Association Test. Debiasing techniques (e.g., project out the gender direction from the embedding space) reduce measured bias but can leave residual bias in extrinsic tasks.

Embedding alignment / Procrustes analysis: To compare embeddings across languages (multilingual alignment) or time (diachronic word change), compute the optimal rotation matrix R that maps one embedding space to another:

min_R ‖AᵀR - Bᵀ‖_F    s.t. RᵀR = I

Solution: R = VUᵀ where UΣVᵀ = SVD(AᵀB). Used in multilingual word embeddings (MUSE) to produce cross-lingual transfer without parallel text.

import numpy as np

def analogy(word2vec, pos1, neg1, pos2):
    """
    Solve: pos1 - neg1 + pos2 = ?
    e.g. analogy(w2v, 'king', 'man', 'woman') -> 'queen'
    """
    target = word2vec[pos1] - word2vec[neg1] + word2vec[pos2]
    target /= np.linalg.norm(target)
    # Exclude input words from results
    exclude = {pos1, neg1, pos2}
    best_word, best_sim = None, -1
    for word, vec in word2vec.items():
        if word in exclude:
            continue
        sim = np.dot(target, vec / (np.linalg.norm(vec) + 1e-9))
        if sim > best_sim:
            best_sim = sim
            best_word = word
    return best_word, best_sim

# Procrustes alignment for two embedding spaces A, B (V x D)
def procrustes_align(A, B):
    """Return rotation R such that A @ R ≈ B."""
    U, _, Vt = np.linalg.svd(A.T @ B)
    R = U @ Vt
    return A @ R  # A aligned to B's space