Embeddings

Deep dive

How it works

In: integer ids [B] (or bags of ids with offsets) · Out: vectors [B, d] · State: the table W (V × d) and per-row optimizer state · Cost: O(B·d) gather, O(B·d) scatter-add, O(V·d) memory

1. Forward (gather): out[b] = W[ids[b]], a memory copy of d values per id.
2. Backward (scatter-add): dW[ids[b]] += dout[b]. Duplicate ids in the batch must accumulate, not overwrite.
3. Update only touched rows. Sparse SGD or row-wise Adagrad (one scalar of state per row instead of d). For LLM token tables, dense AdamW is normal because nearly every row of a 50-150k vocabulary is touched by a large batch.
4. Bags (pooling): recommenders sum or average a variable-length list of ids per example (EmbeddingBag) using offsets, without materializing the per-id vectors.
5. Tying: output logits = h · Wᵀ with the same W. The transformer paper scales the input embedding by √d_model when tied.
6. Hashing trick: row = hash(feature) mod V (optionally a second hash for the sign). Collisions act as noise; size V so the collision rate is tolerable.

ids  [7, 2, 7]          W (V x d)                  dW (only rows 2 and 7 non-zero)
 └─ gather ─► out[0]=W[7], out[1]=W[2], out[2]=W[7]
 ◄─ scatter-add: dW[7] = dout[0] + dout[2],  dW[2] = dout[1]

Code

pub struct Embedding { pub v: usize, pub d: usize, pub w: Vec<f32>, adagrad: Vec<f32> }

impl Embedding {
    pub fn new(v: usize, d: usize, seed: u64) -> Self {
        let mut s = seed;
        let w = (0..v * d).map(|_| {                       // small uniform init, std ~ 0.02
            s ^= s << 13; s ^= s >> 7; s ^= s << 17;
            ((s >> 40) as f32 / (1u64 << 24) as f32 - 0.5) * 0.07
        }).collect();
        Embedding { v, d, w, adagrad: vec![0.0; v] }       // row-wise Adagrad: 1 float per row
    }

    pub fn forward(&self, ids: &[usize]) -> Vec<f32> {
        let mut out = Vec::with_capacity(ids.len() * self.d);
        for &i in ids { out.extend_from_slice(&self.w[i * self.d..(i + 1) * self.d]); }
        out
    }

    /// EmbeddingBag(mean): ids for bag j are ids[offsets[j]..offsets[j+1]].
    pub fn bag_mean(&self, ids: &[usize], offsets: &[usize]) -> Vec<f32> {
        let mut out = vec![0.0f32; (offsets.len() - 1) * self.d];
        for j in 0..offsets.len() - 1 {
            let bag = &ids[offsets[j]..offsets[j + 1]];
            let o = &mut out[j * self.d..(j + 1) * self.d];
            for &i in bag { for c in 0..self.d { o[c] += self.w[i * self.d + c]; } }
            if !bag.is_empty() { o.iter_mut().for_each(|x| *x /= bag.len() as f32); }
        }
        out
    }

    /// Sparse backward + row-wise Adagrad step. Duplicates are merged first (sort + run-length),
    /// so each touched row gets exactly one update with its summed gradient.
    pub fn backward_step(&mut self, ids: &[usize], dout: &[f32], lr: f32, eps: f32) {
        let d = self.d;
        let mut order: Vec<usize> = (0..ids.len()).collect();
        order.sort_unstable_by_key(|&b| ids[b]);
        let mut g = vec![0.0f32; d];
        let mut k = 0;
        while k < order.len() {
            let row = ids[order[k]];
            g.iter_mut().for_each(|x| *x = 0.0);
            while k < order.len() && ids[order[k]] == row {   // scatter-add duplicates
                let b = order[k];
                for c in 0..d { g[c] += dout[b * d + c]; }
                k += 1;
            }
            let g2 = g.iter().map(|x| x * x).sum::<f32>() / d as f32;
            self.adagrad[row] += g2;                            // one accumulator per row
            let step = lr / (self.adagrad[row].sqrt() + eps);
            for c in 0..d { self.w[row * d + c] -= step * g[c]; }
        }
    }

    /// Tied output layer: logits = h · Wᵀ (one row of V logits per hidden vector).
    pub fn tied_logits(&self, h: &[f32]) -> Vec<f32> {
        (0..self.v).map(|r| (0..self.d).map(|c| h[c] * self.w[r * self.d + c]).sum()).collect()
    }
}

/// Hashing trick: FNV-1a of the raw feature string into a fixed number of rows.
pub fn hash_row(feature: &str, rows: usize) -> usize {
    let mut h: u64 = 0xcbf29ce484222325;
    for b in feature.bytes() { h ^= b as u64; h = h.wrapping_mul(0x100000001b3); }
    (h % rows as u64) as usize
}

fn main() {
    let mut emb = Embedding::new(1000, 8, 7);
    let ids = [7usize, 2, 7];
    let out = emb.forward(&ids);
    assert_eq!(&out[0..8], &out[16..24]);                  // same id, same vector
    let before = emb.w.clone();
    let dout = vec![1.0f32; ids.len() * 8];
    emb.backward_step(&ids, &dout, 0.1, 1e-8);
    let changed: Vec<usize> = (0..1000).filter(|&r| emb.w[r * 8..(r + 1) * 8] != before[r * 8..(r + 1) * 8]).collect();
    println!("rows updated: {changed:?}");                 // only [2, 7]
    let bag = emb.bag_mean(&[1, 2, 3, 9], &[0, 3, 4]);
    println!("bag 0 mean[0] = {:.4}, logits len {}", bag[0], emb.tied_logits(&out[0..8]).len());
    println!("user:42 -> row {}", hash_row("user:42", 1 << 20));
}

Complexity and constants

• Forward: one random row read per id. A 4096-dim bf16 row is 8 KiB, so lookups stream at near memory bandwidth. Small rows (d = 16-64 in recommenders, 32-256 B) are latency-bound: one or two cache lines per random access.
• Backward: sort or atomic scatter-add. On GPUs, frequent ids (the token "the", a popular item) cause atomic contention. Sort-then-reduce avoids it.
• Memory for a table with Adam: V·d·(2 + 4 + 4 + 4) bytes with bf16 weights and fp32 master plus m and v. Row-wise Adagrad: V·(d·4 + 4), which is why recommenders use it.
• Matryoshka embeddings: vectors trained so prefixes of 1/2, 1/4, … of the dimensions remain usable.

Efficiency

• Cost: a lookup reads one row: 8 KiB for 4096 bf16 dimensions (bandwidth-bound), 256 B for d = 64 in fp32 (latency-bound, a few cache lines). Training with bf16 weights, fp32 master and Adam state costs V·d·14 bytes (this page); a 10⁹-row, d = 64 fp32 table is 256 GB (arithmetic), which is why recommender tables are sharded. The bound is memory capacity and bandwidth, not FLOPs.
• Levers:
  • Row-wise Adagrad — optimizer state of 4 bytes per row instead of 8 bytes per parameter (this page) — huge sparse tables.
  • int8/int4 tables — 4x/8x smaller than fp32 (arithmetic; FBGEMM supports both) — inference.
  • Hashing trick or compositional embeddings — table size independent of vocabulary (Weinberger 2009; Shi 2020) — very large id spaces.
  • Matryoshka training — up to 14x smaller embeddings at equal ImageNet-1K accuracy (Kusupati 2022) — retrieval and storage.
  • Sort-then-reduce backward instead of atomics — avoids contention on frequent ids (this page) — GPU training.
• Naive → standard → state of the art:
  • One-hot times matrix · O(V·d) per token
  • Row gather, sparse gradient · O(d) per id
  • Sharded, quantized tables with row-wise Adagrad (TorchRec/FBGEMM), Matryoshka prefixes · rows 4-8x smaller, only touched rows updated

Implementation notes

• Pad V to a multiple of 64 or 128 for tensor-core friendly output matmuls (nanoGPT pads GPT-2's 50257 to 50304).
• Initialize with small std (0.02 or 1/√d) when tied to the output. PyTorch's nn.Embedding default is N(0, 1), which makes initial logits huge when tied.
• nn.Embedding(sparse=True) produces sparse gradients that work with SGD, SparseAdam and Adagrad, not with AdamW.
• Huge tables: shard rows across devices (row-wise or table-wise model parallelism) and exchange looked-up vectors with all-to-all, as in DLRM. Cache hot rows on GPU and keep cold rows in host memory.
• Serving: quantize tables row-wise to int8 or int4 with a per-row scale and bias. Memory drops 4x (int8) or 8x (int4) versus fp32; recommenders report small accuracy loss at int8.
• Weight decay on embeddings varies by recipe (nanoGPT decays them, many do not). With decoupled decay on a sparse table, rarely-seen rows shrink toward zero between their updates.

Pitfalls

• Index out of range from a tokenizer/table vocabulary mismatch: on GPU this is a device-side assert or silent garbage.
• Overwriting instead of accumulating gradients for repeated ids.
• Dense Adam on a sparse table: rows not in the batch still move because m keeps a non-zero momentum (and decay shrinks them).
• Tying embeddings and then scaling the input by √d without matching the init: activations start far too large.
• Hash tables sized too small for the id space: heavy collisions merge popular and rare ids.

Variants and alternatives

In production

PyTorch nn.Embedding / nn.EmbeddingBag, TorchRec and FBGEMM (sharded tables, row-wise Adagrad, int8/int4 tables), TensorFlow Recommenders, every LLM token embedding. GPT-2 ties input and output embeddings; Llama 2/3 7B-70B do not.

Implementations

Sources

• Mikolov et al., Distributed Representations of Words and Phrases and their Compositionality, NeurIPS 2013. paper
• Press and Wolf, Using the Output Embedding to Improve Language Models, EACL 2017. paper
• Weinberger, Dasgupta, Langford, Smola, Attenberg, Feature Hashing for Large Scale Multitask Learning, ICML 2009. paper
• Naumov et al., Deep Learning Recommendation Model for Personalization and Recommendation Systems, 2019. paper · code
• Shi et al., Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems, KDD 2020. paper
• Kusupati et al., Matryoshka Representation Learning, NeurIPS 2022. paper · code
• Duchi, Hazan, Singer, Adaptive Subgradient Methods for Online Learning and Stochastic Optimization, JMLR 2011. paper