Skip to content

⚡ Bolt: 최적화된 alpha 파라미터 gradient 연산으로 메모리 할당 병목 제거 #174

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
seonghobae wants to merge 1 commit into from
Closed

⚡ Bolt: 최적화된 alpha 파라미터 gradient 연산으로 메모리 할당 병목 제거 #174

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .jules/bolt.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,3 +33,6 @@
## 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 - Fast reduction of matrix-multiplication-like element-wise products over indices
**Learning:** When calculating an element-wise product followed by a sum over an axis where one operand is broadcasted along factors (e.g. `(e * theta[:, factors]).sum(axis=0)`), numpy allocates a massive intermediate `N x J` array before summation. For large data (e.g. `N=5000, J=500`), this takes seconds and limits scalability.
**Action:** Replace `(A * B[:, factors]).sum(axis=0)` with `(A.T @ B)[np.arange(A.shape[1]), factors]` to eliminate the intermediate array. This allows the linear algebra routines (BLAS) to compute the projection dynamically using dot products, offering an order-of-magnitude speedup.
52 changes: 39 additions & 13 deletions python/fast_mlsirm/objective.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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())
Expand All @@ -109,7 +128,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 * params.theta[:, factors]).sum(axis=0) * a
# Optimized gradient computation: avoid intermediate N x J array allocation during grad_alpha calculation
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
Expand All @@ -129,7 +149,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))
Expand Down Expand Up @@ -195,7 +217,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),
Expand All @@ -205,7 +229,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))
Expand Down
Loading