Repository: pyro-ppl/numpyro

Issue:
[FR] SVI: avoid initializing model params when already specified in guide
(Context: https://forum.pyro.ai/t/sharing-params-between-model-guide-i-e-tied-weights/5033/)

Sharing parameter values between a model and its guide is helpful in some applications, e.g. symmetric autoencoders. This can *almost* be accomplished by defining a param as usual in the guide and calling `numypro.param("shared_site_name")` in the model, but [since `SVI` only replays sample values](https://github.com/pyro-ppl/numpyro/blob/737d7d9ec0dc5d6649f85bbedddf28ca07908e20/numpyro/infer/svi.py#L186-L187) and not parameter values before tracing the model, the return value of this `param()` call defaults to `None` which breaks any downstream code that uses it.

The [desired behavior in SVI is to avoid initializing any model parameters already specified in the guide](https://forum.pyro.ai/t/sharing-params-between-model-guide-i-e-tied-weights/5033/6). I can see a couple of ways of accomplishing this:

1. Explicitly substitute initialized parameter values from the guide before tracing the model. Here's a strawman implementation which uses `substitute()` to avoid changing the semantics of `replay()`: https://github.com/pyro-ppl/numpyro/compare/master...msegado:numpyro:feat-shared-param-init.
2. Change the `replay()` handler to support replaying param values. This should probably be opt-in via an `include_params` kwarg or similar to avoid breaking existing user code, and would also require a change to the docs.

The first seems like the more natural approach to me, but I figured it would be good to ask before opening a PR 🙂 Thoughts?


Discussion:

Resolving patch:
diff --git a/numpyro/infer/svi.py b/numpyro/infer/svi.py
--- a/numpyro/infer/svi.py
+++ b/numpyro/infer/svi.py
@@ -21,7 +21,7 @@
 
 from numpyro.distributions import constraints
 from numpyro.distributions.transforms import biject_to
-from numpyro.handlers import replay, seed, trace
+from numpyro.handlers import replay, seed, substitute, trace
 from numpyro.infer.util import helpful_support_errors, transform_fn
 from numpyro.optim import _NumPyroOptim, optax_to_numpyro
 
@@ -184,9 +184,15 @@ def init(self, rng_key, *args, **kwargs):
         model_init = seed(self.model, model_seed)
         guide_init = seed(self.guide, guide_seed)
         guide_trace = trace(guide_init).get_trace(*args, **kwargs, **self.static_kwargs)
-        model_trace = trace(replay(model_init, guide_trace)).get_trace(
-            *args, **kwargs, **self.static_kwargs
-        )
+        init_guide_params = {
+            name: site["value"]
+            for name, site in guide_trace.items()
+            if site["type"] == "param"
+        }
+        model_trace = trace(
+            substitute(replay(model_init, guide_trace), init_guide_params)
+        ).get_trace(*args, **kwargs, **self.static_kwargs)
+
         params = {}
         inv_transforms = {}
         mutable_state = {}


Test patch:
diff --git a/test/infer/test_svi.py b/test/infer/test_svi.py
--- a/test/infer/test_svi.py
+++ b/test/infer/test_svi.py
@@ -190,6 +190,40 @@ def guide():
     assert_allclose(actual_loss, expected_loss, rtol=1e-6)
 
 
+def test_shared_param_init():
+    shared_init = 1.0
+
+    def model():
+        # should receive initial value from guide when used in SVI
+        shared = numpyro.param("shared")
+        assert_allclose(shared, shared_init)
+
+    def guide():
+        numpyro.param("shared", lambda _: shared_init)
+
+    svi = SVI(model, guide, optim.Adam(0.01), Trace_ELBO())
+    svi_state = svi.init(random.PRNGKey(0))
+    params = svi.get_params(svi_state)
+    # make sure the correct init ended up in the SVI state
+    assert_allclose(params["shared"], shared_init)
+
+
+def test_shared_param():
+    target_value = 5.0
+
+    def model():
+        shared = numpyro.param("shared")
+        # drive the shared parameter toward a target value
+        numpyro.factor("neg_loss", -((shared - target_value) ** 2))
+
+    def guide():
+        numpyro.param("shared", 1.0)
+
+    svi = SVI(model, guide, optim.Adam(0.01), Trace_ELBO())
+    svi_result = svi.run(random.PRNGKey(0), 1000)
+    assert_allclose(svi_result.params["shared"], target_value, atol=0.1)
+
+
 def test_elbo_dynamic_support():
     x_prior = dist.TransformedDistribution(
         dist.Normal(),