⚡ Bolt: numpy 연산 최적화를 통한 그레디언트 계산 속도 향상

.jules/bolt.md

@@ -5,3 +5,7 @@
 ## 2024-05-19 - NumPy Array Allocation and Advanced Vectorization
 **Learning:** `np.sum(x * x)` 패턴은 파이썬 내에서 곱셈을 위한 새로운 중간 배열을 메모리에 할당하고 이후에 그 배열의 합을 구하게 되어 성능 저하를 야기합니다. `np.vdot(x, x)` (또는 `np.einsum`)를 활용하면 중간 메모리 할당을 우회할 수 있어 속도가 비약적으로 증가합니다.
 **Action:** 거대한 배열의 크기나 요소 수와 관련된 최적화 시, `np.sum(x * x)` 대신 `np.vdot(x, x)`를 사용해 오버헤드를 방지합니다.
+
+## 2024-07-01 - NumPy Array Allocation and Advanced Vectorization in gradient calculation
+**Learning:** `(e * a[None, :] * params.theta[:, factors]).sum(axis=0)` allocates a 3D intermediate array for the broadcasted multiplication of `a`, which is highly inefficient for large arrays. Additionally, `(e * (-gamma * distance)).sum()` allocates intermediate arrays for the multiplication by `-gamma` and the element-wise multiplication of `e` and `distance`.
+**Action:** Extract variables that are constant over an axis being summed over, such as pulling `a` out of the sum to make it `a * sum(...)`. Use `np.einsum('ij,ij->j', A, B)` for element-wise multiplication and summation over an axis to avoid intermediate allocations. Use `np.vdot(A, B)` for sum of element-wise multiplication of entire arrays to avoid intermediate allocations.

python/fast_mlsirm/objective.py

@@ -105,7 +105,8 @@ def neg_loglik_and_grad(
     grad_b = e.sum(axis=0)
     grad_alpha = np.zeros_like(params.alpha)
     if free_alpha:
-        grad_alpha = (e * a[None, :] * params.theta[:, factors]).sum(axis=0)
+        # Optimized grad_alpha computation: skip large 3D broadcast by extracting `a` and using np.einsum
+        grad_alpha = a * np.einsum('ij,ij->j', e, params.theta[:, factors])
 
     # Optimized gradient computation: replace loop over dimensions with matrix multiplication
     # np.eye(...)[factors] creates a one-hot encoding (J x D), projecting J items onto D dimensions
@@ -127,7 +128,8 @@ def neg_loglik_and_grad(
         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_tau = float((e * (-gamma * distance)).sum())
+        # Optimized grad_tau computation: replace element-wise sum with vdot to skip intermediate allocations
+        grad_tau = -gamma * float(np.vdot(e, distance))
 
     nll += _add_penalty(params, penalty, free_alpha=free_alpha, uses_space=uses_space)
     grad_theta += penalty.lambda_theta * params.theta