Negative sampling

Deep dive

How it works

In: positive pairs (x, y⁺), a noise distribution q over candidates, k · Out: a sampled loss and gradient touching k+1 candidate rows · State: the alias table (two arrays of length V) and the embedding tables · Cost: O(1) per negative draw, O(k·d) per positive instead of O(V·d)

1. Pick q. Word2vec: count^0.75 / Z. Catalogues: uniform, or popularity^α. In-batch: the batch's own item frequency.
2. Build the sampler. Alias table from q, O(V) once (Alias method sampling).
3. Draw. For each positive take k indices from q; mask any equal to y⁺ (accidental hit).
4. SGNS loss. -log σ(u·v⁺) - Σᵢ log σ(-u·vᵢ⁻); gradient σ(u·v) - label per row, only k+1 rows of v and one of u change.
5. Softmax variant. Logits over {y⁺} ∪ negatives, each corrected by - ln q(y), then cross-entropy: sampled softmax.
6. In-batch. For a batch of B pairs the other B-1 positives are the negatives; q(y) is the probability that y is in a batch (estimated from a streaming counter).
7. Harder negatives. Mix easy draws with ANN-mined near-misses from the current model; mask known positives.
8. Evaluate by full ranking, not the sampled loss.

Estimators. NCE (Gutmann & Hyvärinen 2010): logistic regression of data vs noise on the model's log-density; consistent for unnormalised models, needs q(y) as a number. SGNS (Mikolov et al. 2013): the simplified variant that uses only samples, not their probabilities. Sampled softmax (Jean et al. 2015, importance sampling): approximates the full softmax with the ln q correction; training only, and TensorFlow notes it is generally an underestimate of the full loss.

Where negatives come from

• Mixed. In-batch plus uniform catalogue draws (Yang et al. 2020) cover rare items the batch lacks; each source needs its own correction.
• Hard. DPR adds BM25 near-misses to in-batch negatives; ANCE draws them from an ANN index of the current encoder, refreshed during training (Xiong et al. 2020). A mined negative that is truly relevant gets penalised; filter by a cross-encoder score or rank cut-off.
• Implicit feedback. Unobserved means unknown, not disliked: a draw is a weak label that may be a false negative, so treat negatives as noise, not as dislikes.
• Knowledge graphs. Corrupt a true triple (h, r, t) by replacing h or t with a random entity; RotatE weights the corruptions by their current score (self-adversarial; Sun et al. 2019). Filter corruptions that are true triples.
• Contrastive learning. Same loss, augmented views as positives, the batch as negatives (Contrastive learning).

Code

Alias sampler from count^0.75, SGNS on a two-cluster corpus (tokens 0-9 and 10-19 co-occur only within their cluster), the logQ term. Asserts: 10⁶ draws match the target within 0.003; within-cluster dot products beat cross-cluster; popular vs rare logQ.

use std::time::Instant;
const D: usize = 16; // embedding dim

struct Rng(u64);
impl Rng {
    fn next(&mut self) -> u64 { // xorshift64*
        self.0 ^= self.0 >> 12; self.0 ^= self.0 << 25; self.0 ^= self.0 >> 27;
        self.0.wrapping_mul(0x2545F4914F6CDD1D)
    }
}

/// Vose alias table: O(n) build, O(1) per draw (one bucket pick + one coin).
struct Alias { prob: Vec<f64>, alias: Vec<u32> }
impl Alias {
    fn new(w: &[f64]) -> Self {
        let n = w.len();
        let tot: f64 = w.iter().sum();
        let mut p: Vec<f64> = w.iter().map(|x| x * n as f64 / tot).collect();
        let (mut small, mut large) = (vec![], vec![]);
        for (i, &x) in p.iter().enumerate() { if x < 1.0 { small.push(i) } else { large.push(i) } }
        let mut alias = vec![0u32; n];
        while let (Some(s), Some(l)) = (small.pop(), large.pop()) {
            alias[s] = l as u32;
            p[l] -= 1.0 - p[s];
            if p[l] < 1.0 { small.push(l) } else { large.push(l) }
        }
        for i in small.into_iter().chain(large) { p[i] = 1.0 } // leftovers: rounding
        Alias { prob: p, alias }
    }
    fn draw(&self, r: u64) -> usize {
        let i = (((r >> 32) * self.prob.len() as u64) >> 32) as usize; // bucket
        let u = (r & 0xFFFF_FFFF) as f64 / 4294967296.0; // coin
        if u < self.prob[i] { i } else { self.alias[i] as usize }
    }
}

fn sigmoid(x: f32) -> f32 { 1.0 / (1.0 + (-x).exp()) }
fn dot(a: &[f32; D], b: &[f32; D]) -> f32 { a.iter().zip(b).map(|(x, y)| x * y).sum() }

/// One SGNS step: loss = -log s(u.v) - sum_neg log s(-u.vn).
fn sgns_step(u: &mut [f32; D], v: &mut [[f32; D]], pos: usize, negs: &[usize], lr: f32) {
    let mut gu = [0f32; D];
    for (j, label) in std::iter::once((pos, 1.0f32)).chain(negs.iter().map(|&n| (n, 0.0))) {
        let g = sigmoid(dot(u, &v[j])) - label; // d loss / d logit
        for d in 0..D { gu[d] += g * v[j][d]; v[j][d] -= lr * g * u[d]; }
    }
    for d in 0..D { u[d] -= lr * gu[d]; }
}

/// Sampling-bias correction for in-batch softmax logits (Yi et al. 2019): s - ln q.
fn logq(s: f32, q: f32) -> f32 { s - q.ln() }

fn main() {
    let mut rng = Rng(0x9E3779B97F4A7C15);
    // 1. Alias sampler vs target counts^0.75
    let counts: Vec<f64> = (0..20).map(|i| 1000.0 / (1.0 + (i % 10) as f64)).collect(); // Zipf-ish, 2 clusters
    let target: Vec<f64> = counts.iter().map(|c| c.powf(0.75)).collect();
    let z: f64 = target.iter().sum();
    let al = Alias::new(&target);
    let n = 1_000_000;
    let mut hist = vec![0u32; 20];
    let t = Instant::now();
    for _ in 0..n { hist[al.draw(rng.next())] += 1 }
    let dt = t.elapsed().as_secs_f64();
    for i in 0..20 { assert!((hist[i] as f64 / n as f64 - target[i] / z).abs() < 0.003) }
    println!("alias: {:.0} M draws/s (1e6 draws, incl. loop)", n as f64 / dt / 1e6);


    // 2. SGNS on 2-cluster corpus: (center, context) pairs from the same cluster
    let (v_n, k, lr) = (20, 5, 0.05f32);
    let mut u = vec![[0f32; D]; v_n];
    for r in u.iter_mut() { for x in r.iter_mut() { *x = (rng.next() >> 40) as f32 / 1.6e8 - 0.05 } }
    let mut v = vec![[0f32; D]; v_n];
    let uni = Alias::new(&counts); // data distribution within the corpus
    let mut negs = [0usize; 5];
    let steps = 2_000_000usize;
    let t = Instant::now();
    for _ in 0..steps {
        let c = uni.draw(rng.next());
        let ctx = (uni.draw(rng.next()) % 10) + (c / 10) * 10; // same cluster context
        for x in negs.iter_mut() { *x = al.draw(rng.next()) }
        sgns_step(&mut u[c], &mut v, ctx, &negs[..k], lr);
    }
    let dt = t.elapsed().as_secs_f64();
    println!("sgns: {:.2} M updates/s (dim {D}, k={k}, 1 thread)", steps as f64 / dt / 1e6);
    let (mut win, mut cross, mut nw, mut nc) = (0.0, 0.0, 0, 0);
    for a in 0..v_n { for b in 0..v_n {
        if a / 10 == b / 10 { win += dot(&u[a], &v[b]); nw += 1 } else { cross += dot(&u[a], &v[b]); nc += 1 }
    } }
    let (win, cross) = (win / nw as f32, cross / nc as f32);
    println!("mean u.v within {win:.2} vs cross {cross:.2}");
    assert!(win > cross + 1.0);

    // 3. logQ: equal raw logits, popular q=0.2 vs rare q=0.001
    let (pop, rare) = (logq(2.0, 0.2), logq(2.0, 0.001));
    assert!((pop - (2.0 + 5.0f32.ln())).abs() < 1e-5);
    assert!(((rare - pop) - 200f32.ln()).abs() < 1e-4);
    println!("logq popular {pop:.3} rare {rare:.3}");
}

Measured (Apple M4 Pro, one thread, release build, three runs): about 215 M alias draws/s with the RNG on a 20-entry table; 165-212 M/s on a 10⁶-entry table (12 MB, cache-resident; a variant of this code); about 15 M SGNS steps/s at d = 16, k = 5 on 20 tokens (all in L1: an upper bound).

Complexity and constants

Per positive: (k+1)·d multiply-adds and k+1 random row reads, versus V·d for the full softmax: about V/(k+1) less arithmetic; row gathers decide in practice. Alias table: 12 bytes/entry (f64 + u32). A draw is one RNG step and two reads in one bucket; a CDF binary search needs log₂ V dependent reads.

Efficiency

• Cost: on an M4 Pro, one alias draw is about 5 ns (measured, above) and one SGNS step with k = 5, d = 16 about 65 ns when the tables are in L1. At V = 10⁶, d = 300, rows are 1.2 KB: the step is bound by k+1 cache-missing row reads and writes, not arithmetic.
• Levers:
  • Alias table instead of a CDF search — O(1) vs O(log V) per draw (Walker 1977; Vose 1991) — whenever V is large and q is fixed.
  • Share the k negatives across a mini-batch — one (B×d)·(d×k) matmul replaces B·k scattered reads (design choice, not benchmarked) — GPU or BLAS training.
  • In-batch negatives — zero extra encoder passes, B-1 negatives per positive — needs the logQ correction (Yi et al. 2019).
  • Refresh mined hard negatives periodically, as ANCE does — avoids one ANN query per example — slowly changing encoders.
  • Sub-sample frequent words — 2x-10x speed-up (Mikolov et al. 2013) — skewed co-occurrence data.
• Naive → standard → state of the art:
  • Full softmax over V · O(V·d) per example.
  • Unigram^0.75 negatives from an alias table, k = 5 · O(k·d).
  • In-batch + logQ + uniform + mined hard negatives · no extra encoder pass per in-batch negative.

Implementation notes

• Draw with replacement; masking accidental hits is cheap when V ≫ k (TensorFlow's sampled_softmax_loss defaults remove_accidental_hits=True).
• The logQ term must use the sampler's own normalised q (e.g. the count^0.75/Z value), not raw counts.
• For in-batch negatives estimate q with a streaming counter if popularity drifts; clamp q away from 0 (-ln q blows up).
• One RNG state per worker; a shared one is a contended cache line.

Pitfalls

• In-batch without logQ: popular items are over-penalised, the retriever favours the long tail.
• Hard negatives mined too early or too hard: false negatives dominate; start easy, then anneal.
• A duplicate item in a batch is a negative for itself; mask by id.
• Reading the sampled loss as a likelihood or calibrated score.

Variants and alternatives

In production

• word2vec, gensim (negative = 5, ns_exponent = 0.75 by default): word embeddings.
• In-batch plus logQ in retrieval and recommendation (Two-tower retrieval, Deep recommendation models); DPR and ANCE for passages.
• PyTorch-BigGraph on social and knowledge graphs: num_batch_negs and num_uniform_negs, both 50 by default.

Implementations

Sources

• Mikolov, Sutskever, Chen, Corrado, Dean, Distributed Representations of Words and Phrases and their Compositionality, NeurIPS 2013. paper
• Gutmann, Hyvärinen, Noise-contrastive estimation: A new estimation principle for unnormalized statistical models, AISTATS 2010. paper
• Jean, Cho, Memisevic, Bengio, On Using Very Large Target Vocabulary for Neural Machine Translation, ACL 2015. paper
• Yi et al., Sampling-bias-corrected neural modeling for large corpus item recommendations, RecSys 2019. paper
• Yang, Yi, Cheng, Hong, Zhang, Chi, Mixed Negative Sampling for Learning Two-tower Neural Networks in Recommendations, WWW 2020 Companion. paper
• Karpukhin et al., Dense Passage Retrieval for Open-Domain Question Answering, EMNLP 2020. paper · code
• Xiong et al., Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval, ICLR 2021. paper
• Sun, Deng, Nie, Tang, RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space, ICLR 2019. paper
• Lerer et al., PyTorch-BigGraph: A Large-scale Graph Embedding System, SysML 2019. paper · code
• Walker, An Efficient Method for Generating Discrete Random Variables with General Distributions, ACM TOMS 1977. paper
• Vose, A linear algorithm for generating random numbers with a given distribution, IEEE TSE 1991. paper
• Docs: gensim Word2Vec · tf.nn.sampled_softmax_loss