diff --git a/.jules/bolt.md b/.jules/bolt.md index 73e3fbaf9..83f337447 100644 --- a/.jules/bolt.md +++ b/.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. diff --git a/python/fast_mlsirm/fit.py b/python/fast_mlsirm/fit.py index e185c667e..06ff223b6 100644 --- a/python/fast_mlsirm/fit.py +++ b/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 diff --git a/python/fast_mlsirm/math.py b/python/fast_mlsirm/math.py index a4a444af0..8fcc237cc 100644 --- a/python/fast_mlsirm/math.py +++ b/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: diff --git a/python/fast_mlsirm/objective.py b/python/fast_mlsirm/objective.py index f6c9437d0..b54f033a2 100644 --- a/python/fast_mlsirm/objective.py +++ b/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))