Matrix factorization for collaborative filtering

Deep dive

How it works

Parts:

• interaction triples — (row, col, value), observed cells only, indexed by row and by column
• factors P (rows × k), Q (cols × k) — small random init, never zero
• mean and biases — μ + b_u + b_i absorb rating scale and popularity
• confidence weights — implicit only: c = 1 + α·count, target 1 if seen
• solver — SGD or ALS on squared error + λ(‖p‖² + ‖q‖²)
• scorer — dot product, seen items removed, top-k

1. Model. r̂ = μ + b_u + b_i + p_u·q_i; loss Σ_observed (r − r̂)² + λ(‖p_u‖² + ‖q_i‖² + b_u² + b_i²).
2. SGD (Funk-style). Per cell: e = r − r̂, p += γ(e·q − λp), q += γ(e·p − λq) (use the old p). Shuffle each epoch.
3. ALS. Fix Q; each p_u is a ridge regression: p_u = (Σ qqᵀ + λ·n_u·I)⁻¹ Σ r·q, a k×k solve. Then fix P, re-solve every q_i. Each half-step cannot raise the loss.
4. Implicit feedback. Every cell is target 1 (seen) or 0 with confidence c = 1 + α·count, unseen c = 1. With QᵀC_uQ = QᵀQ + Qᵀ(C_u − I)Q, compute QᵀQ once, add only observed items: O(nnz·k² + n·k³), not U×I.
5. Cold start. New row: fix Q, solve one ridge system on its first interactions (fold-in, O(n·k² + k³)). No interactions: popularity or a bandit. New columns need content features (Embeddings).
6. Retrieve. Score p_u·q_i; exact top-k scans N·k MACs, larger N goes to HNSW or another ANN index (inner product is no metric: append a coordinate or L2-normalise). Drop seen items, pass a few hundred to a ranker.

 R (sparse)        ≈    P (rows×k)  ·  Qᵀ (k×cols)
 . 5 . . 1 .            [ p_u ]        [ q_i ... ]
 3 . . 4 . .
 ALS: fix Q ─► solve each p_u ─► fix P ─► solve each q_i ─► repeat
 serve: p_u ─► dot all q_i (or ANN) ─► drop seen ─► top-k ─► ranker

Code

Hidden rank-4 taste, 300 × 200 cells. Explicit: ALS and Funk-style SGD, held-out RMSE (mean-centred, no biases). Implicit: events where the hidden score is high, ALS with the Gram trick, hit@10 versus popularity.

use std::time::Instant;
const K: usize = 4; // latent dimension
struct Rng(u64);
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 }
    fn f(&mut self) -> f32 { (self.next() >> 40) as f32 / (1u64 << 24) as f32 } // [0,1)
}
// Solve (A + lam*I) x = b by Cholesky (A symmetric positive semi-definite, K x K).
fn solve(a: &[f32; K * K], b: &[f32; K], lam: f32) -> [f32; K] {
    let mut l = [0.0f32; K * K];
    for i in 0..K {
        for j in 0..=i {
            let mut s = a[i * K + j] + if i == j { lam } else { 0.0 };
            for p in 0..j { s -= l[i * K + p] * l[j * K + p]; }
            l[i * K + j] = if i == j { s.sqrt() } else { s / l[j * K + j] };
        }
    }
    let (mut y, mut x) = ([0.0f32; K], [0.0f32; K]);
    for i in 0..K { y[i] = (b[i] - (0..i).map(|p| l[i * K + p] * y[p]).sum::<f32>()) / l[i * K + i]; }
    for i in (0..K).rev() { x[i] = (y[i] - (i + 1..K).map(|p| l[p * K + i] * x[p]).sum::<f32>()) / l[i * K + i]; }
    x
}
type Fac = Vec<[f32; K]>;
fn dot(a: &[f32; K], b: &[f32; K]) -> f32 { (0..K).map(|k| a[k] * b[k]).sum() }
// ALS half-step: re-solve every row of x against fixed y. rows[r] = (col, rating, weight).
// Explicit: weight 1, lam scaled by count. Implicit (Hu et al.): gram = YtY, weight c-1, target 1.
fn half_step(x: &mut Fac, y: &Fac, rows: &[Vec<(usize, f32, f32)>], lam: f32, gram: Option<&[f32; K * K]>) {
    for (r, obs) in rows.iter().enumerate() {
        let mut a = gram.map_or([0.0; K * K], |g| *g);
        let mut b = [0.0f32; K];
        for &(c, rating, w) in obs {
            for i in 0..K {
                b[i] += (if gram.is_some() { w + 1.0 } else { rating }) * y[c][i];
                for j in 0..K { a[i * K + j] += w * y[c][i] * y[c][j]; }
            }
        }
        let n = if gram.is_some() { 1.0 } else { obs.len().max(1) as f32 };
        x[r] = solve(&a, &b, lam * n);
    }
}
fn gram(y: &Fac) -> [f32; K * K] {
    let mut g = [0.0; K * K];
    for v in y { for i in 0..K { for j in 0..K { g[i * K + j] += v[i] * v[j]; } } }
    g
}
fn mk(rng: &mut Rng, n: usize) -> Fac { (0..n).map(|_| [(); K].map(|_| rng.f() - 0.5)).collect() }
fn rmse(u: &Fac, v: &Fac, mu: f32, set: &[(usize, usize, f32)]) -> f32 {
    (set.iter().map(|&(a, b, r)| (mu + dot(&u[a], &v[b]) - r).powi(2)).sum::<f32>() / set.len() as f32).sqrt()
}
fn main() {
    let (nu, ni) = (300, 200);
    let mut rng = Rng(7);
    let (tu, ti) = (mk(&mut rng, nu), mk(&mut rng, ni));
    // 12 % of cells observed; rating = 3 + 2*(p.q) + noise (std about 0.03), split 80/20
    let (mut train, mut test) = (vec![], vec![]);
    for u in 0..nu { for i in 0..ni {
        if rng.f() < 0.12 {
            let r = 3.0 + 4.0 * dot(&tu[u], &ti[i]) + 0.1 * (rng.f() - 0.5);
            if rng.f() < 0.8 { train.push((u, i, r)) } else { test.push((u, i, r)) }
        }
    } }
    let mu = train.iter().map(|t| t.2).sum::<f32>() / train.len() as f32;
    let base = (test.iter().map(|t| (mu - t.2).powi(2)).sum::<f32>() / test.len() as f32).sqrt();
    let (mut urows, mut irows) = (vec![vec![]; nu], vec![vec![]; ni]);
    for &(a, b, r) in &train { urows[a].push((b, r - mu, 1.0f32)); irows[b].push((a, r - mu, 1.0f32)); }
    // --- explicit ALS ---
    let (mut u, mut v) = (mk(&mut rng, nu), mk(&mut rng, ni));
    let t0 = Instant::now();
    for _ in 0..15 { half_step(&mut u, &v, &urows, 0.005, None); half_step(&mut v, &u, &irows, 0.005, None); }
    let (als_ms, als) = (t0.elapsed().as_secs_f64() * 1e3, rmse(&u, &v, mu, &test));
    // --- Funk-style SGD on the same data ---
    let (mut su, mut sv) = (mk(&mut rng, nu), mk(&mut rng, ni));
    let t0 = Instant::now();
    for _ in 0..200 { for &(a, b, r) in &train {
        let e = r - mu - dot(&su[a], &sv[b]);
        for k in 0..K { let (p, q) = (su[a][k], sv[b][k]); su[a][k] += 0.02 * (e * q - 0.02 * p); sv[b][k] += 0.02 * (e * p - 0.02 * q); }
    } }
    let (sgd_ms, sgd) = (t0.elapsed().as_secs_f64() * 1e3, rmse(&su, &sv, mu, &test));
    println!("{} train ratings; test RMSE: mean-only {base:.3}, ALS {als:.3} ({als_ms:.1} ms), SGD {sgd:.3} ({sgd_ms:.1} ms)", train.len());
    assert!(als < 0.1 * base && sgd < 0.2 * base, "factorization must beat the mean by a wide margin");
    // --- implicit ALS: events happen where the hidden taste score is high; confidence c = 1 + 20 ---
    let (mut ptr, mut pte) = (vec![], vec![]);
    for u in 0..nu { for i in 0..ni {
        if rng.f() < 1.0 / (1.0 + (3.0 - 12.0 * dot(&tu[u], &ti[i])).exp()) {
            if rng.f() < 0.8 { ptr.push((u, i)) } else { pte.push((u, i)) }
        }
    } }
    let (mut ur, mut ir) = (vec![vec![]; nu], vec![vec![]; ni]);
    for &(a, b) in &ptr { ur[a].push((b, 1.0, 20.0)); ir[b].push((a, 1.0, 20.0)); }
    let (mut iu, mut iv) = (mk(&mut rng, nu), mk(&mut rng, ni));
    for _ in 0..10 { let g = gram(&iv); half_step(&mut iu, &iv, &ur, 0.1, Some(&g)); let g = gram(&iu); half_step(&mut iv, &iu, &ir, 0.1, Some(&g)); }
    let mut pop = vec![0usize; ni];
    for &(_, b) in &ptr { pop[b] += 1; }
    let hit = |score: &dyn Fn(usize, usize) -> f32| {
        let h = pte.iter().filter(|&&(a, b)| {
            let seen = |j: usize| ur[a].iter().any(|o| o.0 == j);
            (0..ni).filter(|&j| !seen(j) && score(a, j) > score(a, b)).count() < 10
        });
        h.count() as f32 / pte.len() as f32
    };
    let (h_pop, h_als) = (hit(&|_, j| pop[j] as f32), hit(&|a, j| dot(&iu[a], &iv[j])));
    println!("hit@10 on held-out positives: popularity {h_pop:.3}, implicit ALS {h_als:.3}");
    assert!(h_als > 2.0 * h_pop);
}

Measured on an Apple M4 Pro (release, 5 runs, same output each): 5755 train ratings; test RMSE: mean-only 0.654, ALS 0.043 (0.4-0.6 ms), SGD 0.087 (12.4-14.2 ms); hit@10 on held-out positives: popularity 0.084, implicit ALS 0.290. Label noise std about 0.029; SGD untuned.

Complexity and constants

• SGD epoch O(nnz·k); ALS iteration O(nnz·k² + (U+I)·k³), implicit adds O((U+I)·k²) for Gram matrices. Memory (U+I)·k·4 bytes.
• Measured here at k=4: SGD about 12 ns per rating update (200 epochs × 5,755 in about 13.5 ms); ALS about 2.6 ns per rating visit (30 half-steps, about 0.45 ms).

Efficiency

• Cost: One SGD update touches two k-vectors: about 12 ns at k=4 on one M4 Pro core (measured here); latency-bound once factors leave cache. One score is k MACs; a 10⁶-item scan at k=64 reads 256 MB, bandwidth-bound (about 7 ms per core, measured on the Two-tower retrieval page).
• Levers:
  • Parallel ALS rows — near-linear in cores, rows are independent (Zhou et al. 2008) — large nnz.
  • Conjugate gradient solves — a few O(k²) steps replace the O(k³) factorisation (Takács et al. 2011; implicit) — k of order 100 or more.
  • Gram trick — implicit data costs O(nnz), not U×I (Hu et al. 2008) — always for implicit.
  • ANN index over Q — sublinear top-k, recall tradeoff (HNSW) — N above order 10⁶.
• Naive → standard → state of the art:
  • Fill gaps, dense SVD of U×I · U·I memory, biased by the fill
  • ALS or SGD on observed cells, exact scan · O(nnz·k²) per iteration, O(N·k) per query
  • Implicit ALS with CG, multi-core or GPU, plus ANN · O(nnz·k) per iteration, sublinear query

Implementation notes

• CSR by row and by column; f32 factors, k a multiple of 8.
• Ridge scaled by row count (λ·n_u, ALS-WR) or not: libraries differ (Spark scales by count); compare λ within one tool.
• Split by time; judge by top-k metrics (Evaluation metrics).

Pitfalls

• Missing explicit cells read as 0 bias every prediction.
• Zero init stays zero.
• Low RMSE does not give good top-k lists.
• Popular items get large norms and dominate; add biases, normalise or re-rank.

Variants and alternatives

In production

• Netflix Prize rating prediction (Koren et al. 2009), Spark MLlib ALS, implicit.
• Link prediction, drug-target screening, student-question difficulty (Knowledge tracing) use the same completion.

Implementations

Sources

• Koren, Bell, Volinsky, Matrix Factorization Techniques for Recommender Systems, IEEE Computer 2009. paper
• Hu, Koren, Volinsky, Collaborative Filtering for Implicit Feedback Datasets, ICDM 2008. paper
• Zhou, Wilkinson, Schreiber, Pan, Large-Scale Parallel Collaborative Filtering for the Netflix Prize, AAIM 2008. paper
• Takács, Pilászy, Tikk, Applications of the Conjugate Gradient Method for Implicit Feedback Collaborative Filtering, RecSys 2011. paper
• Koren, Factorization Meets the Neighborhood, KDD 2008. paper
• Rendle et al., BPR: Bayesian Personalized Ranking from Implicit Feedback, UAI 2009. paper