Audit: pin MLS2PLM neg-loglik & gradient to canonical equation

docs/papers/mls2plm-canonical-equations.md

@@ -0,0 +1,124 @@
+# MLS2PLM canonical equations (source of truth for the numeric core)
+
+This note pins the exact equations that `fast-mlsirm` implements, so that any
+future performance refactor or "optimization" of the numeric core can be checked
+against the published formula rather than against a previous (possibly already
+drifted) version of the code.
+
+The canonical implementation is a **simple-structure specialization** of the
+Multidimensional Latent Space Item Response Model. See
+`docs/pr25_formula_review.md` for why the simple-structure restriction is a
+model choice, not an implementation detail that may be silently rewritten into
+the full discrimination-vector model.
+
+## Source papers
+
+- Jeon, M., Jin, I. H., Schweinberger, M., & Baugh, S. (2021). *Mapping
+  Unobserved Item-Respondent Interactions: A Latent Space Item Response Model
+  with Interaction Map.* Psychometrika, 86(2), 378-403.
+  DOI: [10.1007/s11336-021-09762-5](https://doi.org/10.1007/s11336-021-09762-5)
+  — defines the Latent Space Item Response Model (LSIRM):
+  `logit P(Y = 1) = theta_p + b_i - gamma * ||u_p - v_i||`.
+- Molenaar, D., & Jeon, M. (2026). *Regularized Joint Maximum Likelihood
+  Estimation of Latent Space Item Response Models.* Psychometrika, 91, 335-359.
+  DOI: [10.1017/psy.2025.10068](https://doi.org/10.1017/psy.2025.10068)
+  — regularized (penalized) joint maximum-likelihood estimation, which is the
+  estimation strategy this package uses (L2 penalties on the parameter blocks).
+- Kang, I., & Jeon, M. (2025). *Multidimensional Latent Space Item Response
+  Models: A Note on the Relativity of Conditional Dependence.* Psychometrika,
+  90(2), 799-826. DOI: [10.1017/psy.2025.5](https://doi.org/10.1017/psy.2025.5)
+  — the multidimensional (MLS2PLM) generalization.
+
+PDFs are not redistributed here (the primary sources are not open-access); cite
+via the DOIs above. This document reproduces only the equations needed to verify
+the code, which are facts and not copyrightable expression.
+
+## Model
+
+For person `p`, item `i` with `factor_id(i) = d(i)`, latent trait vector
+`theta_p`, person latent position `xi_p` and item latent position `zeta_i` in a
+shared latent space:
+
+```text
+r_pi  = sqrt( || xi_p - zeta_i ||^2 + eps_distance )          (regularized distance)
+eta_pi = a_i * theta_{p, d(i)} + b_i - gamma * r_pi           (linear predictor)
+P(Y_pi = 1) = sigmoid(eta_pi)
+```
+
+with the positivity reparameterizations
+
+```text
+a_i   = exp(alpha_i)        (discrimination, > 0)
+gamma = exp(tau)            (interaction weight, > 0)
+```
+
+`eps_distance` (default `1e-8`) smooths the Euclidean distance so that `r_pi`
+and its gradient are finite when `xi_p = zeta_i`. It is applied **inside** the
+square root in both the forward pass and the gradient, keeping `d r_pi / d xi_p
+= (xi_p - zeta_i) / r_pi` exactly consistent.
+
+## Objective (penalized negative log-likelihood)
+
+Over observed entries `O` (mask), with Bernoulli responses `y_pi in {0, 1}`:
+
+```text
+NLL_data = sum_{(p,i) in O} [ softplus(eta_pi) - y_pi * eta_pi ]
+```
+
+This is the standard binary cross-entropy written in a numerically stable form:
+`softplus(eta) - y*eta = -[ y*log P + (1-y)*log(1-P) ]`. The reported
+log-likelihood is `loglik = -NLL_data` (data term only, before the penalty).
+
+The regularized objective adds L2 penalties (Molenaar & Jeon, 2026):
+
+```text
+NLL = NLL_data
+    + 0.5 * lambda_theta * ||theta||^2
+    + 0.5 * lambda_b     * ||b||^2
+    + 0.5 * lambda_alpha * ||alpha - mu_alpha||^2      (only if a_i is free)
+    + 0.5 * lambda_xi    * ||xi||^2                    (only if uses_space)
+    + 0.5 * lambda_zeta  * ||zeta||^2                  (only if uses_space)
+    + 0.5 * lambda_tau   * (tau - mu_tau)^2            (only if uses_space)
+```
+
+(`uses_space` is true for all models except `MIRT`; see "Model variants".)
+
+## Gradient
+
+Let `e_pi = (P(Y_pi = 1) - y_pi)` on observed entries (0 elsewhere). The
+data-term gradients are:
+
+```text
+d NLL / d b_i        = sum_p e_pi
+d NLL / d alpha_i    = sum_p e_pi * a_i * theta_{p, d(i)}   # da_i/dalpha_i = a_i
+d NLL / d theta_{p,k}= sum_{i : d(i)=k} e_pi * a_i
+d NLL / d xi_p       = -gamma * sum_i e_pi * (xi_p - zeta_i) / r_pi
+d NLL / d zeta_i     =  gamma * sum_p e_pi * (xi_p - zeta_i) / r_pi
+d NLL / d tau        = sum_{(p,i)} e_pi * (-gamma * r_pi)   # dgamma/dtau = gamma
+```
+
+The `alpha` and `tau` gradients carry the chain-rule factors from the
+reparameterizations `a_i = exp(alpha_i)` and `gamma = exp(tau)`.
+
+The L2 penalty adds `lambda_* * (param - center)` to each block (`center = 0`
+except `mu_alpha`, `mu_tau`).
+
+## Model variants
+
+`model_flags(model)` toggles two switches:
+
+- `free_alpha` is false for `MLSRM` / `ULSRM` (Rasch-type: `a_i = 1`, so
+  `alpha` has no gradient contribution).
+- `uses_space` is false for `MIRT` (no latent space: `gamma = 0`, `r_pi = 0`,
+  and `xi`, `zeta`, `tau` have no gradient contribution).
+
+## Where this is enforced in code
+
+- `python/fast_mlsirm/objective.py` — `linear_predictor`, `neg_loglik_and_grad`
+  (NumPy reference path) and `_add_penalty`.
+- `crates/mlsirm-core/src/lib.rs` — `neg_loglik_and_grad` (Rust path); kept at
+  parity with the NumPy path by `tests/test_objective.py`.
+- `tests/test_objective.py` — `test_neg_loglik_matches_closed_form_single_entry`
+  and `test_neg_loglik_and_grad_matches_independent_reference` pin the objective
+  and every gradient block to the equations above, computed independently with
+  the Python standard library (no NumPy vectorization shared with production).

tests/test_objective.py

@@ -1,8 +1,210 @@
+import math
+
 import numpy as np
 import pytest
 
 from fast_mlsirm import FitConfig, MLSIRMParams
-from fast_mlsirm.objective import neg_loglik_and_grad, validate_factor_id
+from fast_mlsirm.config import PenaltyConfig as _PenaltyConfig
+from fast_mlsirm.objective import model_flags, neg_loglik_and_grad, validate_factor_id
+
+
+_ZERO_PENALTY = _PenaltyConfig(
+    lambda_theta=0.0,
+    lambda_xi=0.0,
+    lambda_zeta=0.0,
+    lambda_b=0.0,
+    lambda_alpha=0.0,
+    lambda_tau=0.0,
+    mu_alpha=0.0,
+    mu_tau=0.0,
+)
+
+
+def _softplus(x: float) -> float:
+    # Numerically stable log(1 + exp(x)), matches math.py softplus.
+    return max(x, 0.0) + math.log1p(math.exp(-abs(x)))
+
+
+def _sigmoid(x: float) -> float:
+    return 1.0 / (1.0 + math.exp(-x))
+
+
+def _reference_neg_loglik_and_grad(y, factors, params, config):
+    """Independent closed-form reference derived directly from the MLS2PLM
+    canonical equations (docs/papers/mls2plm-canonical-equations.md).
+
+    Pure-Python triple loop with a different structure than the vectorized
+    production path, so it pins the objective, every gradient block, and every
+    summation axis to the published formula rather than to prior code.
+    """
+    free_alpha, uses_space = model_flags(config.normalized_model())
+    penalty = config.penalty
+    eps = config.eps_distance
+
+    n_persons, n_items = y.shape
+    n_traits = params.theta.shape[1]
+    latent_dim = params.xi.shape[1]
+
+    a = [math.exp(params.alpha[j]) if free_alpha else 1.0 for j in range(n_items)]
+    gamma = math.exp(params.tau) if uses_space else 0.0
+
+    nll = 0.0
+    g_theta = [[0.0] * n_traits for _ in range(n_persons)]
+    g_alpha = [0.0] * n_items
+    g_b = [0.0] * n_items
+    g_xi = [[0.0] * latent_dim for _ in range(n_persons)]
+    g_zeta = [[0.0] * latent_dim for _ in range(n_items)]
+    g_tau = 0.0
+
+    for p in range(n_persons):
+        for j in range(n_items):
+            yy = float(y[p, j])
+            if yy == -1.0 or not math.isfinite(yy):  # missing sentinel
+                continue
+            d = int(factors[j])
+            r = 0.0
+            if uses_space:
+                dist2 = eps
+                for k in range(latent_dim):
+                    diff = params.xi[p, k] - params.zeta[j, k]
+                    dist2 += diff * diff
+                r = math.sqrt(dist2)
+            eta = a[j] * params.theta[p, d] + params.b[j] - gamma * r
+            nll += _softplus(eta) - yy * eta
+            e = _sigmoid(eta) - yy
+
+            g_b[j] += e
+            if free_alpha:
+                g_alpha[j] += e * a[j] * params.theta[p, d]
+            g_theta[p][d] += e * a[j]
+            if uses_space:
+                g_tau += e * (-gamma * r)
+                for k in range(latent_dim):
+                    common = gamma * (params.xi[p, k] - params.zeta[j, k]) / r
+                    g_xi[p][k] += e * (-common)
+                    g_zeta[j][k] += e * common
+
+    # Penalty block (Molenaar & Jeon, 2026 regularized JML).
+    for p in range(n_persons):
+        for k in range(n_traits):
+            nll += 0.5 * penalty.lambda_theta * params.theta[p, k] ** 2
+            g_theta[p][k] += penalty.lambda_theta * params.theta[p, k]
+    for j in range(n_items):
+        nll += 0.5 * penalty.lambda_b * params.b[j] ** 2
+        g_b[j] += penalty.lambda_b * params.b[j]
+    if free_alpha:
+        for j in range(n_items):
+            delta = params.alpha[j] - penalty.mu_alpha
+            nll += 0.5 * penalty.lambda_alpha * delta * delta
+            g_alpha[j] += penalty.lambda_alpha * delta
+    if uses_space:
+        for p in range(n_persons):
+            for k in range(latent_dim):
+                nll += 0.5 * penalty.lambda_xi * params.xi[p, k] ** 2
+                g_xi[p][k] += penalty.lambda_xi * params.xi[p, k]
+        for j in range(n_items):
+            for k in range(latent_dim):
+                nll += 0.5 * penalty.lambda_zeta * params.zeta[j, k] ** 2
+                g_zeta[j][k] += penalty.lambda_zeta * params.zeta[j, k]
+        tau_delta = params.tau - penalty.mu_tau
+        nll += 0.5 * penalty.lambda_tau * tau_delta * tau_delta
+        g_tau += penalty.lambda_tau * tau_delta
+
+    return nll, {
+        "theta": np.asarray(g_theta),
+        "alpha": np.asarray(g_alpha),
+        "b": np.asarray(g_b),
+        "xi": np.asarray(g_xi),
+        "zeta": np.asarray(g_zeta),
+        "tau": g_tau,
+    }
+
+
+def test_neg_loglik_matches_closed_form_single_entry():
+    # Single person, single item, 2-D latent space: hand-computable pin of the
+    # canonical MLS2PLM equation eta = exp(alpha)*theta + b - exp(tau)*r.
+    params = MLSIRMParams(
+        theta=np.array([[0.5]]),
+        alpha=np.array([0.2]),
+        b=np.array([0.1]),
+        xi=np.array([[0.3, -0.1]]),
+        zeta=np.array([[-0.2, 0.4]]),
+        tau=0.1,
+    )
+    y = np.array([[1.0]])
+    factors = np.array([0])
+    config = FitConfig(max_iter=1, penalty=_ZERO_PENALTY, eps_distance=1e-8)
+
+    a = math.exp(0.2)
+    gamma = math.exp(0.1)
+    r = math.sqrt((0.3 - (-0.2)) ** 2 + (-0.1 - 0.4) ** 2 + 1e-8)
+    eta = a * 0.5 + 0.1 - gamma * r
+    expected_nll = _softplus(eta) - 1.0 * eta
+    e = _sigmoid(eta) - 1.0
+
+    nll, grad, loglik = neg_loglik_and_grad(y, factors, params, config)
+
+    assert math.isclose(nll, expected_nll, rel_tol=0.0, abs_tol=1e-12)
+    assert math.isclose(loglik, -expected_nll, rel_tol=0.0, abs_tol=1e-12)
+    assert math.isclose(float(grad.b[0]), e, abs_tol=1e-12)
+    assert math.isclose(float(grad.alpha[0]), e * a * 0.5, abs_tol=1e-12)
+    assert math.isclose(float(grad.theta[0, 0]), e * a, abs_tol=1e-12)
+    assert math.isclose(float(grad.tau), e * (-gamma * r), abs_tol=1e-12)
+    # d r / d xi = (xi - zeta)/r ; eta has -gamma*r, so grad_xi = -gamma*e*(xi-zeta)/r
+    assert math.isclose(float(grad.xi[0, 0]), -gamma * e * (0.3 - (-0.2)) / r, abs_tol=1e-12)
+    assert math.isclose(float(grad.xi[0, 1]), -gamma * e * (-0.1 - 0.4) / r, abs_tol=1e-12)
+    assert math.isclose(float(grad.zeta[0, 0]), gamma * e * (0.3 - (-0.2)) / r, abs_tol=1e-12)
+    assert math.isclose(float(grad.zeta[0, 1]), gamma * e * (-0.1 - 0.4) / r, abs_tol=1e-12)
+
+
+def test_neg_loglik_and_grad_matches_independent_reference():
+    # Multi-person / multi-item / multi-trait with a non-trivial factor_id and
+    # the DEFAULT (non-zero) penalty. Pins the objective and every gradient
+    # block/axis to a pure-Python reference derived from the canonical equation.
+    params = MLSIRMParams(
+        theta=np.array([[0.2, -0.3], [0.5, 0.1], [-0.4, 0.25]]),
+        alpha=np.array([0.1, -0.2, 0.05, 0.3]),
+        b=np.array([0.3, -0.1, 0.2, -0.25]),
+        xi=np.array([[0.1, 0.2], [-0.2, 0.4], [0.35, -0.15]]),
+        zeta=np.array([[0.0, -0.1], [0.3, -0.4], [-0.25, 0.2], [0.15, 0.05]]),
+        tau=0.2,
+    )
+    y = np.array(
+        [
+            [1.0, 0.0, 1.0, 1.0],
+            [0.0, 1.0, 0.0, 1.0],
+            [1.0, 1.0, 0.0, 0.0],
+        ]
+    )
+    factors = np.array([0, 1, 0, 1])
+    config = FitConfig(max_iter=1)  # default penalty is non-zero
+
+    nll, grad, loglik = neg_loglik_and_grad(y, factors, params, config)
+    ref_nll, ref_grad = _reference_neg_loglik_and_grad(y, factors, params, config)
+
+    assert math.isclose(nll, ref_nll, rel_tol=0.0, abs_tol=1e-10)
+    assert math.isclose(loglik, -(ref_nll - _reference_penalty_only(params, config)), abs_tol=1e-10)
+    assert np.allclose(grad.theta, ref_grad["theta"], atol=1e-12)
+    assert np.allclose(grad.alpha, ref_grad["alpha"], atol=1e-12)
+    assert np.allclose(grad.b, ref_grad["b"], atol=1e-12)
+    assert np.allclose(grad.xi, ref_grad["xi"], atol=1e-12)
+    assert np.allclose(grad.zeta, ref_grad["zeta"], atol=1e-12)
+    assert math.isclose(float(grad.tau), ref_grad["tau"], abs_tol=1e-12)
+
+
+def _reference_penalty_only(params, config) -> float:
+    free_alpha, uses_space = model_flags(config.normalized_model())
+    penalty = config.penalty
+    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))
+    if free_alpha:
+        delta = params.alpha - penalty.mu_alpha
+        value += 0.5 * penalty.lambda_alpha * float(np.vdot(delta, delta))
+    if uses_space:
+        value += 0.5 * penalty.lambda_xi * float(np.vdot(params.xi, params.xi))
+        value += 0.5 * penalty.lambda_zeta * float(np.vdot(params.zeta, params.zeta))
+        value += 0.5 * penalty.lambda_tau * float((params.tau - penalty.mu_tau) ** 2)
+    return value
 
 
 def test_missing_entries_are_excluded():