⚡ Bolt: Vectorize Newton-Raphson M-step in MMLE estimator

.jules/bolt.md

@@ -40,3 +40,7 @@
 ## 2024-08-01 - Avoid allocating N x J arrays in axis reductions
 **Learning:** Operations like `(e * theta[:, factors]).sum(axis=0) * a` allocate a full N x J array just to compute the elementwise product before summing over the rows. Using dense matrix multiplication followed by integer indexing `(e.T @ theta)[np.arange(e.shape[1]), factors] * a` avoids the massive intermediate allocation and leverages highly optimized BLAS operations.
 **Action:** Replace `(A * B[:, factors]).sum(axis=0)` patterns with dense matrix multiplication `(A.T @ B)[np.arange(A.shape[1]), factors]` to improve speed and reduce memory overhead, specially when computing gradients for parameters across dimensions.
+
+## 2026-08-02 - Vectorize iterative Newton steps across independent dimensions
+**Learning:** During the M-step of MMLE estimation, performing a Python loop over `n_items` to run an iterative inner loop (like Newton-Raphson) creates severe overhead. Independent scalar iterative calculations (like updating item difficulties) are very slow in Python.
+**Action:** Replace outer Python loops over independent elements with a fully vectorized inner loop. Use a boolean array to track active components (e.g. `active = np.ones(n_items, dtype=bool)`) and process only the active subset in each iteration, stopping early for components that converge via boolean masking. This transforms slow Python scalar iterations into fast C-level numpy operations.

python/fast_mlsirm/estimators/mmle.py

@@ -122,30 +122,49 @@ 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
+
+        # Optimized M-step: Vectorize Newton-Raphson across all active items simultaneously
+        # to avoid slow Python loops and per-item array overhead.
+        active = np.ones(n_items, dtype=bool)
+        nodes_sq = nodes * nodes
+
+        for _ in range(25):
+            if not active.any():
+                break
+
+            eta = a_new[active, None] * nodes[None, :] + b_new[active, None]
+            p = _sigmoid(eta)
+
+            n_act = n_iq[active]
+            r_act = r_iq[active]
+
+            w = n_act * p * (1.0 - p)
+            resid = r_act - n_act * p
+
+            g_a = resid @ nodes - ridge_a * a_new[active]
+            g_b = resid.sum(axis=1) - ridge_b * b_new[active]
+
+            h_aa = -(w @ nodes_sq) - ridge_a
+            h_bb = -w.sum(axis=1) - ridge_b
+            h_ab = -(w @ nodes)
+
+            det = h_aa * h_bb - h_ab * h_ab
+            det_ok = np.abs(det) >= 1e-12
+
+            da = np.zeros_like(g_a)
+            db = np.zeros_like(g_b)
+
+            da[det_ok] = (h_bb[det_ok] * g_a[det_ok] - h_ab[det_ok] * g_b[det_ok]) / det[det_ok]
+            db[det_ok] = (h_aa[det_ok] * g_b[det_ok] - h_ab[det_ok] * g_a[det_ok]) / det[det_ok]
+
+            a_new[active] -= da
+            b_new[active] -= db
+            a_new[active] = np.clip(a_new[active], 1e-3, 10.0)
+
+            done = (~det_ok) | ((np.abs(da) + np.abs(db)) < 1e-8)
+            if done.any():
+                active_idx = np.where(active)[0]
+                active[active_idx[done]] = False
 
         a, b = a_new, b_new