Add posterior function to estimate MLE/MAP elbo, responsibility, and smoothed states - #161
Add posterior function to estimate MLE/MAP elbo, responsibility, and smoothed states#161harrisonritz wants to merge 15 commits into
posterior function to estimate MLE/MAP elbo, responsibility, and smoothed states#161Conversation
Add a public `infer_γ` (exported) in `fit_SLDS.jl` that infers the discrete-state responsibilities γₜ(k) = q(zₜ = k) of a fitted SLDS with parameters held fixed (no M-step). It runs the structured-variational E-step, alternating the forward-backward pass over the switching chain q(z) with the Laplace/Kalman smoother over the continuous states q(x) — the classic coordinate-ascent E-step of Ghahramani & Hinton — either until the responsibilities converge (max|Δγ| < tol, capped by max_iter) or for a fixed number of iterations (check_convergence flag). The two smoothers are coupled deterministically: the discrete-layer log-likelihoods are plugged in at the smoothed posterior mean rather than a Monte-Carlo draw, so unlike fit!'s learning-time E-step the iteration is reproducible and γ settles to a fixed point, making the convergence check well behaved. Returns γ matching the input shape family: a K×T matrix for single-trial matrix input, otherwise a Vector of per-trial K×T_i matrices (ragged lengths allowed). Includes tests (shapes, determinism, both stopping modes, K=1 degenerate case, and discrete-state recovery on well-separated regimes) and a docs section.
…es when requested
…mprove numerical stability
|
I'll do my best to review this today, thought it's a big diff so may take me a day or two :) thanks! |
Benchmark Results (Julia v1)Time benchmarks
Memory benchmarks
|
There's no way that return was correct, surely it has to loop over all trials before returning. must have been a mistake, talk about stochastic update
Remove evidence of hacks that didn't work
|
i will review this by EOD |
wtf is this lol |
| @skip_test test_SLDS_posterior_basic() | ||
| @skip_test test_SLDS_posterior_shapes() | ||
| @skip_test test_SLDS_posterior_deterministic_and_modes() | ||
| @skip_test test_SLDS_posterior_K1() | ||
| @skip_test test_SLDS_posterior_recovers_distinct_regimes() |
There was a problem hiding this comment.
probably wanna unskip these?
| z, x, y = rand(rng, slds, fill(tsteps, ntrials)) | ||
|
|
||
| # Single-trial matrix ⇒ bare K × T matrix; wrapping in a vector ⇒ 1-element vector. | ||
| γ_mat = infer_γ(slds, y[1]) |
There was a problem hiding this comment.
many instances of this stale function handle
| the smoothed posterior mean), so the returned responsibilities are reproducible. | ||
|
|
||
| ```@docs | ||
| infer_γ |
| end | ||
|
|
||
| # Core SLDS trial sampling logic | ||
| # Core SLDS trial sampling logic. `ux_trial` / `uy_trial` are the canonicalized |
There was a problem hiding this comment.
seems liek this last predicate was meant to be deleted?
| `ws.opt.grad_buf` and returns it. `ux` (dynamics input, feeds `-Q⁻¹` / | ||
| `A'Q⁻¹` residuals via `Bₖ u`) and `uy` (observation input, feeds the emission | ||
| gradient via `Dₖ v`) are per-trial matrices; `nothing` or zero-row skips them. |
There was a problem hiding this comment.
I'm assuming claude added these types of statements. Probably can mention how inputs are done once and don't need to mention again unless the function is specific to them or has some behavior meant to truly notice
| """ | ||
| posterior(slds, y; max_iter=100, tol=1e-6, check_convergence=true, progress=false) | ||
|
|
||
| Infer the discrete-state responsibilities `γₜ(k) = q(zₜ = k) ≈ p(zₜ = k ∣ y₁:T)` | ||
| of a **fitted** `SLDS`, holding all model parameters fixed. | ||
|
|
||
| This is a post-fit *inference* routine: it estimates the discrete posterior for a | ||
| model whose parameters are already learned, and does **not** update them (there is | ||
| no M-step). | ||
|
|
||
| It runs the structured-variational E-step, alternating between: | ||
|
|
||
| 1. **discrete smoothing** — a forward-backward pass over the switching chain that | ||
| refreshes `q(z)` from the current per-regime log-likelihoods, and | ||
| 2. **continuous smoothing** — the Laplace/Kalman smoother that refreshes `q(x)` | ||
| under the updated responsibilities `γ`, | ||
|
|
||
| until either the responsibilities converge or for a fixed number of iterations. | ||
|
|
||
| # Arguments | ||
| - `y`: observations — a `(obs_dim, T)` matrix (single trial), a | ||
| `(obs_dim, T, ntrials)` array, or a `Vector` of per-trial `(obs_dim, T_i)` | ||
| matrices (ragged trial lengths allowed). | ||
|
|
||
| # Keywords | ||
| - `max_iter::Int=100`: maximum number of E-step (forward-backward + smoother) | ||
| iterations. | ||
| - `tol::Real=1e-6`: convergence tolerance on `max|Δγ|`, the largest absolute change | ||
| in any responsibility between successive iterations (used only when | ||
| `check_convergence=true`). | ||
| - `check_convergence::Bool=true`: the mode flag. `true` iterates *until | ||
| convergence* (`max|Δγ| < tol`), capped at `max_iter` iterations; `false` runs | ||
| *exactly* `max_iter` iterations. | ||
| - `progress::Bool=false`: show a progress bar. | ||
|
|
||
| # Returns | ||
| The responsibilities `γ`: | ||
| - for a single-trial matrix `y`, a `K × T` matrix; and | ||
| - otherwise (3-D array or vector-of-matrices input), a `Vector` of per-trial | ||
| `K × T_i` matrices, one per trial. | ||
|
|
||
| Each column is a probability vector over the `K` discrete states (sums to 1). | ||
| """ |
| function posterior( | ||
| slds::SLDS{T,S,O}, | ||
| data::Data{T}; | ||
| return_γ::Bool=true, |
There was a problem hiding this comment.
if someone is calling posterior--is there really a circumstance in which they dont want gamma? I mean to an extent what sets this apart from smooth is really just the gammas
|
|
||
| # Returns | ||
| The responsibilities `γ`: | ||
| - for a single-trial matrix `y`, a `K × T` matrix; and |
There was a problem hiding this comment.
i don't think this is accurate
| return_x::Bool=false, | ||
| max_iter::Int=100, | ||
| tol::Real=1e-6, | ||
| check_convergence::Bool=true, |
There was a problem hiding this comment.
this also seems a bit unneccesary. I mean we probably should just check convergence and if we don't ust throw a warnign and return given the max_iters
| return_elbo::Bool=true, | ||
| return_x::Bool=false, |
There was a problem hiding this comment.
I think there is a bit of semantic dubiousness going on. we noe effecticely have two ways to get elbo right? so why should this function privilege the "good elbo" ? I wonder if this function need to be integrated more into the core code somewhat like in ssm -- e.g., the same function does the posterior, but depending on the calls, we wi8ther iterate until convergence or do it once etc.
|
I think overall the functions are well written mopdulo some correctness issues. I do think the semantics need to be re-evaluated |
|
Ok, so maybe And then The only hiccup is that you might only want to run the back-and-forth estimator once -- because it might take many iterations to converge -- so nice to have it spit out the posterior and the elbo from the same fit. Maybe some caching would help here. |
yeah that should work. In theory the
do you mean gamma and x? also i don't mind too much if return the elbo as part of a named tuple in |
|
Ok, I'll take a look next week! Meant smooth should now additionally return gamma (and elbo, good idea). |
2 participants
Iterates discrete/continuous smoothing until convergence, then exports elbo, responsibilities, and states. Takes a while to run to completion, so ideally for post-fit eval.
want a better interface to choose what is exported -- maybe an
exportdict?export smoothed cov, but make it optional (big variable)
remove kluge regularization of discrete parameters
maybe we create a function to share with the estep? Eg could have option to run multiple smoothing iterations during estep (iirc you said 2-3 iter can give much (?) better updates). Less surface for estep revisions