You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
fix: VI quality control retrains once, at the caller's observation and settings - #2149
I found this while benchmarking the VI flows for smaller defaults. VIPosterior.train checks the fitted posterior at the end (quality control). When the check failed, the retry it started had three problems:
It sampled and evaluated q at the default observation, not at the x passed to train. So train(x=x_o) on a posterior with another default x failed the check, and the retry then trained q on the default x. The caller got a posterior for the wrong observation, silently.
The retry dropped the caller's settings (x, iteration limits, particles, optimizer arguments) and ran the check again. A check that kept failing retrained until the recursion limit.
It caught every exception, not only a failed check.
Summary
train(x, ..., quality_control=True)
fit q at x
- try evaluate() # samples q at the default x- except Exception:- print(...)- train(learning_rate=lr * 0.1) # default x and settings, check on again+ check q at the x it was trained on+ on QualityControlError:+ warn, retrain from scratch once # caller's x and settings, lr * 0.1, check off+ check again+ still failing: warn "may be unreliable"
The basic checks now raise QualityControlError instead of assert, so they also work under python -O. A q passed as an already-built distribution cannot be rebuilt, so a failed check only warns there. Other errors from the quality metrics now reach the caller.
Evidence
Case
Before
After
train(x=x_o), default x differs
q trained on the default x
q trained on x_o, check passes
Check keeps failing
retrains until the recursion limit
one retry with the caller's settings, then a warning
Already-built q fails the check
raises "Cannot rebuild q" from inside the handler
warns
Three new tests in tests/vi_test.py cover these cases. Each fails on main.
Merge Danger
Door: two-way
Blast Radius: small. It affects only VI training with quality control on, which is the default. Messages are now warnings instead of prints.
…fails
When the quality control of `VIPosterior.train` caught an error, it retrained by calling `train` again with only a smaller learning rate. That retry dropped the caller's `x`, iteration limits and optimizer settings, and it ran quality control again, so a check that kept failing retrained over and over until Python's recursion limit. It also caught every exception, not only a failed check.
The basic checks now raise `QualityControlError` instead of `assert`, and `train` catches only that. A failed check retrains from scratch once, with the caller's settings and a 10 times smaller learning rate, and warns; if the check still fails, or `q` was passed as an already-built distribution that cannot be rebuilt, it warns that the posterior may be unreliable.
The basic checks sampled and evaluated `q` at the default observation. After `train(x=...)` with another observation, sampling there raised, and the old catch-all handler then retrained `q` on the default observation instead of the caller's. The checks now use the observation `q` was trained on.
No actionable comments were generated in the recent review.
Walkthrough
Variational-inference quality checks now use the training observation and raise a dedicated error when checks fail. Training handles that error with a conditional retry that reduces the learning rate and preserves the caller’s training settings.
basic_checks uses the training observation for posterior sampling and log-probability checks. It raises QualityControlError when samples or log probabilities fail the checks. A test verifies that training records the supplied observation.
Conditional training retry sbi/inference/posteriors/vi_posterior.py, tests/vi_test.py
VIPosterior.train catches QualityControlError and retries once at one-tenth the learning rate when q can be rebuilt. The retry preserves the caller’s training settings and disables further quality control. Instance-based q is not retried. Other exception types propagate. Tests cover retry and no-retry behavior.
This change makes VI quality control evaluate at the observation passed to train and retry once with the caller's settings at a reduced learning rate. No concrete merge-blocking risk was identified in the supplied evidence.
Pre-merge checks: 5
Check name
Status
Explanation
Title check
The title clearly and concisely describes the main change: a single VI quality-control retry using the caller’s observation and settings.
Description check
The description is detailed, on-topic, and explains the problem, solution, behavior changes, test coverage, and impact. It does not include the template’s explicit issue-status section or completed ch…
Docstring Coverage
Docstring coverage is 83.33% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentati…
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, on-topic, and explains the problem, solution, behavior changes, test coverage, and impact. It does not include the template’s explicit issue-status section or completed checklist, but the required change information is mostly complete.
Full details: Docstring Coverage
Explanation
Docstring coverage is 83.33% which is sufficient. 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 6 functions across 3 files.
Finishing Touches
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.
File: No configuration file from this repository applied
Commit: Unavailable
Layers: Organization UI
Review profile: CHILL
Plan: Advanced
Run ID: 17f0fc34-8a5a-4365-90d0-525783ad27ea
Commits
Reviewing files that changed from the base of the PR and between a38be4c and b7cecb8.
Files selected for processing: 3
sbi/inference/posteriors/vi_posterior.py
sbi/samplers/vi/vi_quality_control.py
tests/vi_test.py
Included review availability: This review used your included allowance. Your plan provides up to 4 included reviews per hour; 3 remain after this review.
❌ Patch coverage is 83.33333% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.17%. Comparing base (c36f471) to head (b7cecb8). ⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.
The reason will be displayed to describe this comment to others. Learn more.
Approving — this is a clean fix of a real bug.
I checked that the three new tests fail on main and pass here, and that the diagnosis is right: evaluate() never forwards an observation, so basic_checks was sampling at the default x while q had been trained at the caller's. Passing _trained_on through fixes exactly that.
The retry half is the more important fix. The old recursive train() call without quality_control=False would re-run the quality control on every retry, so a persistently failing metric recursed; passing the caller's settings and disabling quality control on the single retry is correct.
Two non-blocking notes, no action needed for this PR:
There is no test for a retry that succeeds (first check fails, second passes). That is the branch the PR exists to make work, so it would be nice to have eventually.
basic_checks still uses a bare assert for the _prior is not None check on line 71 while everything else now raises QualityControlError. That looks deliberate rather than an oversight — a missing _prior is a programming error, not a quality failure — so I would not change it, but the asymmetry is worth a comment in the code.
I also checked the untrained-posterior path (_trained_on is None): it passes x=None into sample(), which then hits the existing "was not fit on the specified observation" ValueError. Same as the old behaviour, and it correctly propagates rather than being swallowed by _quality_control_error, which only catches QualityControlError.
The tests covered a check that keeps failing and an already-built q, but not the case the single retry exists for: the first check fails, the retrained q passes, and train() only warns that it is retraining.
VIPosterior.__init__ raises a ValueError when it finds no prior, and set_q and to() keep a prior in place, so the assert could not fire. A prior set to None by hand still fails with an error that skips the retry, since only QualityControlError triggers one.
Thanks for the review! I did both notes in this PR:
A new test covers a retry that passes: the first check fails, the retrained q passes, and train() only warns that it retrains.
I removed the prior assert instead of commenting on it. VIPosterior already raises when it finds no prior, so the assert could not fire, and a prior set to None by hand still fails without a retry.
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
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.
manuelgloeckler
left a comment
I found this while benchmarking the VI flows for smaller defaults.
VIPosterior.trainchecks the fitted posterior at the end (quality control). When the check failed, the retry it started had three problems:qat the default observation, not at thexpassed totrain. Sotrain(x=x_o)on a posterior with another default x failed the check, and the retry then trainedqon the default x. The caller got a posterior for the wrong observation, silently.x, iteration limits, particles, optimizer arguments) and ran the check again. A check that kept failing retrained until the recursion limit.Summary
The basic checks now raise
QualityControlErrorinstead ofassert, so they also work underpython -O. Aqpassed as an already-built distribution cannot be rebuilt, so a failed check only warns there. Other errors from the quality metrics now reach the caller.Evidence
train(x=x_o), default x differsqtrained on the default xqtrained onx_o, check passesqfails the checkq" from inside the handlerThree new tests in
tests/vi_test.pycover these cases. Each fails onmain.Merge Danger
Door: two-way
Blast Radius: small. It affects only VI training with quality control on, which is the default. Messages are now warnings instead of prints.