[FR] Add inference method for programs with stochastic support

[treigerm]

I was wondering whether there would an interest to add inference methods which support programs with stochastic support? If yes, I would be willing to add inference methods from the Divide, Conquer, Combine framework. I have implementations of both DCC and SDVI in Pyro but the implementations can be very much classed as "research code". Hence, I would want to start from a clean state to integrate the methods into NumPyro properly.

I'm happy to open a draft PR with an initial implementation, if that is of interest.

[martinjankowiak]

hi @treigerm

so it of course depends on details but i'd generally expect that the more exotic a model is (say the further it gets from a model with non-stochastic support) it'll generically be more awkward/cumbersome to represent it succinctly as a model in (num)pyro. do you expect there to be a relatively large space of models for which model specification is more-or-less straightforward and for which inference is tractable?

[treigerm]

Thank you for the response!

Yes, I think there is a reasonably large space of the programs with stochastic support which can be succinctly described in Pyro and for which inference can be done efficiently.

For example, one broad class of models are programs in which the user branches based on the outcomes of discrete sampling statements. In many cases users might have a program which has multiple sub-models that interact in certain ways but this interaction might change based on stochastic control flow. The most basic example of this is probably a mixture model with an unknown number of components. Another example/use case is when the user has several sub-models that get re-used within the main program. Writing the main program as a model with stochastic support is a nice succinct way of representing the different model choices. To give a high-level example of what I mean, here is some pseudo-code of a program with stochastic support:

def submodel1(X):
  """Sub-model with its own sample statements."""
  ...

def submodel2(X):
  """Sub-model with its own sample statements."""
  ...

def model(X, y):
  model_choice1 = sample(Categorical())
  if model_choice1 == 0:
    m1 = submodel1(X[:, 0])
  elif model_choice1 == 1:
    m1 = submodel2(X[:, 0])
  elif model_choice1 == 2:
    m1 = submodel1(X[:, 0]) + submodel2(X[:, 0])

  model_choice2 = sample(Categorical())
  if model_choice2 == 0:
    m2 = submodel1(X[:, 1]) + submodel2(X[:, 0])
  elif model_choice2 == 1:
    m2 = submodel1(X[:, 1] * X[:, 1]) + submodel2(X[:, 1] * X[:, 1])

  sample(Normal(m1 + m2, 1.0), obs=y)

The Divide, Conquer, Combine framework also provides a natural mechanism to allow inference in these programs with stochastic support in a language like numpyro because it breaks down the input program into sub-programs, one for each of the different control-flow paths, and runs inference in each sub-program separately. This then allows you to use a tracing based JIT compiler because the sub-programs have no longer any stochastic control flow.

[martinjankowiak]

@treigerm well certainly feel free to submit a PR! you may want to focus on one of the two algorithms initially. you may also want to start a discussion with an implementation sketch/plan before you dive into a big PR. also it will be helpful if you can (naturally) break up the PR into smaller ones

[treigerm]

Yes, will do! I will open a draft PR with a simple bare bones implementation of DCC so that we have a starting point to discuss design decisions.

I am busy with other things this week but will get to it next week.