Two-tower retrieval

Deep dive

How it works

Parts:

• user/query tower — features to a d-vector (MLP; text encoder in DPR)
• item tower — id and content features to a d-vector, run offline
• in-batch softmax — the batch's other items are the negatives
• ANN index — precomputed item vectors, searched by inner product
• ranker — heavier model on the retrieved few hundred

1. Encode. u = f(user), v = g(item), d of order 10², usually L2-normalised so the dot product is a cosine.
2. Sampled softmax. For B positive pairs: s_ij = u_i·v_j / τ - log q_j, loss = -mean_i log softmax_j(s_ij)[i]. The other B-1 items are free negatives.
3. logQ correction. Item j is in the batch with probability q_j (its popularity), so popular items are over-sampled as negatives and pushed down. Subtracting log q_j from column j removes the bias (Yi et al. 2019); estimate q from counts or a streaming sketch.
4. Temperature τ. Small τ (order of 0.05-0.2) sharpens the softmax; tune it. Mask duplicates of the positive in the batch.
5. Harder negatives. In-batch negatives are mostly easy. Add uniform catalogue samples (Yang et al. 2020) or mined near-misses such as BM25 top results (DPR); see Negative sampling.
6. Serve. Embed all items offline, index them (HNSW, IVF-PQ, ScaNN). Per request: user tower, top-N by inner product, then the ranker.

user feats ─► [user tower] ─► u ┐
                                ├─ u·v / τ − log q ─ softmax over batch, labels on the diagonal
item id/feats ─► [item tower] ─► v ┘
serving:  items ─► v (offline) ─► ANN index;   request ─► u ─► top-N ─► ranker

Code

Linear user tower, item embedding table, cosine with τ = 0.1, in-batch softmax with and without logQ, SGD, brute-force top-10 on held-out clicks. Synthetic: 8 taste groups, Zipf-popular items per group, 20 % off-taste clicks.

const G: usize = 8; // user groups = feature dim
const PER: usize = 50; // items per group
const N: usize = G*PER; // catalogue size
const D: usize = 16; // embedding dim
const B: usize = 128; // in-batch negatives = B-1
const TAU: f32 = 0.1; // softmax temperature

struct Rng(u64); // xorshift64*
impl Rng {
    fn next(&mut self) -> u64 {
        let mut x = self.0;
        x ^= x << 13; x ^= x >> 7; x ^= x << 17;
        self.0 = x;
        x.wrapping_mul(0x2545F4914F6CDD1D)
    }
    fn f(&mut self) -> f32 { (self.next() >> 40) as f32/(1u64 << 24) as f32 }
    fn below(&mut self, n: usize) -> usize { (self.next() % n as u64) as usize }
}

/// One click: noisy one-hot user features, item id.
fn sample(rng: &mut Rng, zipf_total: f32) -> ([f32; G], usize) {
    let g = rng.below(G);
    let mut x = [0.0f32; G];
    for v in x.iter_mut() { *v = 0.3*(rng.f() - 0.5); }
    x[g] += 1.0;
    let g2 = if rng.f() < 0.8 { g } else { rng.below(G) }; // 20 % off-taste clicks
    let mut t = rng.f()*zipf_total; // Zipf(1) rank inside the group
    let mut r = 0;
    while r + 1 < PER && t > 1.0/(r + 1) as f32 { t -= 1.0/(r + 1) as f32; r += 1; }
    (x, r*G + g2)
}
fn unit(z: &[f32]) -> (Vec<f32>, f32) {
    let n = z.iter().map(|v| v*v).sum::<f32>().sqrt() + 1e-9;
    (z.iter().map(|v| v/n).collect(), n)
}
fn unit_back(de: &[f32], e: &[f32], n: f32) -> Vec<f32> {
    let dot: f32 = de.iter().zip(e).map(|(a, b)| a*b).sum();
    de.iter().zip(e).map(|(a, b)| (a - dot*b)/n).collect()
}
struct Towers { wu: Vec<f32>, items: Vec<f32> } // user tower: linear G->D; item tower: embedding table N x D

impl Towers {
    fn user(&self, x: &[f32; G]) -> (Vec<f32>, Vec<f32>, f32) {
        let z: Vec<f32> = (0..D).map(|d| (0..G).map(|g| self.wu[d*G + g]*x[g]).sum()).collect();
        let (e, n) = unit(&z);
        (z, e, n)
    }
    fn item(&self, i: usize) -> (Vec<f32>, f32) { unit(&self.items[i*D..(i+1)*D]) }

    fn step(&mut self, batch: &[([f32; G], usize)], logq: &[f32], lr: f32, correct: bool) {
        let us: Vec<_> = batch.iter().map(|(x, _)| self.user(x)).collect();
        let vs: Vec<_> = batch.iter().map(|&(_, i)| self.item(i)).collect();
        let (mut de_u, mut de_v) = (vec![vec![0.0f32; D]; B], vec![vec![0.0f32; D]; B]);
        for i in 0..B {
            // logits over the batch, minus log q_j; duplicates of the positive masked out
            let lg: Vec<f32> = (0..B).map(|j| {
                if j != i && batch[j].1 == batch[i].1 { return f32::NEG_INFINITY; }
                let s: f32 = (0..D).map(|d| us[i].1[d]*vs[j].0[d]).sum::<f32>()/TAU;
                if correct { s - logq[batch[j].1] } else { s }
            }).collect();
            let m = lg.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            let ex: Vec<f32> = lg.iter().map(|l| (l - m).exp()).collect();
            let z: f32 = ex.iter().sum();
            for j in 0..B {
                let g = (ex[j]/z - if i == j { 1.0 } else { 0.0 })/(TAU*B as f32); // dLoss/ds_ij
                for d in 0..D { de_u[i][d] += g*vs[j].0[d]; de_v[j][d] += g*us[i].1[d]; }
            }
        }
        for i in 0..B {
            let dz = unit_back(&de_u[i], &us[i].1, us[i].2);
            for d in 0..D { for g in 0..G { self.wu[d*G + g] -= lr*dz[d]*batch[i].0[g]; } }
            let dv = unit_back(&de_v[i], &vs[i].0, vs[i].1);
            let row = batch[i].1*D; // duplicates in the batch simply add up
            for d in 0..D { self.items[row + d] -= lr*dv[d]; }
        }
    }
    fn top_k(&self, x: &[f32; G], k: usize, item_emb: &[Vec<f32>]) -> Vec<usize> {
        let u = self.user(x).1;
        let mut s: Vec<(f32, usize)> = (0..N).map(|i| (u.iter().zip(&item_emb[i]).map(|(a, b)| a*b).sum(), i)).collect();
        s.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
        s[..k].iter().map(|p| p.1).collect()
    }
}

fn run(correct: bool, train: &[([f32; G], usize)], test: &[([f32; G], usize)], k: usize) -> f32 {
    let mut rng = Rng(42);
    let mut cnt = vec![1.0f32; N]; // +1 smoothing
    for p in train { cnt[p.1] += 1.0; }
    let tot: f32 = cnt.iter().sum();
    let logq: Vec<f32> = cnt.iter().map(|c| (c/tot).ln()).collect();
    let mut t = Towers { wu: (0..D*G).map(|_| rng.f() - 0.5).collect(), items: (0..N*D).map(|_| rng.f() - 0.5).collect() };
    for s in 0..4000 {
        let batch: Vec<_> = (0..B).map(|_| train[rng.below(train.len())]).collect();
        t.step(&batch, &logq, if s < 3000 { 0.3 } else { 0.1 }, correct);
    }
    let emb: Vec<Vec<f32>> = (0..N).map(|i| t.item(i).0).collect();
    let hits = test.iter().filter(|(x, y)| t.top_k(x, k, &emb).contains(y)).count();
    hits as f32/test.len() as f32
}

fn main() {
    let zt: f32 = (1..=PER).map(|r| 1.0/r as f32).sum();
    let mut rng = Rng(7);
    let train: Vec<_> = (0..20000).map(|_| sample(&mut rng, zt)).collect();
    let test: Vec<_> = (0..2000).map(|_| sample(&mut rng, zt)).collect(); // held-out interactions
    let k = 10;
    let mut cnt = vec![0usize; N];
    for p in &train { cnt[p.1] += 1; }
    let mut order: Vec<usize> = (0..N).collect();
    order.sort_by_key(|&i| std::cmp::Reverse(cnt[i]));
    let pop = test.iter().filter(|p| order[..k].contains(&p.1)).count() as f32/test.len() as f32;
    let plain = run(false, &train, &test, k);
    let corrected = run(true, &train, &test, k);
    println!("recall@{k}: popularity {pop:.3}, in-batch softmax {plain:.3}, + logQ {corrected:.3}");
    assert!(corrected > pop + 0.1, "two-tower must beat popularity");
}

Measured on an Apple M4 Pro (release, about 5 s): recall@10: popularity 0.240, in-batch softmax 0.205, + logQ 0.534; the ceiling of this data is about 0.53. Without logQ the tower avoids popular items, here the right answers; the gap is large because the catalogue is tiny and skewed.

Complexity and constants

• Training: B²·d MACs for logits plus 2B tower passes; the B² logits matrix (B = 4096: 64 MB in f32) is the memory cost.
• Serving: exact search is N·d MACs reading N·d·4 bytes; 10⁸ items × 64 × 4 B = 25.6 GB before quantization.
• Cost drivers: item bytes read per query, index build time, re-embedding rate.

Efficiency

• Cost: One request = user tower forward (small MLP) + one ANN query. Exact scan of 1M items × 64 f32 (256 MB) took about 6.5-7 ms on one M4 Pro core, about 38 GB/s: memory-bandwidth bound (measured here, simple loop with top-10 insertion). A toy training step: about 0.6 ms (measured here).
• Levers:
  • ANN index instead of a scan — visits a small fraction of vectors (order of 10²-10³× fewer at 10⁸ items, recall tradeoff) — N above ~10⁶.
  • int8 or product quantization — 4× (int8), order of 16-64× (PQ), lossy (IVF-PQ) — index must fit in RAM.
  • Larger batch with logQ — more negatives per item-tower pass (Yi et al. 2019) — memory allows B².
• Naive → standard → state of the art:
  • Ranker on every item · N ranker passes per request
  • Two towers + exact scan · one tower pass + N·d MACs, about 7 ms per 10⁶ items per core (measured)
  • Two towers + quantized ANN index · sublinear search, compressed memory

Implementation notes

• Normalise once at index time; store f16/int8.
• Item and user towers must come from the same model version; a mismatch silently destroys recall.

Pitfalls

• No logQ: popular items suppressed.
• logQ assumes batch items are drawn by q; extra uniform negatives need their own correction.
• False negatives (relevant but unclicked items) push good items down.
• Train/serve skew in user features (Feature store).
• Id-only item tower: new items have an untrained row; add content features.

Variants and alternatives

In production

Video recommendation (Covington, Yi), dense passage retrieval for question answering (DPR), product and ad candidate generation; FAISS, ScaNN, pgvector as serving layers.

Implementations

Sources

• Covington, Adams, Sargin, Deep Neural Networks for YouTube Recommendations, RecSys 2016. paper
• Yi, Yang, Hong, Cheng, Hong, Heldt, Kumthekar, Zhao, Wei, Chi, 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
• Guo et al., Accelerating Large-Scale Inference with Anisotropic Vector Quantization, ICML 2020. paper · code