Stan gradient w.r.t ? in advi code

tl;dr
what is the result of stan::model::gradient(m, zeta, tmp_lp, tmp_mu_grad, &ss);?
\zeta value plugged in to \nabla_{\theta} \log p(\mathbf{X}, \theta) OR \nabla_{\zeta} \log g(\mathbf{X}, \zeta)?


It seems this ADVI code,
stan::model::gradient(m, zeta, tmp_lp, tmp_mu_grad, &ss);
is calculating formula (equation 5, below) from the paper
\nabla_{\theta} \log p(\mathbf{X}, \theta) \nabla_{\zeta} T^{-1}(\zeta)+\nabla_{\zeta} \log \left|\operatorname{det} J_{T^{-1}}(\zeta)\right|
= \nabla_{\zeta} \log g(\mathbf{X}, \zeta)
( g(\mathbf{X}, \zeta) = p\left(\mathbf{X}, T^{-1}(\zeta)\right)\left|\operatorname{det} J_{T^{-1}}(\zeta)\right|)

as mu_grad equals tmp_mu_grad here.

According to stan::model::gradient (f, x, fx, grad_fx) documentation, it calculates the value and the gradient of the specified function at the specified argument.

So, from the above, gradient of tmp_lp calculated from the model m would return \zeta value plugged in to \nabla_{\theta} \log p(\mathbf{X}, \theta), not \nabla_{\zeta} \log g(\mathbf{X}, \zeta).

From here, I assumed that m would return lp function of original parameters (\theta, cont_params_) not \zeta.

This stan manual says lp__ is not just p(\theta, x) but evaluation of the model on the unconstrained scale.

So is it a correct understanding that lp__ returned by models in stan c++ are always the function of unconstrained parameters and therefore \nabla_{\zeta} \log g(\mathbf{X}, \zeta) is the correct answer?

I think the gradients here will be the log density with respect to the unconstrained space. So that’s \nabla_{\zeta} p(x | \theta)p(\theta) |J_{T^{-1}}(\theta)| where \theta = T^{-1}(\zeta)

I think in this notation, \theta are the model parameters, \zeta are the unconstrained parameters, and \phi are the parameters of the variational distribution, q(\zeta | \phi).

1 Like