Soss relating to the broader PPL ecosystem

Specific Domains Probabilistic programming
1

Hello Julia PPL,

I’ve been thinking some more about how Soss might relate to other PPL work. Hope I might get your thoughts on this, and how it might benefit the community as a whole.

The big idea is that Soss is probabilistic glue.

Say we have a relatively generic model like

m = @model x, d1, d2
    z ~ d1(x)
    y ~ d2(x,z)
end

This acts like a parameterized family of distributions, similar to Normal, etc. And like Normal, specifying parameters produces a distribution, in this case a Soss JointDistribution. These are handy, because it’s very easy to rand the whole thing or build predictive distributions that remove ancestors of given variables (there’s really more to it than that, but that’s the idea).

Like a distribution, you can (not always, but often) call rand, logpdf, and some other things. In Soss we do this by passing responsibility to the components, and generating code to connect the pieces in the right way. We will be considering special cases where code can be rewritten, but we can always fall back on this per-component approach. At inference time, you can specify any variables you want and reason about the rest.

Typically, d1 and d2 would be Distributions. But they can really be anything where the required methods are available: Soss models, or even if principle Turing or Gen models.

It can go the other way as well; any PPL that requires a logpdf can call a Soss model. This could make it easy, for example, to wrap a Gen model in Turing, or vice-versa, by using Soss as the glue.

This opens up some interesting possibilities. For example, the design of Turing makes it difficult to connect to MLJ, but it seems this may not be a problem for Soss. So a solution could be to solve it for Soss, then connect Soss with Turing. I’d expect the benefits to be the same as those of any glue code.

I’m still feeling this out, but it What do you think?

5 Likes
2

One issue I think with this is that Turing doesn’t really directly support the Distributions API to the same extent as Soss, so it would be very difficult to have a Turing-within-Soss model. It would be moderately trivial I think to do it the other way for a couple of our samplers, though I don’t know how Soss would work in the case of parameter transforms.

3

Really? I wasn’t aware of this. I’ve had some trouble with Distributions, and Gen doesn’t seem to use it at all. This makes me think we need to either Build a separate library specializing in PPL needs, or determine what needs to change for Distributions to be a better fit.

What parts of this cause problems for you? For me, it’s mostly the way the types are set up. I’d love to have an alternative where distribution structs are relatively unconstrained (IMO constraints belong in methods, not structs), new methods are very easy to implement, and the types are set up in way that’s optimized for PPL.

Still, I don’t see this as a limitation. Inference in Soss works by calling primitives (logpdf, rand, xform, etc) on the components, though this can be specialized as needed.

For parameter transforms, do you mean like the ones needed for HMC? I’m using TransformVariables.jl because of the support for named tuples, but I have a function xform that takes a distribution and returns the appropriate transformation. I think Bijectors.jl has a similar approach.

4

I suppose by “Distributions API” I mostly mean the suite of rand, logpdf, and a couple others you might need. In Turing’s current iteration, the model typically needs a sampler to access a parameter, and so you can’t just call rand(model) or logpdf(model, some_parameters). My understanding was that Soss did that as kind of a first order thing.

Yeah, I think we need to have a unified interface on the model side, or at the very least an agreement on what you should be able to call on a model. I think rand and logpdf are the two that should definitely be included.

5

If I’m understanding you correctly, this is also the core idea behind Gen’s Generative Function Interface (GFI) :slight_smile: The GFI is just the set of methods that are supported by models, including (as you write) “rand, logpdf, and a couple others you might need.” For models that call other models, these methods are implemented recursively, by delegating to the callee’s implementation. However, the user can always implement any of these methods in a custom way, so this “recursive” compositional implementation is just a helpful default.

As you note, such an interface is useful because you can then “glue together” models written in different languages, as long as each language knows how to implement the GFI for models written within it.

We used this in Gen to “glue together” models written in a Dynamic language, a Static language, and a TensorFlow-based language, but in principle, you could also use it to glue Soss models, Turing models, and Gen models together. In fact, in an earlier draft of the paper, we had called the GFI the “FFI” instead – a kind of probabilistic programming foreign function interface.

I’d love to have an alternative where distribution structs are relatively unconstrained (IMO constraints belong in methods, not structs), new methods are very easy to implement, and the types are set up in way that’s optimized for PPL.

I agree! One reason Gen settled on this design was precisely to allow different models to use different structs under the hood.

Yeah, I think we need to have a unified interface on the model side, or at the very least an agreement on what you should be able to call on a model. I think rand and logpdf are the two that should definitely be included.

The GFI is our proposed “unified interface”: in addition to rand (which we call simulate) and logpdf (which we call assess), we have several others, including generate (which proposes the values of some variables in the model given constraints on the others), update (which accepts a trace of the model and some updates to make, and efficiently computes a log probability ratio between the old trace and the new trace – useful for Metropolis-Hastings), and others useful for back-propagation. We’ve found that this interface gives programmers the primitives necessary to implement an inference library for Monte Carlo and variational inference that works with any model. Not that this interface is the be-all/end-all – I’m sure it’ll continue to evolve!

It’d be great to see the Julia PPL community coalesce around this kind of implementation strategy (treating models as black-box “probabilistic computations,” which you interact with via a well-defined interface), whether or not Gen’s interface is to people’s liking. I think it’s a useful idea, not just for model inter-op but also for user-friendly inference programming (the inference algorithm implementor doesn’t have to learn about the internals of any particular PPL’s implementation, just the interface that models promise to satisfy). But of course, there’s also value in continuing to research new implementation strategies that might not fit easily within this framework.

1 Like
6

I really need to find time to dig into Gen, especially since the approaches are so similar (well, from a really high level anyway). The idea of these working together or merging is pretty exciting. And what I’m really after is to be able to use something that satisfies the goals of Soss, so even if Gen obviates it that could be ok.

Much more to say about this, but it’s sprint time for a few days. More later :slight_smile:

2 Likes