Extending programmable VI to parametric discontinuities

Ascertain the extent to which the probabilistic programming language and ADEV-based modular variational inference transformations presented in this paper can be cleanly extended to support models and objectives with parametric discontinuities, leveraging methods such as reparameterization for affine discontinuities and techniques for differentiating integral expressions with parametric discontinuities, while preserving unbiased gradient estimation guarantees.

Background

The paper’s design forbids parametric discontinuities—expressions that compute discontinuous functions of input parameters—similar to constraints in Pyro, ProbTorch, and Gen. This restriction simplifies unbiased gradient estimation and compositional correctness but excludes important modeling cases where discontinuities arise.

Recent research has developed gradient estimators for certain classes of discontinuities (e.g., affine cases) and broader techniques for systematically differentiating integrals with parametric discontinuities. Integrating these techniques into the paper’s compositional program transformations (density, sim, and ADEV-based differentiation) could expand expressivity, but the authors state uncertainty about feasibility and scope.

References

No parametric discontinuities. A key limitation of our language, shared by Pyro, ProbTorch, and Gen, is that parametric discontinuities (expressions that compute discontinuous functions of the input parameters) are not permitted. Variational inference is possible in these settings, and \citet{lee2018reparameterization} proposed a gradient estimator that can be automated for a restricted PPL with affine discontinuities. More recently, \citet{bangaru2021systematically} and \citet{michel2024distributions} have presented techniques for differentiating integral expressions with parametric discontinuities. It is not yet clear to what extent the design we present could be cleanly extended to exploit these techniques.

— Probabilistic Programming with Programmable Variational Inference  (2406.15742 - Becker et al., 2024) in Discussion, Limitations (No parametric discontinuities)