Configurable RMS computation timing within CAVI cycle

[jc-macdonald]

Summary

RMS (reconstruction error) is currently computed at a single fixed point in the coordinate-ascent cycle: after the scores and loadings updates, but before the noise variance update, using the new A and S but the old μ. This asymmetry is the root cause of the period-2 oscillation documented in _rms_oscillation_trace.py.

This issue adds a configurable option to control where in the CAVI cycle RMS is evaluated, enabling systematic comparison in Tier 1 of the convergence design matrix (Factor 5).

---

Current behavior (T0)

RMS is computed in _converge.py::_log_and_check_convergence(), which is called from the main iteration loop after scores + loadings + bias + noise updates. However, the RMS value passed to this function is computed earlier — after S and A are updated but before the noise variance V is updated (see _full_update.py). This means:

• S: current iteration (new)
• A: current iteration (new)
• μ: current iteration (new, but computed from previous iteration's err_mx)
• V: previous iteration (old)

The μ update uses err_mx which was computed with the previous iteration's μ, creating a one-step lag that generates the 2-cycle.

---

Proposed timing variants

ID	When RMS is computed	What's new vs. old	Notes
T0	After S+A+μ, before V (current)	S,A,μ new; V old	Asymmetric — causes oscillation
T1	After full cycle (S, A, μ, V)	All parameters at iteration-end	Symmetric — likely eliminates oscillation
T2	Average of pre- and post-cycle	Period-2 average	Equivalent to smoothed RMS but computed differently
T3	After A only (before μ and V)	S,A new; μ,V old	Isolates loadings convergence

---

What to add

• New option rms_timing: str with values "default" (T0), "end_of_cycle" (T1), "average" (T2), "after_loadings" (T3). Default: "default" (no behavior change).
• Modify the RMS computation call site in _full_update.py (or the main loop in _pca_full.py) to compute RMS at the selected point.
• For T1: move/defer the RMS call to after the noise variance update.
• For T2: compute RMS both before and after the full update; store the average.
• For T3: compute RMS after loadings but before bias.
• Store the timing variant used in lc["_rms_timing"] for trace provenance.

---

Estimated scope

• ~80–120 lines in _full_update.py / _pca_full.py (restructure the RMS call sites)
• ~20 lines in options defaults
• ~150–200 lines of tests (verify each timing variant produces different traces on the known-oscillating dense case, verify T1 eliminates oscillation)
• Total: ~250–350 lines

Design matrix references

• Factor 5: RMS computation timing (T0–T3)
• Tier 1: Oscillation root cause (24 runs)

Notes

• This is independent of Convergence criterion overhaul: relative ELBO, composite criterion, patience window #54 (convergence criterion overhaul) — Convergence criterion overhaul: relative ELBO, composite criterion, patience window #54 adds new criteria for checking convergence; this changes when the RMS value itself is computed.
• Pre-centering (R1, R2 in Factor 7) already eliminates oscillation by a different mechanism. This issue investigates whether timing alone is sufficient.
• Results from this issue inform whether the default RMS computation point should be changed permanently (a separate decision).

[jc-macdonald]

Root cause identified

The period-2 RMS oscillation is caused by a one-iteration lag in the bias (mean) update:

1. _bias_and_center computes d_mu from err_mx of the previous iteration
2. It re-centers the data with this increment
3. Scores and loadings re-fit to the new centering
4. The new err_mx reflects the updated centering — but the gradient direction can flip because A and S were fitted to the old geometry
5. Next iteration reverses the mean correction → stable 2-cycle

This is a coordinate descent coupling issue between the mean update and the (A, S) updates.

Documented workaround

Center your data before fitting. Use MissingAwareStandardScaler (or AutoEncoder) as a preprocessing step:

from vbpca_py import VBPCA, MissingAwareStandardScaler

scaler = MissingAwareStandardScaler()
x_scaled = scaler.fit_transform(x, mask=mask)
model = VBPCA(n_components=5, maxiters=100)
model.fit(x_scaled, mask=mask)

Pre-centered data eliminates the oscillation entirely, even with bias=True. This is confirmed by the existing diagnostic scripts (_rms_oscillation_trace.py, _convergence_analysis.py).

A note has been added to the Known limitations section of the README on main (commit e3814a3).

Possible future fixes

• Auto-precenter internally (decouple mu from the iterative loop)
• Mu damping (under-relax the mean update)
• RMS smoothing (2-point moving average for convergence detection only)