diff --git a/.jules/bolt.md b/.jules/bolt.md index cb92f7396..a2a8d88a7 100644 --- a/.jules/bolt.md +++ b/.jules/bolt.md @@ -44,3 +44,9 @@ ## 2026-08-04 - Matrix-vector reductions for MMLE quadrature nodes **Learning:** In the NumPy MMLE reference fallback, expressions such as `(resid * nodes[None, :]).sum(axis=1)` materialize an item-by-node intermediate array. The mathematically equivalent matrix-vector product `resid @ nodes` avoids that broadcast temporary and can use the configured NumPy linear-algebra backend. Runtime gains depend on matrix shape, memory layout, BLAS implementation, and threading, so no universal percentage improvement should be claimed without a reproducible benchmark. **Action:** Prefer a matrix-vector product for equivalent quadrature-node reductions when dtype, shape, and numerical parity are preserved. Keep Rust as the primary production path, retain the NumPy implementation as a tested reference fallback, and benchmark representative workloads before making quantitative performance claims. +## 2025-05-19 - Using Matrix Multiplication to Avoid Axis Reduction Overhead +**Learning:** In heavily used iteration loops (like MMLE E-step or M-step), calling `.sum(axis=...)` adds a non-trivial python dispatch and array allocation overhead. NumPy translates these array reductions via slow inner C paths rather than leveraging highly optimized BLAS. Using a dense vector of ones and `matmul` (`@`) performs the same reduction much faster by keeping operations inside the BLAS library level. +**Action:** Replace `array.sum(axis=1)` with `array @ np.ones(array.shape[1])` to reduce reduction overhead when iterating extensively. Preallocate `ones` when used within the loop. +## 2025-05-19 - Using BLAS for Axis Reductions over N-dimension +**Learning:** In heavily used iteration loops (like calculating gradients in `objective.py` across thousands of persons), calling `.sum(axis=0)` incurs unnecessary internal numpy reduction overhead compared to highly optimized linear algebra primitives. Using matrix multiplication with a pre-allocated vector of ones (`ones @ array`) keeps the operation entirely inside the fast BLAS library level. For large dimension N, this speeds up reduction considerably. +**Action:** Replace `array.sum(axis=0)` with `ones @ array` (where `ones = np.ones(array.shape[0])`) to reduce reduction overhead when repeatedly aggregating across the outer dimension in gradients or likelihoods. Preallocate `ones` to avoid repeated memory allocations. diff --git a/python/fast_mlsirm/estimators/mmle.py b/python/fast_mlsirm/estimators/mmle.py index 6393d7f82..93774ce52 100644 --- a/python/fast_mlsirm/estimators/mmle.py +++ b/python/fast_mlsirm/estimators/mmle.py @@ -83,6 +83,10 @@ def fit_mmle_2pl( nodes, weights = gauss_hermite_nodes(n_nodes) # (Q,), (Q,) log_weights = np.log(weights) + # Optimization: define matrix multipliers to avoid .sum(axis=1) intermediate overheads + ones_1d = np.ones(n_nodes, dtype=np.float64) + ones_2d = ones_1d[:, None] + rng = np.random.default_rng(seed) # Init: a=1, b from observed item log-odds of endorsement. p_item = (y_filled * obs_f).sum(0) / np.clip(obs_f.sum(0), 1.0, None) @@ -109,7 +113,7 @@ def fit_mmle_2pl( # Normalize across nodes (log-sum-exp) max_lj = log_joint.max(axis=1, keepdims=True) stab = np.exp(log_joint - max_lj) - denom = stab.sum(axis=1, keepdims=True) + denom = stab @ ones_2d posterior = stab / denom # (n_persons, Q) person_loglik = max_lj[:, 0] + np.log(denom[:, 0]) total_loglik = float(person_loglik.sum()) @@ -144,10 +148,10 @@ def fit_mmle_2pl( # Optimization: Replace element-wise multiply and axis reduction with dense matrix multiplication # to avoid large intermediate array allocations g_a = resid @ nodes - ridge_a * ai - g_b = resid.sum(axis=1) - ridge_b * bi + g_b = (resid @ ones_1d) - ridge_b * bi h_aa = -(w @ nodes_sq) - ridge_a - h_bb = -w.sum(axis=1) - ridge_b + h_bb = -(w @ ones_1d) - ridge_b h_ab = -(w @ nodes) det = h_aa * h_bb - h_ab * h_ab diff --git a/python/fast_mlsirm/objective.py b/python/fast_mlsirm/objective.py index 59f6ff316..ddad00ace 100644 --- a/python/fast_mlsirm/objective.py +++ b/python/fast_mlsirm/objective.py @@ -174,7 +174,9 @@ def neg_loglik_and_grad( loglik = -nll e = (pi - y) * observed - grad_b = e.sum(axis=0) + # Optimization: Use BLAS matrix multiplication instead of slow axis reduction loops for summing N dimension + ones_n = np.ones(e.shape[0], dtype=e.dtype) + grad_b = ones_n @ e grad_alpha = np.zeros_like(params.alpha) if free_alpha: # Optimized gradient computation: avoid intermediate N x J array allocation in element-wise multiplication @@ -197,7 +199,8 @@ def neg_loglik_and_grad( sum_e_over_d = e_over_d.sum(axis=1, keepdims=True) grad_xi = -gamma * (params.xi * sum_e_over_d - np.dot(e_over_d, params.zeta)) - sum_e_over_d_j = e_over_d.sum(axis=0, keepdims=True).T + # Optimization: Use BLAS matrix multiplication to avoid slow axis reduction overhead + sum_e_over_d_j = (ones_n @ e_over_d)[:, None] grad_zeta = gamma * ( np.dot(e_over_d.T, params.xi) - params.zeta * sum_e_over_d_j )