Repository navigation
Conversation
The guidance methods (affine_classifier_free, universal, interval, prior_guide) took the score from the estimator's forward pass. For FMPE, forward returns the velocity, so guided FMPE sampling used the velocity as if it were the score. The guidance methods now call the estimator's score, which is the forward pass for score estimators and unchanged for NPSE.
All posteriors draw initial samples for map() with their sample() method, which shows a progress bar by default. map() did not pass its show_progress_bars on, so map(show_progress_bars=False), the default, still showed one. The base, ensemble and vector field posteriors now pass it on.
iid methods and guidance change the score, and only SDE sampling uses the score. With sample_with="ode", iid observations failed with a reshape error, and guidance was ignored without a warning: interval guidance returned unguided samples. Both now raise a ValueError that points to sample_with="sde".
FlowMatchingEstimator.score multiplied the velocity by a time of batch shape, so one time per sample failed for several dimensions and returned a (batch, batch) tensor for one dimension. It now broadcasts the time over the batch and event dimensions, as forward() does. Guidance reads the score since the previous commit, so it relies on this.
The error now says whether iid observations or guidance triggered it, says that both adjust the score that only SDE sampling uses, and points to sample_batched for independent observations. The docstrings of iid_method and guidance_method say that they need sample_with="sde". The test moves next to the other VectorFieldPosterior API test in base_test.py, and the neutral guidance test keeps the cases that can fail (VE and FMPE). The slow prior_guide test builds its FMPE posterior with sample_with="sde". FMPE samples with the ODE by default, so this test ignored the guidance before and now runs guided FMPE.
PriorGuide normalized the component weights with sum(dim=1). The sampler passes inputs of shape (samples, 1, dim), so dim 1 was the singleton axis and every weight became 1, which dropped the fitted mixture weights for K > 1. It now normalizes over the last dimension. A test checks that both input layouts give the same score.
FMPE samples with the ODE by default, which does not use the adjusted score of guidance or iid methods and now raises for them. The vector field tutorial and the how-to guide said FMPE supports both without saying how; they now name sample_with="sde".
VectorFieldPosterior.map() optimizes the potential of one observation. With a default x of several rows, it failed with a reshape error (ODE posterior, initial samples from the posterior) or an autograd error (every other case), on both samplers. It now raises a ValueError that names the number of rows. The docstrings of iid_method and guidance_method say that they need a posterior built with sample_with="sde", not the deprecated per-call argument, and the FMPE score docstring says that it takes one time for all samples or one per sample.
FMPE trains on theta_t = (1 - t) theta_0 + (t + sigma_min) eps. For this path the score is (-(1 - t) v - theta_t) / ((t + sigma_min) (1 + sigma_min)) and the squared diffusion is 2 (t + sigma_min) (1 + sigma_min) / (1 - t). Both lacked the factor (1 + sigma_min), so `score` was 0.1% too large. The two errors cancelled in the drift of the reverse SDE, which is why sampling was unaffected; consumers of `score` alone (guidance, iid methods, Langevin correctors) got the scaled score. The Tweedie denoiser (x + s^2 score) / m of guidance multiplies this error by up to 100 near t = 1. New tests check, for all vector field estimators, that the forward SDE reproduces the time marginals (dm/dt = f m and d(s^2)/dt = 2 f s^2 + g^2), and that the flow matching score equals the exact score when the velocity is exact. The docstrings state the noise path with sigma_min, and write g(t) instead of the undefined LaTeX macro \g(t).
mean_t_fn clamped time at 0.99 (m >= 0.01, so denoising a uniform prior does not divide by zero), while score and std_fn used the unclamped time. Above t = 0.99, score therefore described a different marginal than mean_t_fn and std_fn. With the corrected score, I + s^2 J in the jac_gauss precision is exactly zero at t = 1 and torch.inverse raises; before, the 0.1% score bias hid this as a negative precision. score, std_fn, mean_t_fn, drift_fn and diffusion_fn now all clamp at the class attribute effective_t_max (0.99), so above it they describe the marginal at 0.99. drift_fn and diffusion_fn already clamped 1 - t; the diffusion numerator now uses the clamped time too. ODE sampling and log_prob use forward and are unchanged. The exact score test now compares score with mean_t_fn and std_fn, also at t = 0.995 and 1.0, and the forward SDE test takes the cutoff from effective_t_max.
IntervalGuidance scaled its log-sigmoid walls by scale_factor / (m^2 s^2), with m = mean_t_fn and s = std_fn. The Simformer paper (Gloeckler et al. 2024, Sec. 3.4) and its reference code scale by 1 / s^2. The extra 1 / m^2 is 1 for VE, but about 150 for VP and subVP near t = 1 and 10^4 for FMPE, where m is clamped at 0.01. Samples stayed inside the interval, but far from the truncated posterior. A new test samples a Gaussian posterior N(0, 0.25) with its exact score and a lower bound at 0, and compares the samples with the half-normal. With the old scale, the mean was 0.39 to 0.45 posterior sd off for VP and subVP and 1.6 sd for FMPE, with an sd 21% to 174% too large; now every estimator is within 0.11 sd and 5%. VE is unchanged.
…idance IntervalGuidanceConfig compared lower_bound >= upper_bound inside an `and` chain, so tensor bounds with more than one element raised "Boolean value of Tensor with more than one value is ambiguous". The check now compares the bounds as given (Python floats as floats, so they keep their precision, and tensors elementwise) and raises its ValueError when any lower bound is not below its upper bound. The guidance already handled per-dimension bounds. The docstring now says that a bound can be a tensor, with inf for no bound, and a new test checks that infinite bounds guide like a masked-out dimension and that unordered tensor bounds raise.
interval_fn expanded the mask to the shape of its first input and stored the result in the config. A second call with another batch size then raised in expand, for example when an IntervalGuidance object is called directly, as its docstring shows. torch.where broadcasts the mask, so the expand is not needed. Sampling through a posterior was not affected, because it builds a new config for every call. A new test calls one guidance object with two batch sizes.
IntervalGuidance subtracted the configured bounds from the input and masked with the configured mask as given. Tensor bounds or a mask on the CPU therefore failed when the guidance ran on CUDA or MPS ("Expected all tensors to be on the same device"); only Python float bounds and 0-dim tensors worked. Both now move to the input's device and dtype at each call, so the guidance also keeps working after the potential is moved.
A gpu-marked test compares the guided score on the accelerator with the CPU result, for tensor bounds and a mask created on the CPU.
…g SDE functions drift_fn, diffusion_fn and mean_t_fn took an optional effective_t_max, while score and std_fn always use the class attribute. Passing it to one of them therefore mixed two marginals again, which the shared cutoff exists to prevent. No caller in sbi passes it. Passing it now raises a FutureWarning and still replaces the cutoff for that call. One helper, _clamp_time, clamps time for all five functions. The reasons for the cutoff (the SDE is singular at t = 1, and iid methods and guidance divide by m = 1 - t) move from the argument docs to the class attribute.
The comment on FlowMatchingEstimator.effective_t_max keeps the two reasons for the cutoff in two lines. The mean_t_fn docstring shows the kernel standard deviation with sigma_min, as std_fn does. Two test comments keep to their constraint and state the measured error of the truncated-posterior test.
JacCorrectedScoreFn inverts I + s^2 J, with J the Jacobian of the score. For a posterior much narrower than the noise, near the end of the diffusion this matrix is about m^2 Sigma / s^2, below the rounding error of its unit diagonal, and float32 can round it to exactly zero; torch.inverse then raised a singular-matrix error. With the corrected FMPE score this happened for a posterior sd of 0.01 with 50 SDE steps; before, the score error made the matrix slightly negative instead, so the call ran with wrong precisions. The precision now comes from an eigendecomposition, and eigenvalues whose magnitude is below 10 machine epsilons are set to that floor. Other eigenvalues, including clearly negative ones from trained networks, keep their value, so results elsewhere do not change. A float32 test with the exact score of a narrow Gaussian posterior checks that the iid score stays finite at the cutoff time.
The comment in loss() described the path without sigma_min, unlike the class docstring and the code below it.
…nce tests Two helpers replace repeated setup: one builds the default FMPE or NPSE estimator, one replaces an estimator's score by the exact score of N(x, var I). Its mean depends on x, so the score has one row per observation, as the iid methods expect; the float32 jac_gauss test now combines its three observations as it would with a network, and still fails without the eigenvalue floor.
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Flow Matching SDE and score calculations sbi/neural_nets/estimators/flowmatching_estimator.py, tests/vf_estimator_test.py |
Flow Matching now uses a shared effective time cutoff and updated noise-adjusted formulas for scores and SDE functions. Tests cover score behavior, cutoff warnings, and SDE identities. |
Guidance score and numerical handling sbi/inference/potentials/vector_field_adaptor.py, tests/score_samplers_test.py, tests/linearGaussian_vector_field_test.py |
Guidance adaptors use the estimator’s .score() method. Interval guidance updates bound validation, scaling, and tensor device and dtype handling. Prior mixture normalization and Jacobian covariance inversion also change. Tests cover guidance scores, sampling, and numerical behavior. |
Posterior sampling and MAP constraints sbi/inference/posteriors/vector_field_posterior.py, sbi/inference/posteriors/base_posterior.py, sbi/inference/posteriors/ensemble_posterior.py, docs/advanced_tutorials/19_vector_field_methods.ipynb, docs/how_to_guide/25_choosing_vector_field_options.ipynb, tests/base_test.py |
Vector-field posteriors reject ODE sampling with IID observations or guidance, and reject MAP with multiple default observations. Posterior-based MAP initialization forwards progress-bar settings. Documentation and tests describe and verify these behaviors. |
Priority:
Estimated code review effort: 3 (Moderate) | ~25 minutes
Change: Bug fix
Possibly related PRs
- sbi-dev/sbi#2106: Changes MAP initialization behavior and errors across posterior types, including
VectorFieldPosterior. - sbi-dev/sbi#2088: Addresses IID vector-field diffusion sampling behavior.
- sbi-dev/sbi#2098: Updates numerical stability for Jacobian-corrected IID score calculations.
Merge Risk: 
3d346
-
No actionable merge-blocking risk was established for the sampling and guidance changes; the PR is mergeable after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Title check The title clearly and concisely describes the primary fixes for guided and IID sampling of vector field posteriors. Description check The description is detailed and covers the changes, motivation, evidence, risks, and affected behavior. It does not include the repository checklist or an explicit issue-closing statement, but it is o… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Full details: Description check
-
Explanation
The description is detailed and covers the changes, motivation, evidence, risks, and affected behavior. It does not include the repository checklist or an explicit issue-closing statement, but it is otherwise complete and directly relevant.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 20 functions across 9 files. (2 skipped: 2 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Comment @coderabbitai help to get the list of available commands.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2111 +/- ##
==========================================
+ Coverage 90.24% 90.29% +0.05%
==========================================
Files 142 142
Lines 14845 14871 +26
==========================================
+ Hits 13397 13428 +31
+ Misses 1448 1443 -5
Flags with carried forward coverage won't be shown. Click here to find out more.
|
manuelgloeckler
left a comment
manuelgloeckler
left a comment
There was a problem hiding this comment.
Replying to your question: the m**2 was on the wrong side, but 1/std**2 is not the right answer either — it is incomplete.
The m**2
The guidance gradient is taken w.r.t. the detached denoised input (vector_field_adaptor.py:333):
input = input.detach().clone().requires_grad_(True)
score = torch.autograd.grad(config.guidance_fn(input, condition, m, std).sum(), input, ...)[0]So ∂θ̂/∂x = (I + std²∇s)/m never enters, and the guidance term is added straight to the x-space score. Under that, a constant x-push g corresponds to a θ₀-space displacement of g·std²/m relative to the denoising precision m²/std². So:
| scale | displacement in θ₀ units |
|---|---|
k/(m²·std²) (main) |
k/m³ |
k/std² (this PR) |
k/m |
k·m/std² |
k — independent of m |
Neither is m-invariant. m² in the denominator has the wrong sign — I could not derive it from the chain rule, from evaluating the wall on x_t/m, or from precision matching, all of which give 1/m, 1/m and m² respectively. Combined with no justifying comment at the point it was introduced and no test pinning the scale, dropping it is the right call.
But the residual 1/m is not negligible. I ran your own setup (exact N(0, 0.5²) score, lower bound 0, scale_factor=0.5, universal guidance with the same log-sigmoid wall, 2000 SDE samples, seed 1) on the PR head with no source edits. Entries are sample mean / half-normal mean (ideal 1.00):
| scale | ve | vp | subvp | fmpe |
|---|---|---|---|---|
k/(m²·std²) main |
1.03 | 1.56 | 1.49 | 2.89 |
k/std² this PR |
1.03 | 1.10 | 1.10 | 1.13 |
k·m/std² |
1.03 | 1.00 | 0.98 | 0.97 |
VE is identical across all three because m ≡ 1 there — which is also why interval guidance looked correct only for VE and wrong for every other method. That signature is what you would expect from an m-dependent error.
Why requesting changes
k/std² is a strict improvement and I would take it over main any day. But it leaves a systematic ~10-13% bias on VP, subVP and FMPE, and shipping a new default that we know is biased 13% means a follow-up fix later. scale_factor * m / (std**2 + 1e-6) is within ~3% on all four and is m-invariant under the current detached gradient. Alternatively keep k/std² and include the ∂θ̂/∂x factor, which is the more principled fix but a larger change.
Two smaller notes
- The test comment attributing the bias "up to 0.14 sd" to Tweedie looks misattributed. The same Tweedie step with
m/std²gives at most ~0.035 sd, so most of that bias comes from the scale, not from Tweedie. - Worth checking the "as in [1]" docstring claim — if [1] used VE, then
1/std²andm/std²coincide there and the citation does not decide between them.
(Posted as a comment on the scale line itself as well.)
| def like_input(value: Union[float, Tensor]) -> Tensor: | ||
| return torch.as_tensor(value, dtype=input.dtype, device=input.device) | ||
|
|
||
| scale = config.scale_factor / (std**2 + 1e-6) |
There was a problem hiding this comment.
Dropping m**2 is right, but 1/std**2 is still not m-invariant, because the guidance gradient is taken w.r.t. the detached θ̂ (line 333), so ∂θ̂/∂x never enters.
Relative to the denoising precision m**2/std**2, the induced displacement in θ₀ units is:
k/(m²·std²)→k/m³k/std²→k/mk·m/std²→k(m-independent)
Measured on the PR head with an exact Gaussian score (2000 SDE samples, seed 1), sample mean / half-normal mean, ideal 1.00:
| scale | ve | vp | subvp | fmpe |
|---|---|---|---|---|
k/(m²·std²) |
1.03 | 1.56 | 1.49 | 2.89 |
k/std² |
1.03 | 1.10 | 1.10 | 1.13 |
k·m/std² |
1.03 | 1.00 | 0.98 | 0.97 |
VE is identical everywhere because m ≡ 1. Suggest scale_factor * m / (std**2 + 1e-6).
Summary
When looking into the VF tests to finish #1461 I found a couple of problems in the FMPE methods. Guidance and iid data gave wrong samples with vector field posteriors. Most errors hit FMPE.
prior_guidewas wrong with every estimator, and interval guidance with all but VE. This PR fixes them, and the ODE sampler now refuses what it cannot do.Guidance read the velocity as the score. All guidance methods took the score from the estimator's forward pass. For NPSE, the forward pass returns the score. For FMPE, it returns the velocity, so guided FMPE samples diverged.
The FMPE score was 0.1% too large. FMPE trains on θ_t = (1 − t) θ_0 + (t + σ) ε, with σ =
noise_scale. For this path, the score and the diffusion of the forward SDE both have the factor (1 + σ), which was missing:The two errors cancel in the drift of the reverse SDE, so plain SDE sampling was almost exact. Guidance and the iid methods use the score alone and divide by m = 1 − t, which multiplies the error by up to 100 near t = 1. All FMPE SDE functions now cut time off at one value,
effective_t_max = 0.99, so they describe the same marginal at every time.Interval guidance made its walls too steep. It scaled them by
scale_factor / (m² s²). The Simformer paper (Gloeckler et al. 2024, Sec. 3.4) uses a scale proportional to 1 / s². The extra 1 / m² is 1 for VE, up to about 150 for VP and subVP, and 10⁴ for FMPE. Samples stayed in the interval, but far from the truncated posterior. The scale is nowscale_factor / s².@manuelgloeckler, was the m² in the wall scale intended? I found no reason for it in PR #1482 or PR #1752, and the measurements below favour the scale of the paper.
prior_guideignored its mixture weights. It normalized them over the wrong dimension of the sampler's(samples, 1, dim)inputs, so every weight became 1.The ODE sampler ignored guidance and failed on iid data. FMPE samples with the ODE by default, and the ODE does not use the adjusted score of guidance and iid methods. Guidance had no effect and gave no warning, and iid data failed with a reshape error. Both now raise a
ValueErrorthat names the reason and points tosample_with="sde".map()with iid data raises too. The vector field tutorial and the how-to guide now say that FMPE needssample_with="sde"for guidance and iid data.Smaller fixes from the same work:
FlowMatchingEstimator.scoreaccepts one time per sample, asforwarddoes.map(show_progress_bars=False)hides the progress bar of its initial samples, for all posteriors.jac_gaussno longer raises a singular-matrix error for posteriors much narrower than the noise. There, float32 rounds the matrix it inverts to zero; the corrected FMPE score made this reachable.effective_t_maxargument ofdrift_fn,diffusion_fnandmean_t_fnis deprecated. All SDE functions share the class attribute.Anyone who uses guidance or iid data with FMPE, or
prior_guideor interval guidance with any estimator, needs this.Evidence
Exact estimators for Gaussian posteriors, so every error comes from the sampler. Errors are in posterior standard deviations (sd).
mainprior_guide, VP and subVPprior_guide, FMPEscorewith the exact velocitymap()with iid dataValueErrorthat names the reasonThe new tests fail on
mainfor the estimators that each fix concerns. One test samples a Gaussian posterior with its exact score and a lower bound, and compares the samples with the half-normal.Merge Danger
Door: two-way
Blast Radius: small
FMPE SDE samples change by about 0.1%, and iid and guided FMPE samples a little more. ODE samples and
log_probdo not change. Guided samples with FMPE,prior_guidesamples, and interval-guided samples with VP and subVP change, because they were wrong. A user who tunedscale_factorfor the old walls gets softer walls with VP, subVP and FMPE.Code that combines ODE sampling with guidance or iid data, or that calls
map()with iid data, now raises.prior_guidestays 0.3 sd off with VE: it approximates the Jacobian of the denoiser by m I, which fits the wide VE marginals worst. A follow-up PR fixes that.