Add GaussianizedFlowMatching estimator (flow-matching SBI with global ZCA whitener)#112
Open
cweniger wants to merge 15 commits into
Open
Add GaussianizedFlowMatching estimator (flow-matching SBI with global ZCA whitener)#112cweniger wants to merge 15 commits into
cweniger wants to merge 15 commits into
Conversation
A new falcon estimator: a Gaussian whitening preconditioner (conditional + unconditional, with prior-width covariance clamping) feeding two flow-matching residual models, with importance sampling adapted to standard-normal latent space. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
A Gaussian-preconditioned flow-matching posterior estimator: a learned conditional + unconditional Gaussian (GaussianFullCov-style MLP mean, EMA full covariance with prior-width clamping) whitens the prior's standard-normal latent space, and two conditional flow-matching velocity fields learn the residual non-Gaussian structure in the ~N(0,I) whitened space. Sampling and density feed Flow's importance-sampling machinery, adapted to standard-normal latent space: the latent prior N(0,I) is added analytically to the weights (replacing the hypercube penalty) and a prior-width truncation backstops runaway proposals. New: - estimators/flow_matching.py: velocity field, EMA, Euler sampling, CNF density (exact + Hutchinson divergence, chunked). - estimators/gaussianized_flow_matching.py: the estimator + _GaussianizedFlow. - examples/06_gaussianized_flow_matching: moderate 5D bimodal demo. - examples/02_bimodal/config_gfm.yml: extreme-scale stress config. Status: implemented, registered, integrated; trains and samples end-to-end. CNF density validated (normalizes to 1); coarse multimodal structure recovered. Known limitation: under-sharpens within-mode width ~2-4x because mu(x) prediction error smears the whitened target; fails on extreme needle posteriors. See plans/GAUSSIANIZED_FLOW_MATCHING.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Replace the conditional Gaussian mu-MLP -- which over-fit the finite buffer
per-observation (memorising which mode each x maps to, collapsing the posterior
to unimodal) -- with a single global EMA mean/covariance whitener, and fold the
conditional + marginal flows into ONE flow-matching network selected by a
learnable NULL token (classifier-free style). Fits dynamic SBI: no per-x
estimator to over-fit, and the global whitener zooms as the buffer concentrates.
Whiten with the symmetric (ZCA) square root Sigma^{-1/2} = V diag(lambda^{-1/2}) V^T
rather than PCA (.) @ V / sqrt(lambda): under a near-isotropic buffer the
eigenvectors are arbitrary and jitter/sign-flip every step, spinning the flow's
target into a rotation-averaged blob (right variance, no modes; sometimes NaN).
The symmetric root is a function of Sigma alone, so the jitter cancels and the
whitened target stays in the original coordinate frame. Log-det unchanged.
The 5D bimodal example now recovers a clean, fully-resolved 32-mode posterior
(modes +/-0.5, width ~0.13 vs true 0.1, balanced); FM loss 1.36 -> 0.62.
Also: add examples/06.../analyze.py (corner plot + per-dim recovery), clean the
example config (drop obsolete Gaussian-MLP keys, use_best_models on), and set
x_obs to zeros for a symmetric bimodal target.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…not FM
The whitened FM loss is ~scale-invariant, so it stays the same order of magnitude
and is a poor checkpoint/early-stop criterion. Add a proper held-out NLL in
theta_lat space as the selection metric:
log p(theta_lat | c) = cnf_logprob(whiten(theta_lat), c) + whitener.logdet()
whose -1/2 sum log lambda term is exactly the compression/log-volume contribution
(the analogue of the Gaussian NLL's 1/2 log det Sigma), so the metric falls as the
posterior genuinely sharpens. val_step now reports this as "loss" (drives both the
base early-stopping and on_epoch_end best-model selection); FM stays the optimized
training loss and is still logged on both sides for like-for-like comparison.
whiten_logdet is logged so the zoom is directly visible.
The selection density only needs to *rank* checkpoints, so val_nll uses cheap
Hutchinson divergence (val_n_probe=1, noise averages over the val set) and few steps
(val_density_steps=16) -- ~9x cheaper than the exact/high-step IS density, ~28ms vs
253ms per call. Importance sampling still uses the accurate density.
Example: switch to a dynamic-SBI config (simulate_count, concentrating buffer).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The velocity field concatenated the summary s = E(x) raw with [w, time_embed], with no input normalization -- unlike GaussianFullCov, which whitens its conditioning via EMA _input_mean/_input_std before the mean MLP. An offset or odd scale in s then swamps the unit-scale w and buries the fine x-dependence that pins down a zoomed posterior, so significant zooming degrades. Add the same EMA diagonal whitening of the summary in _WhitenedFlow (_cond_mean/_cond_std, updated in training only) and apply it to the conditional branch in training_loss / val_nll / importance sampling / discard mask. The null (marginal) token stays learnable and now lives in this normalized conditioning space (its zero init = the mean of normalized s). Stats start at (0,1) so it is identity until they populate; frozen outside training. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- gaussianized_flow_matching: add grad_clip (clip_grad_norm_ between backward/step, default 1.0) and a layernorm toggle threaded through to VelocityField; drop dead hidden_dim/num_layers/flow_enabled kwargs. - flow_matching: LayerNorm on hidden layers only (input cat + output Linear left unnormalized); merge the two divergence helpers into one _vel_and_div; strip debug instrumentation from fm_loss. - example 06: config/model update for the current dynamic-SBI zoom test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JYDMB1wbxQqv9kUwWNBdNz
- flow_matching.fm_loss: optional antithetic sampling -- mirror w0->-w0 and t->1-t (4 variants/target), an unbiased ~40%-lower-variance estimator of the same FM objective. Toggle threaded through as `antithetic` (default True). - example 06 config: current well-behaved dynamic-SBI zoom settings (lr 5e-4, antithetic on by default, grad_clip/layernorm off). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JYDMB1wbxQqv9kUwWNBdNz
…trace plot - estimator: add tight_proposal_mode (log_target = gamma*(cond-marg)+prior) and optional periodic shrink-and-perturb of the live velocity head (EMA/deployed model, whitener, embedding, null token left untouched). - example config: retune (max_epochs 1000, gamma 1.0, tiny whitener momentum, betas [0.5,0.5], proposal_mixture_beta 0.9, tight_proposal_mode on). - example model: tighten SIGMA/SHIFT toward a near-delta posterior. - analyze.py: add buffer z0-over-training trace plot (ported from buffer.ipynb). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…a<1) Adds proposal-only truncation to stabilize importance-sampled proposals when gamma<1. Inside _importance_sample we EMA-estimate a scalar floor _logp_cond_floor = the truncation_alpha-quantile of log_prob_cond under the posterior (pooled weighted quantile of the mixture candidates -> EMA; no loss), using the UNtruncated posterior weights so the estimate stays unbiased. Proposal candidates with log_prob_cond <= floor get the out-of-bounds penalty -> ~zero resample weight. Posterior sampling is never truncated. Why it works: self-normalized-IS weight variance is dominated by the max weight, which sits in the flow's poorly-modeled low-density tail; gamma<1 broadens the target toward the prior and pumps weight into exactly that tail. A tiny alpha (cuts ~0.1% of posterior mass) removes a large share of weight variance (mass-tail != weight-tail) -> outsized stabilization. Density-domain truncated importance sampling (cf. Ionides 2008 / PSIS). - new kwargs: truncation_alpha (default 1e-6 ~5sigma; 0 disables, bit-identical), truncation_momentum (0.05), truncation_warmup_epochs (0) - persisted as a python float in logp_cond_floor.pth (existence-checked load, back-compatible); logged as logp_cond_floor - guardrails: asserts a single shared observation per call; relies on proposal_mixture_beta < 1 to keep floor estimation honest - example config: gamma=0.2 + truncation_alpha=1e-3, max_epochs=3000, lr=1e-3 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…t obs - raystore: buffer/n_training now reports the actual trainable pool (TRAINING + DISFAVOURED), matching the dataloader (cached_loader) and the min_samples floor in rotate_sample_buffer. It previously counted only TRAINING status, so it spuriously read below min_samples once samples were disfavoured (which are still trained on until tombstoned). n_disfavoured is kept as the eviction-pending subset. - GFM _compute_discard_mask: condition on the TARGET observation (stashed as _target_summary), not each buffer sample's own simulated x_i. Matched (z_i, x_i) pairs are posterior-typical, so log_prob_cond was always well above the floor and nothing was ever discarded; conditioning on x_obs makes the floor test meaningful and evicts old broad-proposal samples. Evaluate under _best_flow (consistent with how the floor is estimated); add a stochastic p-gate to amortize the CNF-density cost. - drop the dead log_ratio_threshold knob (superseded by the floor). - example: tighter posterior (SIGMA/SHIFT), gamma=0.2 + truncation/discard config. Note: WIP left in place -- floor-weight formula is under experiment (marginal/prior terms commented out) and an unused `info` import remains. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Mark the floor-weight line (cond-only vs full posterior weights) as an open experiment, and remove the now-unused `info` import left over from debug traces. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #112 +/- ##
=========================================
- Coverage 10.33% 9.08% -1.26%
=========================================
Files 30 32 +2
Lines 3927 4468 +541
=========================================
Hits 406 406
- Misses 3521 4062 +541
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
VelocityField gains two optional input transforms, both off by default (0 = bit-identical to previous behavior): - w_embed_dim > 0: linear up-projection of w before the input concat, so a low-dim parameter is not variance-shadowed by a high-dim summary. - cond_embed_dim > 0: small MLP bottleneck (Linear-SiLU-Linear) that compresses the conditioning summary; applied to the null token too, keeping conditional and marginal branches consistent. Both are threaded through _WhitenedFlow and GaussianizedFlowMatching. Also comments out the single-shared-observation assert in the proposal truncation: evidence re-simulated per sample from a broadcast observation (e.g. tokens derived from the observed x) arrives as identical rows, and obs_batch == num_samples makes the expand a no-op. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… field per_param_nets=True replaces the single MLP trunk with one independent MLP per parameter, each seeing the full (w, t, s) input and predicting one velocity component — decouples per-component FM-loss gradients so hard components cannot crowd out easy ones. Input projections stay shared. focus_dims=[i, ...] (debugging) restricts the flow to a subset of whitened coordinates: non-focus coords are hidden from the input and zeroed in the output velocity, and fm_loss averages over focus components only. The masked-loss minimizer is the exact marginal velocity field for the focus coords (nuisances marginalized over full simulation variance); masked coords stay at their N(0,1) base draw, i.e. the whitener's Gaussian buffer approximation, keeping sampling/density/importance weights consistent. Both default off (bit-identical behavior, state-dict compatible). The 11D focus_dims test (Mc alone) confirms network capacity is sufficient when trained on a single component — pointing at cross-parameter interference, not capacity, as the joint-training failure mode. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…rows cond_per_param=True (requires per_param_nets) splits the conditioning summary into param_dim equal slices and routes slice i exclusively to velocity leg i — pair with an embedding emitting one block per parameter (TransformerEmbedding query pooling, n_queries=param_dim) for a private readout->leg pipeline per parameter after the shared embedding body. Guardrails: incompatible with cond_embed_dim; cond_dim must be divisible by param_dim. _importance_sample now collapses per-sample duplicated condition rows (evidence re-simulated from a broadcast observation, e.g. tokens from the shared observed x) to batch 1. This makes the shared-observation truncation assumption hold, stores _target_summary single-row — fixing a crash in the discard path (expand of a 16-row stash against batch 256 when discard_samples=true) — and cuts the proposal-embedding forward by the chunk factor. _compute_discard_mask additionally guards with [:1]. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a new SBI estimator,
GaussianizedFlowMatching, plus its example and a small buffer-metrics fix. The design is deliberately simple and dynamic-SBI-friendly:_GlobalWhitener): a single EMA-tracked mean/cov of the standard-normal latent, whitened with the symmetric rootSigma^{-1/2}(eigenvalues clamped to the prior variance). No per-observation MLP, so nothing to over-fit; in sequential SBI the buffer concentrates each round and the whitener zooms automatically._WhitenedFlow): one velocity field used both conditionally (on the embeddings = E(x)) and unconditionally (learnable NULL token). Conditional density = posterior proposal, null density = marginal. Trains like MSE regression — no invertibility/Jacobian fragility — and feeds the same importance-sampling machinery asFlow.Notable features
truncation_alpha-quantile floor onlog_prob_cond; proposal candidates below it are penalised so the proposal never emits from the flow's poorly-modelled tail. Stabilisesgamma < 1proposals (self-normalized-IS weight variance is dominated by the extreme-density tail).discard_samplesevicts buffer samples in the target-posterior tail, conditioned on the target observationx_obs(not each sample's own simulatedx_i), using_best_flowfor consistency with the floor.buffer/n_trainingfix (raystore.py): now reports the actual trainable pool (TRAINING + DISFAVOURED) that the dataloader draws and thatmin_samplesfloors — previously countedTRAINING-only and spuriously read belowmin_samples.Files
src/falcon/estimators/gaussianized_flow_matching.py,flow_matching.py,estimators/__init__.pyexamples/06_gaussianized_flow_matching/(config, model, analyze.py, data) andexamples/02_bimodal/config_gfm.ymlsrc/falcon/core/raystore.py(n_training accounting)plans/GAUSSIANIZED_FLOW_MATCHING.mdTest plan
alpha=0bit-identical no-op; warmup gate; posterior-mode safety; all-below self-heal).falcon launchruns clean with truncation + discard active; floor estimated, logged, and checkpoint round-trips.Open items (WIP, noted in code)
# TODO: figure out what works best— cond-only vs full posterior weights).🤖 Generated with Claude Code