Improving differentiation for algebraic solver

[charlesm93]

Description

There are two ways to improve the way in which we compute the Jacobian matrix of an algebraic solver:

• compute the Jacobian of the system with respect to the parameters and the unknown simultaneously, to use their shared expression tree.
• compute the Jacobian-adjoint vector product directly, and thence only do one sweep of reverse-mode AD.

The two options are competing but should be fairly straightforward to implement and compare.
Note the procedure here will work for any algebraic solver where the sensitivities are computed via the implicit function theorem.

Expected Output

Increased speed when solving algebraic equations and computing their sensitivities. Will test on classical PK steady state problem and find the mode of a Laplace approximation.

Discussion on forum

https://discourse.mc-stan.org/t/better-computation-of-the-jacobian-matrix-for-algebraic-solver/8593

Current Version:

v2.19.1

[bob-carpenter]

If f:R^N -> R^M, then we can compute a Jacobian-adjoint with N directional derivative calculations (forward mode) or with M gradient calculations (reverse mode) and an explicit Jacobian-adjoint multiply. You seem to be indicating it can be done in a single sweep of reverse mode, but each sweep of reverse mode only computes derivatives with respect to a single output component of f.

[charlesm93]

The reason we require M gradient calculations in reverse mode is because of the initial cotangents we use. Suppose the output is y, with adjoint v = (a, b, c), and that the input is x.

With Jacobian(), the initial cotangents are (1, 0, 0), (0, 1, 0) and (0, 0, 1). This gives J, and then we multiply it by v, to get Jv, the adjoint of x. But instead, we could use one initial cotangent, (a, b, c), and directly get Jv in one sweep. This avoids redundant calculations. My understanding is that controlling the initial cotangent can be done with the adjoint-Jacobian method @bbbales2 developed.

[charlesm93]

We have three options here.

Let f(x, y) the function for which we want the root, x in R^M is the unknown, and y parameters in R^N. Let v be the adjoint of the output. The implicit function theorem gives us dx/dy = - [df/dx]^-1 df/dy. The options are:

1. compute df/dx and df/dy separately --> 2M sweeps
2. compute df/d(x + y) together --> M sweeps
3. compute df/dx and df/dy * v --> M + 1 sweeps

According to @betanalpha, when N, the number of parameters, becomes large the third option becomes better. I've encountered both regimes (M > N, and M < N), but we'll need to pick. Option 2 is the simplest to implement.

[bob-carpenter]

The function grad(v) initialize's v's adjoint to 1, and then does the backward pass. If all we need to do is initialize other adjoints, it'd be easy to write a more general grad(v). The adjoint-Jacobian thing from Ben is just a way to define a lazy chain() method for reverse mode. Rather than storing the Jacobian and multiplying it by the adjoint and adding it to the operand adjoints, we lazily evaluate incrementing the operand adjoints with the Jacobian-adjoint product. That allows us to get away with a lot less memory. It also saves virtual function calls on multivariate functions by managing non-propagating vari instances and putting all the work on a single vari.

[wds15]

Have you, @charelsm93, considered to apply the same trick as in the ode speedup pr which you reviewed recently? I mean, it sounds as if you do similar things here which is repeated gradients to get the Jacobian. So Parameters stay always the same and you can take advantage of that. Probably this is an additional thing to do here? For non stiff ode integration this gave me 20% speed for the sie example.

[charlesm93]

@bob-carpenter Indeed we would need a more general grad(v). I'm guessing it would have other applications too.

@wds15 That's a very good point. The trick applies to both options. For option (1) we use the fact the input x is constant and for option (2) the fact that both x and y remain constant. I still want to do this incrementally: first one of the above-described improvements, then your trick on top of it.

[betanalpha]

We don’t necessarily have to use grad(v) here — the Jacobian-adjoint product can be computed by implementing the nested sweeps by hand as I originally did in the ODE code.

…

On Jun 6, 2019, at 7:22 AM, Charles Margossian ***@***.***> wrote: @bob-carpenter <https://github.com/bob-carpenter> Indeed we would need a more general grad(v). I'm guessing it would have other applications too. @wds15 <https://github.com/wds15> That's a very good point. The trick applies to both options. For option (1) we use the fact the input x is constant and for option (2) the fact that both x and y remain constant. I still want to do this incrementally: first one of the above-described improvements, then your trick on top of it. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#1257?email_source=notifications&email_token=AALU3FVC5DCROHM76O7DZ2TPZDXPDA5CNFSM4HS32JH2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODXCRHAA#issuecomment-499454848>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AALU3FU5VZKLP6OF26RTCU3PZDXPDANCNFSM4HS32JHQ>.

[charlesm93]

Per the discussion during yesterday's Stan meeting, I'm going with option (2).

[bob-carpenter]

On Jun 6, 2019, at 11:44 PM, Michael Betancourt ***@***.***> wrote: We don’t necessarily have to use grad(v) here — the Jacobian-adjoint product can be computed by implementing the nested sweeps by hand as I originally did in the ODE code.

That's what I meant---writing a new function like grad(v) that computes an adjoint gradient. I don't think it's a good idea to try to generalize grad(v).

[charlesm93]

That's what I meant---writing a new function like grad(v) that computes an adjoint gradient. I don't think it's a good idea to try to generalize grad(v).

You mean in general, or for this specific issue?

[charlesm93]

I implemented option 2 -- fairly straightforward -- and ran some performance test.

The first test is based on pharma steady states, as described in the appendix of the autodiff review paper (https://arxiv.org/abs/1811.05031). The second test involves finding the mode of a conditional distribution, as required when doing INLA on a Poisson process with a latent Gaussian variable.

These tests show two things:

• the time to compute the Jacobian is small next to the time required to solve the algebraic system. For the Steady State problem, solving + differentiating takes about 0.03 s for 200 states, while constructing the Jacobian matrix requires 0.0075 s. But the latter is constructed using the grad function and looping over the 200 outputs. In practice, the Jacobian matrix is only calculated once, and this takes on the order e-05 s. Solving the equations however has order e-02s. Same with the INLA problem: for 500 states, the solver needs 633 seconds, and the Jacobian calculation is e-06 s.
• This is more intriguing: I did not observe any speed-up when looking at the time to calculate the Jacobian matrix. See the this graph for the steady state problem. The red line is the current solve, the blue one the new implementation. For the INLA problem I went down from 3e-6 to 2e-6, but I only ran the experiment once and the difference is likely due to noise.

Hmmm... did I do something wrong? The number of sweeps is halved and there is shared computation so there should be some speedup. In any case, this change provides at best a marginal improvement. The code is written and the PR review should not take very long, but I'm not sure that it's worth it.

[bob-carpenter]

I'd be happy to take a look at the code diff if you can point me to something manageable. The trick to answering your second question is to figure out what's dominating computation. If you speed up an operation that only takes 5% of the computation, you won't see much overall gain. Often computation's dominated by memory or branch prediction (and how the code gets generated to exploit CPU features) rather than by sheer number of operations.

[charlesm93]

@bob-carpenter See the code for the chain method on lines 66 - 73 in algebra_solver.hpp.

There are some changes in algebra_system, lines 61 - 74, where the () operator is modified to allow a vector containing both unknowns and parameters to be passed as the input variable. This is required for the call to Jacobian().

I think you're right: costs are dominated by memory and branch prediction, as opposed to number of operations. I'm running other tests that corroborate this conjecture.

[bob-carpenter]

One thing you can do is avoid [i], which checks bounds, and use .coeffRef(i). We should be doing this in a lot of places in the code.

The bigger issue is whether you can automate a lot of the stuff you built by hand using adjoint-Jacobian apply. The solver's chain() method is just doing that:

  void chain() {
    for (int j = 0; j < y_size_; j++)
      for (int i = 0; i < x_size_; i++)
        y_[j]->adj_ += theta_[i]->adj_ * Jx_y_[j * x_size_ + i];
  }

The Jacobian itself is calculated using a map and stored in J_x_y_ until the reverse pass:

   Map<MatrixXd>(&Jx_y_[0], x_size_, y_size_)
        = -mdivide_left(fx.get_jacobian(theta_dbl),
                        f_y(fs, f, theta_dbl, value_of(y), dat, dat_int, msgs)
                            .get_jacobian(value_of(y)))

Can you wait to construct the Jacobian until the chain() method is called? It should be more efficient as it'll cut down on the preallocated memory. Then you can also define it using holistic matrix operations rather than in a loop.