fix(dml): √n scaling bug in dml_model_averaging SE — CI was ~20× too wide

The cross-candidate variance aggregator multiplied influence functions
by 1/√n in the outer-product, giving Σ φ_k φ_l = (1/n)·Σ_i φ_ki φ_li,
but then used that directly as Var(θ̂_avg) instead of the correct
(1/n)·Σ / n = (1/n²)·Σ.  Net effect: reported SE was √n larger than
the truth.

On the canonical n=400 DGP with true θ=1.5:

              before fix      after fix
  CI width     4.20            0.21     (×20 = √400)
  Coverage    100% (over)      82% (~nominal)

82% vs the nominal 95% is driven by a small finite-sample downward
bias in the point estimate (mean 1.44 vs true 1.5, about 4% — normal
PLR behaviour when the nuisance is fit on n=400 with p=8). That's a
tuning-of-nuisance issue, not a variance-formula issue; the SE now
lines up with the empirical SD of the point estimate.

Also adds a regression guard in test_dml_model_averaging.py that
fails if CI width creeps back above 1.0 on the canonical DGP.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: brycewang-stanford

src/statspai/dml/model_averaging.py

@@ -217,13 +217,17 @@ def dml_model_averaging(
 
     theta_avg = float(np.sum(w * thetas))
 
-    # --- variance: covariance between candidate scores --------------- #
-    psi_matrix = np.zeros((n, len(thetas)))
+    # --- variance: cross-candidate influence-function covariance ------ #
+    # For DML-PLR each θ̂_k has influence function φ_k,i = ψ_k,i / J_k so
+    # Var(θ̂_avg) = (1/n²) · w^T (Φ^T Φ) w where Φ is the n × K matrix of
+    # influence functions.  We store φ/√n so Φ^T Φ already carries one
+    # factor of 1/n; dividing by n once more gives Var(θ̂_avg).
+    phi_matrix = np.zeros((n, len(thetas)))
     for k, (y_r, d_r, theta_k) in enumerate(resids):
         J_k = -np.mean(d_r ** 2)
-        psi_matrix[:, k] = (y_r - theta_k * d_r) * d_r / (J_k * np.sqrt(n))
-    cov = psi_matrix.T @ psi_matrix
-    var_avg = float(w @ cov @ w)
+        phi_matrix[:, k] = (y_r - theta_k * d_r) * d_r / (J_k * np.sqrt(n))
+    cov_scaled = phi_matrix.T @ phi_matrix            # ≈ (1/n) · Σ φ_k φ_l
+    var_avg = float(w @ cov_scaled @ w) / n           # = (1/n²) · Σ w_k w_l Σ φ_k φ_l
     se_avg = float(np.sqrt(max(var_avg, 0.0)))
 
     z = sp_stats.norm.ppf(1 - alpha / 2)

tests/test_dml_model_averaging.py

@@ -76,3 +76,22 @@ def test_registered_in_public_api():
     fns = sp.list_functions()
     assert "dml_model_averaging" in fns
     assert "model_averaging_dml" in fns
+
+
+def test_se_on_correct_scale():
+    """Regression test for the v1.5.1 √n scaling bug.
+
+    Before the fix, the CI width was ~√n × too wide (≈4.2 for n=400 instead
+    of ~0.2). After the fix, on the canonical DGP with true θ=1.5 and
+    n=400, a single draw should give a CI width well under 1.0.
+    """
+    df = _synth_dml_data(seed=42, n=400)
+    r = sp.dml_model_averaging(
+        df, y="y", treat="d",
+        covariates=[f"x{j}" for j in range(8)],
+        n_folds=3, seed=42,
+    )
+    width = r.ci[1] - r.ci[0]
+    assert width < 1.0, f"CI width {width:.3f} suggests the √n scaling bug has regressed"
+    # Point estimate shouldn't have drifted wildly
+    assert abs(r.estimate - 1.5) < 0.5