Comparing stan backend to Julia / native cpp (ctsem, continuous time nonlinear state space modelling)

Developers
1

Just thought it might interest some here – I’m migrating the ctsem software GitHub - cdriveraus/ctsem: Hierarchical continuous time state space modelling · GitHub (now with GUI for anyone really curious! GitHub - cdriveraus/ctsemGUI: GUI for ctsem package · GitHub ) from the rstan backend to Julia, and made a few comparisons along the way. I was surprised (and a little disappointed!) how well the stan version held up – the ctsem stan model is really pushing the stan code structure further than I think it was intended (particularly to avoid recompiling models except in complex nonlinear cases).

The continuous time extended kalman filter over multiple subjects involves a lot of repeated small - moderate size matrix operations (cholesky, exp, custom solves for the diffusion), and it’s primarily used for evaluating the log prob / gradient (hmc is too slow, though importance sampling from the mode works well in low to moderate dims) and I really thought a re-write would more dramatically improve things because of this. Both Julia and native cpp achieve ~3x the stan performance for gradient calculation on 20 dimensional nonlinear systems, which is very nice, but required custom adjoint gradient functions – none of the autodiff packages could handle a reverse mode pass on it. The julia backend anyway buys a lot more flexibility going forwards, avoids very long model recompiles when a user specifies anything complicated, and means I can retire the rstan dependency at some point, so I’m not unhappy with the change, to be clear ;) Mostly just impressed at the speed of the stan code, which wouldn’t have surprised me for simpler models, but did surprise me for my fairly horrendous model code ( ctsem/inst/stan/ctsm.stan at master · cdriveraus/ctsem · GitHub for those that like horror films) – but maybe it doesn’t surprise others?

2

If you wouldn’t mind testing, I’d be really interested to see how @seantalts’s stanli interpreter compares on a big model like this - since you get the usability benefit of no recompilation! If there’s a big perf difference that could also help identify some room for optimisation.

You can try it out in a new R package I’m working on:

# pak::pak("andrjohns/stanr")

mod <- stanr::stan_model("inst/stan/ctsm.stan", backend = "stanli")
fit <- mod$sample(data = ...)
3

Thanks for the review. It makes sense stans AD is faster until you write custom adjoint gradient functions. The idea of a library like Stan math is to avoid having to write
custom adjoint gradient functions :-P . If you think the adjoint functions you wrote would be nice for others to use you could also backport them to Stan math to get the same 3x speedup while being in the Stan ecosystem

4

@stevebronder not really sure what’s involved in backporting to stan math so I asked claude for an overview of what I have that is useful in case you or anyone wants to look at it or point in an easy direction :)

README.rmd (8.5 KB)

@andrjohns stanr seems to silently break rstan, so comparison is a bit too fiddly for me to look at sorry, but if you want to try the code below is as far as I got and should get you close.

```

pak::pak(“andrjohns/stanr”)
library(stanr);library(ctsem)
set.seed(1)

cmat ← diag(.5,10) + 1
cmat=t(chol(cmat %*% t(cmat)))

cmat2 ← diag(.5,10) + 1
cmat2[5:10,] ← cmat2[5:10,] * -1
cmat2=t(chol(cmat2 %*% t(cmat2)))
cov2cor(tcrossprod(cmat2))

gm ← ctModel(type=‘omx’,LAMBDA = diag(1,10),DRIFT=diag(-1,10),
T0VAR=cmat,
DIFFUSION=diag(1,10),Tpoints=2)
d1 ← data.frame(ctGenerate(ctmodelobj = gm,n.subjects = 1000,burnin = 0))

gm ← ctModel(type=‘omx’,LAMBDA = diag(1,10),DRIFT=diag(-1,10),
T0VAR=cmat2,
DIFFUSION=diag(1,10),Tpoints=2)
d2 ← data.frame(ctGenerate(ctmodelobj = gm,n.subjects = 1000,burnin = 0))

d2$id ← d2$id + 2000
d ← rbind(d1,d2)
d$TI1 ← 0
d$TI1[d$id > 2000] ← 1

m ← ctModel(LAMBDA = diag(1,10),manifestNames = paste0(‘Y’,1:10),
type=‘ct’,
TIpredNames = ‘TI1’)

f ← ctFit(datalong = d,model= m,priors=TRUE,cores=1,verbose=0,fit=F)

ctsf=ctsem:::stan_reinitsf(model = ctsem:::stanmodels$ctsm, data=f$standata)

system.time(rstan:::grad_log_prob(ctsf,rep(0,580)))

mod ← stanr::stan_model(f$stanmodeltext, backend = “stanli”)
stanrfit ← mod$sample(data = f$standata)
5

We’ve put a lot of work into efficiency tuning for autodiff. It has really helped that virtual function calls continue to get faster, so putting everything onto a stack and essentially interpreting it as reverse mode does continues to get faster. There have also been a ton of small improvements in very important things like memory locality and vectorization (in the CPU sense, not the Stan sense) and lazy evaluation (implemented with expression templates on the C++ side). Most of this was way beyond what we were capable of when coding v1, in part because we weren’t even up to C++11.

In addition to speed of the model log density and gradient evaluations, there have been a few algorithmic improvements like Nutpie and Walnuts, which I’d suggest adopting if you’re rebuilding Nuts on the Julia side.

6

Your Stan code could be made much faster by using vector and matrix operations rather than loops. For example, ctsmgen.stan has:

matrix sqkron_prod(matrix mata, matrix matb){
    int d=rows(mata);
    matrix[d*d,d*d] out;
    for (k in 1:d){
      for (l in 1:d){
        for (i in 1:d){
          for (j in 1:d){
            out[ d*(i-1)+k, d*(j-1)+l ] = mata[i, j] * matb[k, l];
          }
        }
      }
    }
    return out;
  }

I’m not sure where this goes in the code, but one should almost never expand Kronecker products explicitly. We don’t support lazy Kronecker in Stan, but it’s really wasteful in autoidff to do this as you have it. It should be doubly-nested loop using scalar * matrix assigned to a block on the inside. And if they’re covariance matrices, it shouldn’t repeat the above-and-below diagonals, but rather should just transpose and copy by block.

I’m not really good enough at the linear algebra, but I think the solves could be greatly simplified as it looks like they’re using Kronecker products.

This kind of code can be made a lot faster by grouping operations and copying rather than repeating math.

   matrix constraincorsqrt1(matrix mat){ //converts from unconstrained lower tri matrix to cor
    int d=rows(mat);
    matrix[d,d] o;
    vector[d] ss = rep_vector(0,d);
    vector[d] s = rep_vector(0,d);
    real r;
    real r3;
    real r4;
    real r1;
    
    for(i in 1:d){
      for(j in 1:d){
        if(j > i) {
          ss[i] +=square(mat[j,i]);
          s[i] +=mat[j,i];
        }
        if(j < i){
          ss[i] += square(mat[i,j]);
          s[i] += mat[i,j];
        }
      }
      s[i] += 1e-5;
      ss[i] += 1e-5;
    }
    
    
    for(i in 1:d){
      o[i,i]=0;
      r1=sqrt(ss[i]);
      r3=(abs(s[i]))/(r1)-1;
      r4=sqrt(log1p_exp(2*(abs(s[i])-s[i]-1)-4));
      r=(r4*((r3))+1)*r4+1;
      r=(sqrt(ss[i]+r));
      for(j in 1:d){
        if(j > i)  o[i,j]=mat[j,i]/r;
        if(j < i) o[i,j] = mat[i,j] /r;
      }
      o[i,i]=sqrt(1-sum(square(o[i,]))+1e-5);
    }
    
    return o;
  }

Let me give it a go for everyone’s benefit. I didn’t test this, so there are almost certainly typos and brainos in it, but the general idea of the refactor is sound.

  matrix constraincorsqrt1(matrix mat) {
    double EPSILON = 1e-5;
    int d = rows(mat);
    vector[d] ss;
    vector[d] s;

    matrix mat_zero_diag = mat;
    for (i in 1:d)
      mat_zero_diag[i, i] = 0;
      ss[i] = sum(square(mat_zero_diag[ , i])) + EPSILON;
      s[i] += sum(mat_zero_diag) + EPSILON;
    } 

    vector[d] r1 = sqrt(ss);
    vector[d] r3 = abs(s) ./ r1 - 1;
    vector[d] r4 = sqrt(log1p_exp(2 * (abs(s) - s - 1) - 4));
    vector[d] r = sqrt(r1 + (r4 .* r3 + 1) .* r4 + 1);
    matrix[d, d] o = rep_matrix(0, d, d);
    for (i in 1:d) {
      o[1:(i - 1), i] = mat[i, 1:(i - 1)] / r;
    }
    o += o';
    for (i in 1:d) {
      o[i,i] = sqrt(1 - sum(square(o[i, ])) + EPSILON);
    }
    return o;
  }

Mainly, I’ve replaced looped operations with grouped operations. There’s overhead to construct intermediates, but it’s generally a win from fewer virtual function calls in the backward pass and it also lets us call vectorized versions of things like the square root and absolute functions, which are much faster on modern hardware. I assume these same issues pertain to coding in Julia. It’s really hard for the compiler to figure any of this out even with scalars, but the autodiff makes it almost impossible.

Having said that, I haven’t tested and the proof’s always in the pudding.

I think the bigger win here is figuring out how to not expand the Kronecker factors or enter into the four nested loops.

7

The kronecker was always a big issue, one of the wins from the julia switch was getting rid of it by using complex schur / Lyupanov equation solves (this could probably be done in modern stan now but still not rstan). It was probably 5 years ago that i wrote the kronecker expansion you pointed at, but at the time it was much faster than the naive route involving kronecker products - I wouldn’t have kept the ugly thing otherwise. And yes good spot on the symmetry :)

8

The main issue is not ever expanding Kronecker products. If, for example, you have two positive-definite matrices and Kronecker them, then the relevant solves factor through solves of the input matrices, which will be much smaller. So the trick’s to never expand into the full Kronecker form, but do all the math that would have been implicated. For example, we can use a matrix normal rather than Kroneckering everything. Given an M x N matrix y with a matrix location parameter and independent covariance matrices for the rows and columns, the density is

\textrm{vec}(y) \sim \textrm{multiNormal}(\textrm{vec}(\mu), \Sigma^\textrm{col} \otimes \Sigma^\textrm{row}),

where the \textrm{vec} operation concatenates a matrix into a vector in column-major order.

If you were to actually evaluate everything by unfolding the Kronecker products, it’d involve an \mathcal{O}\!\left((N \times M)^3\right) operation to solve the product, whereas everything you need is in the \mathcal{O}\!\left(N^3 + M^3\right) time for the individual solves. This isn’t even considering the \mathcal{O}\!\left((N \times M)^2\right) memory. The key relations for evaluating the normal density are

(A \otimes V)^{-1} = A^{-1} \otimes V^{-1},

and

| \textrm{det}(A \otimes B) | = | \textrm{det}(A)|^N \cdot | \textrm{det}(B)|^M.
10

@Bob_Carpenter I suspect the SDE diffusion use of the Kronecker is a bit different to what you’re familiar with, but maybe I’m missing something :)

For a linear continuous-time SDE,

dx(t) = A x(t)\,dt + L\,dW(t),

with instantaneous diffusion covariance

Q = LL^\top,

the discrete-time process-noise covariance over an interval \Delta is

Q_\Delta = \int_0^\Delta e^{As} Q e^{A^\top s}\,ds.

That integral is the fundamental quantity. The Kronecker approach is just one way of evaluating it. Using

\operatorname{vec}(AXB) = (B^\top \otimes A)\operatorname{vec}(X),

we get

\operatorname{vec}(Q_\Delta) = (A\oplus A)^{-1} \left[ e^{(A\oplus A)\Delta} - I \right] \operatorname{vec}(Q),

where

A\oplus A = I\otimes A + A\otimes I.

So yes, explicitly constructing the d^2 \times d^2 Kronecker matrix is computationally unattractive. The reason it is there is that the more natural alternatives aren’t straightforwardly available in Stan or were slower than the custom solve I wrote when tested years ago.

For example, the same quantity can be obtained from the Sylvester/Lyapunov equation

A Q_\Delta + Q_\Delta A^\top = e^{A\Delta} Q e^{A^\top\Delta} - Q,

which avoids constructing the d^2 \times d^2 system if you have a proper Sylvester/Lyapunov solver. As far as I know, Stan doesn’t.

The other standard option is the Van Loan method, which uses a block matrix exponential. That also avoids the Kronecker expansion, but requires a larger matrix exponential.

So the Kronecker isn’t really being used here in a way that can simply be replaced by grouped matrix operations; it’s being used to convert the continuous-time SDE diffusion integral into a solve that Stan can actually perform.