Use provenance tracking to compute downstream costs in TraceGraph_ELBO

[ordabayevy]

This proposes to use ProvenanceTensor to track cost terms influenced by each non-reparameterizable sample site. It simplifies the code and in theory should have a more fine-grained Rao-Blackwellization.

[martinjankowiak]

@ordabayevy cool! i think you need to update the docstring to reflect the (potentially) finer level of dependency tracking

[martinjankowiak]

you should probably run some of the tests that are currently skipped in tests/integration_tests/test_conjugate_gaussian_models.py on your local machine

[ordabayevy]

Tests pass on the local machine:

$ pytest test_conjugate_gaussian_models.py
=========================== test session starts ===========================
platform linux -- Python 3.8.12, pytest-7.1.1, pluggy-1.0.0
benchmark: 3.4.1 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /home/ordabayev/repos/pyro, configfile: setup.cfg
plugins: xdist-2.5.0, benchmark-3.4.1, forked-1.4.0, nbval-0.9.6
collected 12 items

test_conjugate_gaussian_models.py ............                      [100%]

===================== 12 passed in 799.56s (0:13:19) ======================

[ordabayevy]

I'm happy to go through the PR code with anyone interested in reviewing it 😄

[fehiepsi]

I'm interested in. Thanks Yerdos!

[ordabayevy]

Awesome @fehiepsi ! Do you want to set up a Zoom meeting for Friday or Saturday?

[fehiepsi]

LGTM! Could you also run the attend, infer, and repeat example to see if we still get desirable result?

[ordabayevy]

I ran the air.py example using the new TraceGraph_ELBO and got 99.1% accuracy. Parameters used are the same as from the example:

python main.py -n 200000 -blr 0.1 --z-pres-prior 0.01 --scale-prior-sd 0.2 --predict-net 200 --bl-predict-net 200 --decoder-output-use-sigmoid --decoder-output-bias -2 --seed 287710

Graphs below show ELBO and Accuracy for the new (orange) and old (blue; not finished) TraceGraph_ELBO. They overlap very closely.

ELBO

Accuracy

[fehiepsi]

Thanks Yerdos! This is really nice.

[ordabayevy]

Thanks for reviewing @fehiepsi !