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.

11

I think we may be talking at cross purposes here. Your code does two things. It uses skron_prod to actually unfold a Kronecker product. This is almost always a bad idea compared to being lazy about it. The problem as you point out is that you call ksolve on that structure. What you need to do to make this efficient is pass the multiplicands you passed to skron_prod directly to ksolve and then exploit the Kronecker structure lazily.

Even in the form it’s in now, the Stan code can be made a lot faster. For instance, you can use blocked assignments in skron_prod which replaces M * N scalar products with a single matrix-scalar product, which is a lot faster in reverse mode.

If it turns out there’s no way to do that, which would be surprising, but not impossible, then you’re definitely not going to be able to use Stan for this at any scale.

12

I got quite confused by this thread @Bob_Carpenter and asked for help, making sure I’m not missing something etc – I appreciate you answer a lot of stuff so this is no dig! Pasting claude direct in case you’re curious or anyone else reads and gets confused:

@Bob_Carpenter I think there’s a mixup about which code path is which. ksolve isn’t the consumer of sqkron_prod output – it’s the alternative to it. ksolve is the continuous-time branch, and it already avoids any d^2 \times d^2 expansion by solving the symmetric d(d+1)/2 system directly. sqkron_prod only appears in the discrete-time branch, which isn’t the typical use here.

The other thing is that neither solve is a Kronecker product, so the identities don’t apply:

  • continuous time: the operator is the Kronecker sum A \oplus A = I \otimes A + A \otimes I. That solve is the Lyapunov equation, and (A \oplus A)^{-1} doesn’t factor. The part that does factor lazily, e^{(A \oplus A)\Delta} = e^{A\Delta} \otimes e^{A\Delta}, is exactly what gets used to reduce to the Lyapunov form in the first place.
  • discrete time: add_diag(-sqkron_prod(A,A),1) \ vec(Q) is (I - A \otimes A)^{-1}\mathrm{vec}(Q), i.e. identity minus a Kronecker product – a Stein operator, also not factorable that way.

So without a Sylvester/Schur solver in Stan there’s nothing lazy to exploit in the inverse, which is why the ugly expansion survived on the discrete path.

That said, your blocked-assignment point is well taken and exact, not just a refactor gamble – for fixed (i,j) the inner loops fill a contiguous block:

out[(d*(i-1)+1):(d*i), (d*(j-1)+1):(d*j)] = mata[i,j] * matb;

d^4 scalar vari nodes down to d^2 matrix-scalar ops. Will fix.

And for the discrete case there is actually a way to skip the big solve entirely, just not via Kronecker structure – Smith’s doubling iteration for X = AXA' + Q: X_0 = Q, A_0 = A, then X_{k+1} = X_k + A_k X_k A_k', A_{k+1} = A_k^2. Quadratic convergence whenever the stationary covariance exists, so ~10-15 iterations of quad_form_sym on d \times d, no d^2 \times d^2 object anywhere, and plain Stan that reverse mode handles fine. O(d^3 \log(1/\epsilon)) instead of d^6. So your instinct was right, just on the branch I don’t use much :)

On “not going to be able to use Stan at any scale” – the numbers in post 1 were the Stan version landing within ~3x of hand-written adjoint gradients in C++/Julia on 20-dim nonlinear systems, so it’s holding up fine. That’s what surprised me enough to post.

13

Thanks for clarifying. I misunderstood reading the quadruply nested loop code. My linear algebra isn’t the strongest and I’ve never even heard of a Kronecker sum.

It can still be a problem with scaling for reverse-mode autodiff applied directly as the memory requirements grow pretty quickly. I imagine being able to hand-write adjoints would give you much more memory headroom. Early on we were lazy in some of our evaluations in Stan just for memory and then eventually we factored into more efficient matrix autodiff with analytic gradients expressed using matrix operations. That doesn’t take up nearly so much memory as autodiffing through something like a matrix product or much worse, a solve of any kind.

14

Yes you’re right i didn’t mention the memory front but that was always pretty demanding and it’s much better with the custom adjoint stuff…

15

Do you mean a Sylvester solver and Schur decomposition or is there a solver in particular that you are referencing that uses the Schur (besides Sylvester)?

If this would be useful for you I can look into adding a Sylvester solver to Stan. jax has one so it should be a reasonable thing to port over

16

I was looking for this some years ago and had some cpp half ready for when rstan added complex number support, but that never eventuated and code is much cheaper now! Who knows maybe I port it back or use stan math somehow again but for now I got a few things that have been on the wish list for a while - faster performance and compile, laplace approx, hmc initialisation around the mode with auto metric etc… strange times

Edit: sorry yes Lyupanov solve (simplified Sylvester) with complex schur decomp was one of the puzzle pieces

17

A grown-up way of solving the Lyapunov equation would be useful to me, I’ve had issues doing it by hand.

18

If you mean the continuous time form, the ksolve function I posted was the best I could do in stan some years ago. I guess better is probably possible by now.

 matrix ksolve(matrix A, matrix Q){
    int d= rows(A);
    int d2= (d*d-d)%/%2;
    matrix[d+d2,d+d2] O;
    vector[d+d2] triQ;
    matrix[d,d] AQ;
    int z=0; //z is row of output
    for(j in 1:d){//for column reference of solution vector
      for(i in 1:j){ //and row reference...
        if(j >= i){ //if i and j denote a covariance parameter (from upper tri)
          int y=0; //start new output row
          z+=1; //shift current output row down
          
          for(ci in 1:d){//for columns and
            for(ri in 1:d){ //rows of solution
              if(ci >= ri){ //when in upper tri (inc diag)
                y+=1; //move to next column of output
                
                if(i==j){ //if output row is for a diagonal element
                  if(ri==i) O[z,y] = 2*A[ri,ci];
                  if(ci==i) O[z,y] = 2*A[ci,ri];
                }
                
                if(i!=j){ //if output row is not for a diagonal element
                  if(y==z) O[z,y] = A[ri,ri] + A[ci,ci]; //if column of output matches row of output, sum both A diags
                  if(y!=z){ //otherwise...
                    // if solution element we refer to is related to output row...
                    if(ci==ri){ //if solution element is a variance
                      if(ci==i) O[z,y] = A[j,ci]; //if variance of solution corresponds to row of our output
                      if(ci==j) O[z,y] = A[i,ci]; //if variance of solution corresponds to col of our output
                    }
                    if(ci!=ri && (ri==i||ri==j||ci==i||ci==j)){//if solution element is a related covariance
                      //for row 1,2 / 2,1 of output, if solution row ri 1 (match) and column ci 3, we need A[2,3]
                      if(ri==i) O[z,y] = A[j,ci];
                      if(ri==j) O[z,y] = A[i,ci];
                      if(ci==i) O[z,y] = A[j,ri];
                      if(ci==j) O[z,y] = A[i,ri];
                    }
                  }
                }
                if(is_nan(O[z,y])) O[z,y]=0;
              }
            }
          }
        }
      }
    }
    
    z=0; //get upper tri of Q
    for(j in 1:d){
      for(i in 1:j){
        z+=1;
        triQ[z] = Q[i,j];
      }
    }
    triQ=-O \ triQ; //get upper tri of asymQ
    
    z=0; // put upper tri of asymQ into matrix
    for(j in 1:d){
      for(i in 1:j){
        z+=1;
        AQ[i,j] = triQ[z];
        if(i!=j) AQ[j,i] = triQ[z];
      }
    }
    return AQ;
  }

Qasym = ksolve(A, Q);