Expression specific correlation in a model

Hello,

I’m working with a model of the form:

A s + n = b

where:

  • s is unknown,

  • n is noise,

  • b is observed, and

  • A is parametrized with latent variables that I aim to infer.

To estimate the parameters in A, I’m using a NUTS kernel. My implementation of the log likelihood is theoretically correct, but it’s not converging as expected. The log likelihood has the form:

(A^\top N^{-1} b)^\top (A^\top N^{-1} A)^{-1} (A^\top N^{-1} b)

where N is the noise covariance. When I compute the gradient of the negative log likelihood with the correct parameters, the result is approximately zero, suggesting correctness, but inference does not perform well.

For further context:

  • The matrix A has three parameters: a, b, and c.

  • The parameters are conditionally independent in their priors, but correlations arise through the likelihood. In particular, parameters a and b are more strongly correlated than c.

  • Interestingly, when I fix a to the correct value, the model converges well.

I would appreciate any insights into:

  1. Whether my log likelihood implementation might benefit from changes to improve convergence, or if there’s a common alternative for expressing such models.

  2. How to express the specific correlation pattern I observe between a, b, and c in a way that might assist with inference.

Thank you for any pointers or suggestions!