From 292fce1f4cfe239b33f46c461051fed824608996 Mon Sep 17 00:00:00 2001 From: seonghobae <8172694+seonghobae@users.noreply.github.com> Date: Sun, 26 Jul 2026 02:41:01 +0000 Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Vectorize=20MMLE=20Newton-R?= =?UTF-8?q?aphson=20M-step=20using=20active=20masks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .jules/bolt.md | 4 ++ python/fast_mlsirm/estimators/mmle.py | 73 ++++++++++++++++++--------- 2 files changed, 53 insertions(+), 24 deletions(-) diff --git a/.jules/bolt.md b/.jules/bolt.md index 73e3fbaf9..1cb08a0d5 100644 --- a/.jules/bolt.md +++ b/.jules/bolt.md @@ -33,3 +33,7 @@ ## 2025-05-19 - Dot product scalar gradients allocation **Learning:** During gradient calculation, `float((e * (-gamma * distance)).sum())` creates two full-size `(N, J)` arrays: one for the scaled distance and one for the element-wise multiplication before reduction. **Action:** Replace `(A * B).sum()` with `np.vdot(A, B)` when scalar reduction is needed over matrix multiplication (where `B` can incorporate scalars naturally like `-gamma * np.vdot(A, B)`). This entirely avoids the 2D array allocation overhead and yields order-of-magnitude improvements in scalar gradient components. + +## 2025-05-19 - Vectorizing Item-Level Iterative Algorithms (Newton-Raphson) +**Learning:** In optimization loops (like the Newton-Raphson M-step in MMLE), iterating over hundreds or thousands of items individually with a Python `for` loop causes severe performance overhead. While some items converge quickly, the loop overhead for the entire set of items dominates. +**Action:** Vectorize iterative item-level updates by replacing the Python loop with a simultaneous update over an active boolean mask (`active = np.ones(n_items, dtype=bool)`). Only compute updates for items that haven't converged yet (`active = new_active`), filtering calculations with advanced indexing and native NumPy vector operations. diff --git a/python/fast_mlsirm/estimators/mmle.py b/python/fast_mlsirm/estimators/mmle.py index 222b977fd..62289142b 100644 --- a/python/fast_mlsirm/estimators/mmle.py +++ b/python/fast_mlsirm/estimators/mmle.py @@ -121,30 +121,55 @@ def fit_mmle_2pl( a_new = a.copy() b_new = b.copy() - for i in range(n_items): - ai, bi = a[i], b[i] - # Newton steps on the item's expected log-likelihood over nodes. - for _ in range(25): - eta = ai * nodes + bi - p = _sigmoid(eta) - w = n_iq[i] * p * (1.0 - p) - resid = r_iq[i] - n_iq[i] * p - g_a = float((resid * nodes).sum()) - ridge_a * ai - g_b = float(resid.sum()) - ridge_b * bi - h_aa = -float((w * nodes * nodes).sum()) - ridge_a - h_bb = -float(w.sum()) - ridge_b - h_ab = -float((w * nodes).sum()) - det = h_aa * h_bb - h_ab * h_ab - if abs(det) < 1e-12: - break - da = (h_bb * g_a - h_ab * g_b) / det - db = (h_aa * g_b - h_ab * g_a) / det - ai -= da - bi -= db - ai = float(np.clip(ai, 1e-3, 10.0)) - if abs(da) + abs(db) < 1e-8: - break - a_new[i], b_new[i] = ai, bi + active = np.ones(n_items, dtype=bool) + nodes_sq = nodes * nodes + + # Vectorized Newton steps over all active items simultaneously + for _ in range(25): + if not active.any(): + break + + ai = a_new[active] + bi = b_new[active] + + eta = ai[:, None] * nodes[None, :] + bi[:, None] + p = _sigmoid(eta) + + n_iq_act = n_iq[active] + r_iq_act = r_iq[active] + + w = n_iq_act * p * (1.0 - p) + resid = r_iq_act - n_iq_act * p + + g_a = (resid * nodes).sum(axis=1) - ridge_a * ai + g_b = resid.sum(axis=1) - ridge_b * bi + + h_aa = -(w * nodes_sq).sum(axis=1) - ridge_a + h_bb = -w.sum(axis=1) - ridge_b + h_ab = -(w * nodes).sum(axis=1) + + det = h_aa * h_bb - h_ab * h_ab + valid = np.abs(det) >= 1e-12 + + da = np.zeros_like(ai) + db = np.zeros_like(bi) + + da[valid] = (h_bb[valid] * g_a[valid] - h_ab[valid] * g_b[valid]) / det[valid] + db[valid] = (h_aa[valid] * g_b[valid] - h_ab[valid] * g_a[valid]) / det[valid] + + ai -= da + bi -= db + ai = np.clip(ai, 1e-3, 10.0) + + a_new[active] = ai + b_new[active] = bi + + not_converged = (np.abs(da) + np.abs(db)) >= 1e-8 + keep_going = valid & not_converged + + new_active = np.zeros(n_items, dtype=bool) + new_active[active] = keep_going + active = new_active a, b = a_new, b_new