⚡ Bolt: 그레이디언트 연산 및 수학 유틸리티 벡터화 성능 최적화

.jules/bolt.md

@@ -33,3 +33,10 @@
 ## 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 - Gradient alpha intermediate allocations
+**Learning:** In Python numerical gradients, operations like `(e * theta[:, factors]).sum(axis=0)` create massive intermediate arrays of shape (N, J). For large datasets, this severely impacts performance and memory.
+**Action:** Replace such calculations with transposed matrix multiplication and index selection (e.g., `(e.T @ theta)[np.arange(e.shape[1]), factors]`) to leverage BLAS optimization and avoid the large intermediate array allocation.
+
+## 2025-05-19 - Vectorizing standard math operations on 2D arrays
+**Learning:** Functions like `standardize` that assume 1D inputs force loops over 2D array columns, causing performance drops due to Python iteration overhead.
+**Action:** Upgrade math utilities to seamlessly support 2D arrays natively with  aggregations and fallback boolean masking, eliminating Python loops in caller code.

python/fast_mlsirm/fit.py

@@ -7,8 +7,12 @@
 from .backend import normalize_device, resolve_backend
 from .config import FitConfig
 from .math import logit, normalize_latent_positions, standardize
-from .objective import (model_flags, neg_loglik_and_grad, prepare_response,
-                        validate_factor_id)
+from .objective import (
+    model_flags,
+    neg_loglik_and_grad,
+    prepare_response,
+    validate_factor_id,
+)
 from .types import FitResult, MLSIRMParams
 
 
@@ -58,7 +62,9 @@ def fit(
             best = candidate
 
     if best is None:
-        raise RuntimeError("Optimization failed to find a valid fit.")  # pragma: no cover
+        raise RuntimeError(
+            "Optimization failed to find a valid fit."
+        )  # pragma: no cover
     return best
 
 
@@ -222,8 +228,7 @@ def _initial_params(
     denom = np.maximum(observed @ item_mask.astype(np.float64), 1)
     x = ((y * observed) @ item_mask.astype(np.float64)) / denom
 
-    for d in range(n_dims):
-        theta[:, d] = standardize(x[:, d])
+    theta = standardize(x)
 
     item_counts = np.maximum(observed.sum(axis=0), 1)
     item_means = (y * observed).sum(axis=0) / item_counts

python/fast_mlsirm/math.py

@@ -21,11 +21,26 @@ def logit(p: np.ndarray | float, eps: float = 1e-6) -> np.ndarray:
 
 def standardize(x: np.ndarray) -> np.ndarray:
     x = np.asarray(x, dtype=np.float64)
-    mean = np.nanmean(x)
-    sd = np.nanstd(x)
-    if not np.isfinite(sd) or sd < 1e-12:
-        return np.zeros_like(x, dtype=np.float64)
-    return (x - mean) / sd
+
+    if x.ndim == 1:
+        mean = np.nanmean(x)
+        sd = np.nanstd(x)
+        if not np.isfinite(sd) or sd < 1e-12:
+            return np.zeros_like(x, dtype=np.float64)
+        return (x - mean) / sd
+
+    # Vectorized 2D standardization across columns
+    mean = np.nanmean(x, axis=0)
+    sd = np.nanstd(x, axis=0)
+
+    # Handle zero/invalid standard deviations safely
+    valid = np.isfinite(sd) & (sd >= 1e-12)
+    sd_safe = np.where(valid, sd, 1.0)
+
+    result = (x - mean) / sd_safe
+    result[:, ~valid] = 0.0
+
+    return result
 
 
 def normalize_latent_positions(params: MLSIRMParams) -> MLSIRMParams:

python/fast_mlsirm/objective.py

@@ -2,13 +2,20 @@
 
 import numpy as np
 
-from .backend import load_rust_core, normalize_backend, normalize_device, resolve_backend
+from .backend import (
+    load_rust_core,
+    normalize_backend,
+    normalize_device,
+    resolve_backend,
+)
 from .config import FitConfig, PenaltyConfig
 from .math import sigmoid, softplus
 from .types import MLSIRMParams
 
 
-def prepare_response(responses: np.ndarray, mask: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]:
+def prepare_response(
+    responses: np.ndarray, mask: np.ndarray | None = None
+) -> tuple[np.ndarray, np.ndarray]:
     y = np.asarray(responses, dtype=np.float64)
     if y.ndim != 2:
         raise ValueError("responses must be a 2D matrix")
@@ -59,9 +66,11 @@ def linear_predictor(
 
     if uses_space:
         # Optimized distance computation: replace O(N*J*D) 3D broadcast with O(N*J) 2D dot product
-        xi_sq = np.einsum('ij,ij->i', params.xi, params.xi)
-        zeta_sq = np.einsum('ij,ij->i', params.zeta, params.zeta)
-        dist_sq = xi_sq[:, None] + zeta_sq[None, :] - 2 * np.dot(params.xi, params.zeta.T)
+        xi_sq = np.einsum("ij,ij->i", params.xi, params.xi)
+        zeta_sq = np.einsum("ij,ij->i", params.zeta, params.zeta)
+        dist_sq = (
+            xi_sq[:, None] + zeta_sq[None, :] - 2 * np.dot(params.xi, params.zeta.T)
+        )
         dist_sq = np.maximum(dist_sq, 0.0)
         distance = np.sqrt(dist_sq + eps_distance)
         gamma = params.gamma
@@ -84,10 +93,18 @@ def neg_loglik_and_grad(
 ) -> tuple[float, MLSIRMParams, float]:
     config = config or FitConfig()
     requested_backend = normalize_backend(backend)
-    normalized_backend = resolve_backend(requested_backend) if requested_backend == "auto" else requested_backend
+    normalized_backend = (
+        resolve_backend(requested_backend)
+        if requested_backend == "auto"
+        else requested_backend
+    )
     if normalized_backend == "rust":
-        resolved_device = normalize_device(device if device is not None else config.rust_device)
-        return _neg_loglik_and_grad_rust(responses, factor_id, params, config, mask, resolved_device)
+        resolved_device = normalize_device(
+            device if device is not None else config.rust_device
+        )
+        return _neg_loglik_and_grad_rust(
+            responses, factor_id, params, config, mask, resolved_device
+        )
 
     model = config.normalized_model()
     penalty = config.penalty
@@ -99,7 +116,9 @@ def neg_loglik_and_grad(
 
     free_alpha, uses_space = model_flags(model)
     a = params.a if free_alpha else np.ones_like(params.alpha)
-    eta, distance = linear_predictor(params, factors, model=model, eps_distance=config.eps_distance)
+    eta, distance = linear_predictor(
+        params, factors, model=model, eps_distance=config.eps_distance
+    )
     pi = sigmoid(eta)
     entry_loss = (softplus(eta) - y * eta) * observed
     nll = float(entry_loss.sum())
@@ -109,7 +128,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 gradient computation: avoid massive N x J intermediate array allocation
+        # using a fast dot product e.T @ theta and selecting the required factor indices.
+        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
@@ -129,7 +150,9 @@ def neg_loglik_and_grad(
         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
-        grad_zeta = gamma * (np.dot(e_over_d.T, params.xi) - params.zeta * sum_e_over_d_j)
+        grad_zeta = gamma * (
+            np.dot(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))
@@ -195,7 +218,9 @@ def _neg_loglik_and_grad_rust(
         device,
     )
     grads = MLSIRMParams(
-        theta=np.asarray(gradients["theta"], dtype=np.float64).reshape(params.theta.shape),
+        theta=np.asarray(gradients["theta"], dtype=np.float64).reshape(
+            params.theta.shape
+        ),
         alpha=np.asarray(gradients["alpha"], dtype=np.float64),
         b=np.asarray(gradients["b"], dtype=np.float64),
         xi=np.asarray(gradients["xi"], dtype=np.float64).reshape(params.xi.shape),
@@ -205,7 +230,9 @@ def _neg_loglik_and_grad_rust(
     return float(objective), grads, float(loglik)
 
 
-def _add_penalty(params: MLSIRMParams, penalty: PenaltyConfig, free_alpha: bool, uses_space: bool) -> float:
+def _add_penalty(
+    params: MLSIRMParams, penalty: PenaltyConfig, free_alpha: bool, uses_space: bool
+) -> float:
     # Optimized penalty calculation: replace np.sum(x * x) with np.vdot(x, x) to avoid intermediate array allocation
     value = 0.5 * penalty.lambda_theta * float(np.vdot(params.theta, params.theta))
     value += 0.5 * penalty.lambda_b * float(np.vdot(params.b, params.b))