diff --git a/.jules/bolt.md b/.jules/bolt.md index 73e3fbaf9..c8e44e48c 100644 --- a/.jules/bolt.md +++ b/.jules/bolt.md @@ -33,3 +33,8 @@ ## 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. + + +## 2024-05-19 - Fast matrix-matrix products in NumPy +**Learning:** Using element-wise multiplication followed by sum reduction over an axis (e.g. `(e_over_d.T * xi).sum(axis=...)` or similar forms) can be much slower than equivalent matrix multiplication `e_over_d.T @ xi`. +**Action:** Rewrite scalar reductions of broadcasted products into matrix multiplication `@` wherever possible. For instance, `xi * sum_e_over_d - np.dot(e_over_d, zeta)` is fast but we can convert `np.dot` to `@` directly or write `e_over_d @ zeta` for clarity and performance. diff --git a/python/fast_mlsirm/objective.py b/python/fast_mlsirm/objective.py index e43df9809..f4a3cb566 100644 --- a/python/fast_mlsirm/objective.py +++ b/python/fast_mlsirm/objective.py @@ -133,7 +133,9 @@ def neg_loglik_and_grad( grad_b = e.sum(axis=0) grad_alpha = np.zeros_like(params.alpha) if free_alpha: - grad_alpha = (e * params.theta[:, factors]).sum(axis=0) * a + # Optimized alpha gradient: avoid N x J array allocation by using dense matrix multiplication + # and advanced integer indexing instead of element-wise multiplication and summation. + grad_alpha = (e.T @ params.theta)[np.arange(e.shape[1]), factors] * a # Optimized gradient computation: replace loop over dimensions with matrix multiplication # We embed 'a' directly into the projection matrix to avoid a JxD intermediate array allocation during multiplication @@ -150,10 +152,10 @@ def neg_loglik_and_grad( # Optimized gradient computation: avoid 3D array creation, use 2D matrix multiplication instead e_over_d = e / distance 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)) + grad_xi = -gamma * (params.xi * sum_e_over_d - e_over_d @ params.zeta) sum_e_over_d_j = e_over_d.sum(axis=0, keepdims=True).T - grad_zeta = gamma * (np.dot(e_over_d.T, params.xi) - params.zeta * sum_e_over_d_j) + grad_zeta = gamma * (e_over_d.T @ params.xi - params.zeta * sum_e_over_d_j) # Optimized gradient computation: avoid intermediate array allocation by using vdot grad_tau = float(-gamma * np.vdot(e, distance))