Make metropolis elemwise updates independent of each other #6678
Conversation
ricardoV94
force-pushed
the
fix_batched_metropolis
branch
from
c693453 to
c785d01
Compare
ricardoV94
changed the title
Not updating q0 after each elemwise update rendered subsequent proposals dependent on the previous ones.
ricardoV94
force-pushed
the
fix_batched_metropolis
branch
from
c785d01 to
9a1014e
Compare
| np.testing.assert_allclose(trace["x"].mean(("draw", "chain")), mu, rtol=0.1) | ||
| np.testing.assert_allclose(trace["x"].var(("draw", "chain")), mu, rtol=0.2) |
There was a problem hiding this comment.
The relative error was ~2x worse before the changes
Codecov Report
Additional details and impacted files@@ Coverage Diff @@
## main #6678 +/- ##
===========================================
- Coverage 91.99% 80.28% -11.71%
===========================================
Files 94 94
Lines 15944 15945 +1
===========================================
- Hits 14667 12801 -1866
- Misses 1277 3144 +1867
|
|
@bwengals you rejected the null in my PR! https://github.com/pymc-devs/pymc/actions/runs/4722680748/jobs/8377620494?pr=6678 Can I add a seed? |
|
Maybe NHST works, I guess this should happen 1/100 times. Yes absolutely, I should have done that. |
| q_temp = q0d.copy() | ||
| # Shuffle order of updates (probably we don't need to do this in every step) | ||
| np.random.shuffle(self.enum_dims) | ||
| for i in self.enum_dims: | ||
| q_temp[i] = q[i] | ||
| accept_rate_i = self.delta_logp(q_temp, q0d) | ||
| q_temp_, accepted_i = metrop_select(accept_rate_i, q_temp, q0d) | ||
| q_temp[i] = q_temp_[i] | ||
| q_temp[i] = q0d[i] = q_temp_[i] |
There was a problem hiding this comment.
If I understand correctly, this basically means that you are performing Gibbs updates at each batch dimension iteratively. So if you have a vector variable
If I my understanding is correct, then this looks ok to me
There was a problem hiding this comment.
Yes that's exactly it!
Only thing is we don't distinguish between core and batched dimensions, everything is unrolled as one long vector.
We don't perform elemwise_update if there is any multivariate discrete variable in the sampler, as those are usually impossible to update one value at a time (e.g., Multinomial).
For multivariate continuous, we implicitly assume they have a transform that makes each core dim more or less independent (or it's naturally so without a transform).
|
@bwengals I tried to seed the test, and the first seed I gave it also failed, so there could be something wrong going on. I'll push the seeded test in a separate PR so you can have a look. |
|
Ah ok thanks @ricardoV94, I commented there before seeing this. I'll try running it locally and see |
Not updating q0 after each elemwise update rendered subsequent proposals dependent on the previous ones.
This issue was revealed when testing the GeneralizedPoisson introduced in pymc-devs/pymc-extras#143
📚 Documentation preview 📚: https://pymc--6678.org.readthedocs.build/en/6678/