TuringLang / TuringLang/AbstractMCMC.jl

Adding facility for separation of the log likelihood and log prior via a Joint struct

Open
#75 7 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
108
Forks
20
PR merge metrics
No merged PRs in 30d

Description

Many of the packages that build upon AMCMC define model's with a log density function, in general this is the joint density consisting of the sum of the log prior and likelihood, for the MCMCTempering and some other situations (@yebai made me aware of these I cannot remember them off the top of my head but was assured I might have some luck in advocating for this as I believe it has a reasonably wide use case), I propose we add something like the following to AMCMC:

struct Joint{Tℓprior, Tℓll} <: Function
    ℓprior      :: Tℓprior
    ℓlikelihood :: Tℓll
end

function (joint::Joint)(θ)
    return joint.ℓprior(θ) .+ joint.ℓlikelihood(θ)
end

The purpose is to allow a user to - if they desire - define the log prior and likelihood separately and pass this in to a logdensity model in AMH / as the log density function in AHMC etc. In most cases there is not much motivation to do this as you would simply sum your components in the density function, but having them separable is critical for MCMCTempering at least to work, and in general is a cheap and very simple interface to facilitate operation on the prior / likelihood.

I am open to a bit of discussion on the design of this, as it is a bit of a weird thing to expose to a user potentially, unless they are using tempering or something like that, so I would imagine it would just be documented and flagged up as an option to users rather than being explicitly encouraged, but perhaps other people have an idea for how the same functionality can be achieved in a nicer way in terms of design.

In tempering, we use the user-defined Joint in the model to then apply a temperature to result in the following TemperedJoint, hopefully this illustrates what I mean:

struct TemperedJoint{Tℓprior, Tℓll, T<:AbstractFloat} <: Function
    ℓprior      :: Tℓprior
    ℓlikelihood :: Tℓll
    β           :: T
end

function (tj::TemperedJoint)(θ)
    return tj.ℓprior(θ) .+ (tj.ℓlikelihood(θ) .* tj.β)
end

Currently I need to add both the Joint and TemperedJoint to any tempering implementations that depend on MCMCTempering (aside from Turing which works definitely as the user never defines the logdensity anyway), so if Joint could go in AMCMC, I could add TemperedJoint to MCMCTempering and the requirements for adding tempering to samplers becomes even more trivial.

Woud like to hear thoughts, relatively simple PR if people are happy to add it.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Review the AbstractMCMC interface and the current MCMCTempering implementations to understand how log density functions are passed and transformed. Compare the proposed Joint and TemperedJoint roles, then clarify the API and documentation requirements with maintainers; done means an agreed design that supports tempering without duplicating Joint implementations.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend-api-design
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.