GLM EM M-step uses Local-only marginal for shared-state emission, breaking EM monotonicity

[edeno]

Reported by: investigation triggered by Phase 2 PR #30, which surfaced this through the new _DetectorBase EM-monotonicity and final-E-step warnings.

Symptom

tests/test_em_monotonicity.py::TestEncodingUpdateGuards::test_sorted_spikes_glm_encoding_runs_end_to_end (NonLocalSortedSpikesDetector with sorted_spikes_algorithm="sorted_spikes_glm", max_iter=3) trips:

UserWarning: EM did not converge after max_iter=3 iterations (final log-likelihood change=-6.057e+00, tolerance=1.000e-04).
UserWarning: Final E-step log-likelihood decreased by 6.117e+00 (tolerance=1.000e-04). This indicates an inconsistency between the E-step and the M-step output...

Per-iteration marginal log-likelihoods (instrumented run on the simulated input from make_simulated_data(seed=42, n_neurons=5)):

iter	source	LL	Δ
0	initial E-step	-4124.96
1	after M-step 1	-4072.65	+52.31
2	after M-step 2	-4078.71	-6.06
3	final E-step	-4084.82	-6.12

The log-likelihood decreases monotonically after the first iteration, which violates the EM monotonicity guarantee. The KDE counterpart (test_sorted_spikes_kde_encoding_with_rollback) does NOT trip these warnings.

Root cause

Structural mismatch between the GLM M-step's objective and the marginal log-likelihood it is supposed to monotonically improve:

1. The M-step at src/non_local_detector/models/base.py:1907-1957 refits the GLM using only the Local-state marginal posterior as weights (acausal_state_probabilities[:, local_state_index], line 1908-1910) and the animal's actual position as the design-matrix input.

2. The fitted place_fields / coefficients are then used unchanged as the per-bin firing-rate emissions for every state — Local, Non-Local Continuous, Non-Local Fragmented, No-Spike — via self.encoding_model_[likelihood_name[:2]] at base.py:3967.

3. The HMM's expected complete-data log-likelihood for the place field θ in this model is:

   Σ_t Σ_{state,bin} γ_t(state,bin) · log p(spikes_t | bin, θ)

   using the joint acausal posterior over state × bin. The current M-step replaces that with:

   Σ_t γ_t(Local) · log p(spikes_t | pos_t, θ)

   — wrong weights AND wrong design points for the Non-Local components. This is not the EM M-step for this emission, so the monotonicity guarantee no longer holds.

Evidence ruling out alternative causes

• Not a BFGS-tolerance bug. Instrumenting fit_poisson_regression (sorted_spikes_glm.py:137-206) confirms BFGS does improve the local Q-objective on every M-step; the returned loss is strictly below the previous iter's coefficients' loss. Q under Local-only weighting is monotone — yet the marginal LL drops, the diagnostic signature that the optimized Q ≠ EM's Q.
• Not a weights-source confusion. Weights are read as acausal_state_probabilities[:, local_state_index] (the marginal of the Local state, shape (n_time,)), not from acausal_posterior (the joint, shape (n_time, n_bins)). Instrumentation confirms weight_sum shifts 55,742 → 74,040 → 84,344 across iters as Local-mass grows — the M-step is fitting a moving target with a one-state slice of the posterior.
• Not an is_training leakage bug. The mask is consistently applied at base.py:3727-3752 for both KDE and GLM.
• Why KDE doesn't trip the warnings. KDE's M-step is a closed-form weighted kernel density estimator. The same structural mismatch exists, but the KDE place-field perturbation between iters is small enough and well-conditioned enough that the marginal LL stays monotone within tolerance. The GLM's regularized spline-basis fit is much more sensitive: low-spike-count neurons get coefficients dragged by the L2 prior in ways that change rapidly when weight_sum shifts.

Recommended fix (design decisions needed)

Three options, in increasing order of invasiveness:

1. Aggregate over all states that share the emission. Pass acausal_state_probabilities[:, [Local, Non-Local-Continuous, Non-Local-Fragmented]].sum(-1) as the weight instead of the Local marginal alone. Still ignores the Non-Local states' design-point mismatch (they don't condition on the animal's actual position) but is a much smaller bias than the current one. Smallest change; possibly enough for a smoke-level fix.
2. Treat the GLM M-step as a partial M-step and add Q-rollback. Refuse the M-step if it doesn't increase the marginal LL on the smoother's posterior — fall back to the previous iter's coefficients. Restores monotonicity at the cost of slower EM. The KDE path already has a rollback guard (test_sorted_spikes_kde_encoding_with_rollback exists, see base.py encoding-update with rollback).
3. Redesign the M-step for shared emissions. Integrate over the joint posterior (state, bin) correctly for each state's emission contribution. Largest change; requires algorithmic redesign.

Secondary: warm-start BFGS from the previous iter's coefficients (currently cold-starts from [log(avg_rate), 0, …] at sorted_spikes_glm.py:182-187). Wouldn't fix the structural mismatch but would reduce iteration-to-iteration noise.

Recommended PR scope

For now, leave the warnings in place in PR #30 — they correctly surface a real algorithmic mismatch. Either:

• Option 1 above is a candidate for a small follow-up PR if smoothing the symptom is desired before tackling the redesign.
• Option 3 is the right long-term fix but is non-trivial work.

The test could be marked pytest.warns(UserWarning, match="did not converge|decreased") to acknowledge the known behavior, or xfail pending the M-step redesign.

Reproduction

cd non_local_detector
uv run pytest src/non_local_detector/tests/test_em_monotonicity.py::TestEncodingUpdateGuards::test_sorted_spikes_glm_encoding_runs_end_to_end -v -W default

The two UserWarnings appear in the captured warnings list.