Skip to content

Fit shared latent covariance models with joint EM - #30

Merged
jamgochiana merged 10 commits into
mainfrom
feat/joint-latent-em
Oct 29, 2026
Merged

jamgochiana merged 10 commits into
mainfrom
feat/joint-latent-em

Conversation

@jamgochiana

@jamgochiana jamgochiana commented Oct 1, 2026 •

Copy link
Copy Markdown
Collaborator

Fit LatentCovariance with standard latent-variable CovarianceEM, either as a weighted single Gaussian or inside EM with exactly Tied(:F,:D). Mixture components learn shared loadings F and residual variances D, their own full or diagonal A_k, and unrestricted means: Σ_k=F A_k F′+Diagonal(D). PCAEM remains the separate fixed-projection pipeline.

For fixed responsibilities and updated means, a small posterior-moment helper uses latent-space Cholesky solves. The joint M-step pools mass-weighted cross/second moments to update F, updates each A_k, and pools expected squared residuals to update D. variance_floor constrains D. One joint inner report describes this covariance optimization.

The implementation follows the direct matrix-equation style of existing LRD CovarianceEM, reuses _variance and _objective, and shares one core between single-Gaussian and mixture fitting. It uses the existing public workspaces and outer EM dispatch. No custom buffer struct or optimizer framework is introduced. fit! continues from current parameters; docs explain initialization, identifiability, local convergence and the inner sweep budget.

Workspace validation checks F/D ties, latent shape and positive-definite A. Zero responsibilities are omitted before arithmetic, and component arrays remain independently owned. Includes the shared empty-observation mixture prediction fix also in #31. Other partial-sharing patterns remain deferred. No new dependencies or version bump.

Validation

  • 816 assertions pass locally, including independent dense posterior/one-sweep references, full/diagonal A, weighted replication, one-component equivalence, continuation, monotonicity, variance floors, zero-responsibility extreme values, and high loading/noise likelihood against a dense BigFloat reference.
  • Documentation builds with executed examples; README and fitting guide describe the method and limitations.
  • Fresh GitHub Julia 1.11, prerelease and documentation checks pass on b363efd.
  • A local simulation of squash-merging Add banded precision Gaussians and exact fitting #31 followed by Fit shared latent covariance models with joint EM #30 combines cleanly: 1,813 assertions pass, and combined docs build.

Performance follow-up

Reusable buffers, blocked accumulation and shared LRD/latent posterior kernels are deferred to #34, which preserves the prototype and its benchmark evidence. This PR retains sample-sized intermediate arrays and centers one component at a time.

A cached MNIST check (10k train / 10k validation, p=784, K=10, rank 30, seed 1, Gaussian dequantization σ=0.01) uses three outer iterations and three inner sweeps, Julia 1.11.5 / BLAS1, and two fresh fits after a small dispatch warm-up. Median fit time is 17.23 s, cumulative allocations 60.69 GB, and mean validation log likelihood 966.2014811920. Covariances, means, weights, training history and per-image validation scores match the original direct implementation exactly. Timings include initialization and in-fit GC; loading, validation scoring and I/O are excluded. Cumulative allocations are not peak RAM; this checks a fixed iteration budget, not convergence. Benchmark assets remain outside the repository.

@jamgochiana
jamgochiana changed the base branch from main to fix/fitting-cleanup October 1, 2026 21:32
Base automatically changed from fix/fitting-cleanup to main October 7, 2026 05:29
@jamgochiana
jamgochiana merged commit 17fdad3 into main Oct 29, 2026
4 checks passed
@jamgochiana
jamgochiana deleted the feat/joint-latent-em branch October 29, 2026 23:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant