Computing the Laplace approximation in JAX

This is a two step process and, to be honest, all of the steps are pretty standard. So (hopefully) this will not be too tricky to implement. For simplicity, I’m not going to bother with the dividing and multiplying by , although for very large data it could be quite

import jax.numpy as jnp
from jax.scipy.optimize import minimize
from jax.scipy.special import expit
from jax import jacfwd, grad
from jax import Array
from typing import Callable, Tuple, List, Set, Dict

def laplace(f: Callable, x0: Array) -> Array:
    nx = x0.shape[0]
    mode, *details = minimize(lambda x: -f(x), x0, method = "BFGS")
    H =  -1.0 * jacfwd(grad(f))(mode)
    return mode, H

There are a few things worth noting here. There’s not really much in this code, except to note that jax.scipy.optimize.minimize finds the minimum of , so I had to pass through the negative of the function. This change also propagates to the computation of the Hessian, which is computed as the Jacobian of the gradient of f. 

Depending on what needs to be done with the Laplace approximation, it might be more appropriate to output the log-density rather than just the mode and the Hessian, but for the moment we will keep this signature.

Let’s try it out. First of all, I’m going to generate some random data from a logistic regression model. This is going to use Jax’s slightly odd random number system where you need to manually update the state of the pseudo-random number generator. This is beautifully repeatable⁷ unlike, say, R or standard numpy, where you’ve got to pay a lot of attention to the state of the random number generator to avoid oddities.

from jax import random as jrandom

def make_data(key, n: int, p: int) -> Tuple[Array, Array]:
  key, sub = jrandom.split(key)
  X = jrandom.normal(sub, shape = (n,p)) /jnp.sqrt(p)

  key, sub = jrandom.split(key)
  beta = 0.5 * jrandom.normal(sub, shape = (p,))
  key, sub = jrandom.split(key)
  beta0 = jrandom.normal(sub)


  key, sub = jrandom.split(key)
  y = jrandom.bernoulli(sub, expit(beta0 + X @ beta))

  return (y, X)

An interesting side-note here is that I’ve generated the design matrix to have standard Gaussian columns. This is not a benign choice as gets big. With very high probability, the columns of will be almost⁸ orthonormal, which means that this is the best possible case for logistic regression. Generally speaking, design matrices from real⁹ data have a great deal of co-linearity in them and so algorithms that perform well on random design matrices may perform less well on real data.

Ok, so let’s fit the model! I’m just going to use priors on all of the s.

from functools import partial
n = 100
p = 5

key = jrandom.PRNGKey(30127)
y, X = make_data(key, n, p)

def log_posterior(beta: Array, X: Array, y: Array) -> Array:
    assert beta.shape[0] == X.shape[1] + 1

    prob = expit(beta[0] + X @ beta[1:])
    
    return (
      jnp.sum(y * jnp.log(prob) + 
      (1-y) * jnp.log1p(-prob)) - 
      0.5 * jnp.dot(beta, beta)
    )


post_mean, H = laplace(
  partial(log_posterior, X = X, y = y),
  x0 =jnp.zeros(X.shape[1] + 1)
)

post_cov = jnp.linalg.inv(H)

Let’s see how this performs relative to MCMC. To do that, I’m going to build and equivalent PyMC model.

import numpy as np
import pymc as pm
import pandas as pd

with pm.Model() as logistic_reg:
  beta = pm.Normal('beta', 0, 1, shape = (p+1,))
  linpred = beta[0] + pm.math.dot(np.array(X), beta[1:])
  
  pm.Bernoulli(
    "y", 
    p = pm.math.invlogit(linpred),
    observed = np.array(y)
  )
  posterior = pm.sample(
    tune=1000, 
    draws=1000, 
    chains=4, 
    cores = 1)

# I would like to apologize for the following pandas code.
tmp = pm.summary(posterior)
tmp = tmp.assign(
  laplace_mean = post_mean, 
  laplace_sd = np.sqrt(np.diag(post_cov)), 
  Variable = tmp.index
)[["Variable", "mean", "laplace_mean", "sd", "laplace_sd"]]

with pd.option_context('display.precision', 3):
    print(tmp)

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Sequential sampling (4 chains in 1 job)
NUTS: [beta]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 5 seconds.

        Variable   mean  laplace_mean     sd  laplace_sd
beta[0]  beta[0]  0.249         0.234  0.235       0.229
beta[1]  beta[1] -0.964        -0.914  0.435       0.428
beta[2]  beta[2] -1.710        -1.616  0.490       0.470
beta[3]  beta[3] -0.975        -0.926  0.423       0.416
beta[4]  beta[4] -0.739        -0.716  0.470       0.457
beta[5]  beta[5]  0.637         0.609  0.481       0.475

Well that’s just dandy! Everything is pretty¹⁰ close. With 1000 observations it’s identical to within 3 decimal places.

Speeding up the computation

So that is all well and dandy. Let’s see how long it takes. I am interested in big models, so for this demonstration, I’m going to take . That said, I’m not enormously interested in seeing how this scales in (linearly), so I’m going to keep that at the fairly unrealistic value of .

import timeit
def hess_test(key, n, p):
  y, X = make_data(key, n , p)
  inpu = jnp.ones(p+1)
  def hess():
    f = partial(log_posterior, X = X, y = y)
    return -1.0 * jacfwd(grad(f))(inpu)
  return hess

n = 1000
p = 5000
key, sub = jrandom.split(key)
hess = hess_test(sub, n , p)
times = timeit.repeat(hess, number = 5, repeat = 5)
print(f"Autodiff: The average time with p = {p} is {np.mean(times): .3f}(+/-{np.std(times): .3f})")

Autodiff: The average time with p = 5000 is  3.222(+/- 0.379)

That doesn’t seem too bad, but the thing is that I know quite a lot about logistic regression. It is, after all, logistic regression. In particular, I know that the Hessian has the form where is a diagonal matrix that has a known form.

This means that the appropriate comparison is between the speed of the autodiff Hessian and how long it takes to compute for some diagonal matrix X.

Now you might be worried here that I didn’t explicitly save and , so the comparison might not be fair. But my friends, I have good news! All of that awkward key, sub = jrandom.split(key) malarkey has the singular advantage that if I pass the same key into make_data that I used for hess_test, I will get the exact same generated data! So let’s do that. For I’m just going to pick a random matrix. This will give a minimum achievable time for computing the Hessian (as it doesn’t do the extra derivatives to compute properly).

If you look at that code and say but Daniel you used the wrong multiplication operator, you can convince yourself that X * d[:, None] gives the same result as jnp.diag(d) @ X. But it will be faster. And it uses such beautiful¹¹ broadcasting rules.

y, X = make_data(key, n , p)
key, sub = jrandom.split(key)
d = jrandom.normal(sub, shape = (n,))
mm = lambda: X.T @ (X * d[:, None])
times = timeit.repeat(mm, number = 5, repeat = 5)
print(f"Symbolic (minimum possible): The average time with p = {p} is {np.mean(times): .3f}(+/-{np.std(times): .3f})")

Symbolic (minimum possible): The average time with p = 5000 is  0.766(+/- 0.014)

Oh dear. The symbolic derivative¹² is a lot faster.

Speeding this up is going to take a little work. The first thing we can try is to explicitly factor out the linear transformation. Instead of passing in the function , we could pass in such that for some matrix . In our case would have a diagonal Hessian. Let’s convince ourselves of that with a small example. As well as dropping the intercept, I’ve also dropped the prior term.

g = lambda prob: jnp.sum(y * jnp.log(prob) + (1-y) * jnp.log1p(-prob))
key, sub2 = jrandom.split(key)
y, X = make_data(sub2, 5, 3)
b = X @ jnp.ones(3)
D = -1.0 * jacfwd(grad(g))(b)
print(np.round(D, 1))

[[0.7 0.  0.  0.  0. ]
 [0.  3.7 0.  0.  0. ]
 [0.  0.  7.8 0.  0. ]
 [0.  0.  0.  4.9 0. ]
 [0.  0.  0.  0.  0.3]]

Wonderfully diagonal!

def hess2(g, A, x):
  # 
  b = A @ x
  D = -1.0 * jacfwd(grad(g))(b)
  H = A.T @ (A * jnp.diag(D)[:, None])
  return H

y, X = make_data(sub, n, p)
g = lambda prob: jnp.sum(y * jnp.log(prob) + (1-y) * jnp.log1p(-prob))
x0 = jnp.ones(p)
h2 = lambda: hess2(g, X, x0)
times = timeit.repeat(h2, number = 5, repeat = 5)
print(f"Separated Hessian: The average time with p = {p} is {np.mean(times): .3f}(+/-{np.std(times): .3f})")

Separated Hessian: The average time with p = 5000 is  0.975(+/- 0.163)

Well that’s definitely better.

Now, we might be able to do even better than that if we notice that if we know that is diagonal, then we don’t need to compute the entire Hessian, we can simply compute the Hessian-vector product where is the vector of ones. Just as we computed the Hessian by computing the Jacobian of the gradient, it turns out that we can compute a Hessian-vector product by computing a Jacobian-vector product jvp of the gradient. The syntax in JAX is, honestly, a little bit gross here¹³, but if you want to read up about how it works the docs are really nice¹⁴.

This observation is going to be useful because jacfwd computes the Jacobian by computing Jacobian-vector products. So this observation is saving us a lot of work.

from jax import jvp
def hess3(g, A, x):
  # 
  b = A @ x
  D = -1.0 * jvp(grad(g), (b,), (jnp.ones(n),))[1]
  H = A.T @ (A * D[:, None])
  return H

h3 = lambda: hess3(g, X, x0)
times = timeit.repeat(h3, number = 5, repeat = 5)
print(f"Compressed Hessian: The average time with p = {p} is {np.mean(times): .3f}(+/-{np.std(times): .3f})")

Compressed Hessian: The average time with p = 5000 is  0.879(+/- 0.082)

This is very nearly as fast as the lower bound for the symbolic Hessian. There must be a way to use this.

Getting to know jaxprs

A jaxpr is a transformation of the python code for evaluating a JAX function into a human-readable language that maps types primitives through the code. We can view it using the jax.make_jaxpr function.

Let’s look at the log-posterior function after partial evaluation to make it a single-input function.

from jax import make_jaxpr

lp = partial(log_posterior, X=X, y=y)
print(make_jaxpr(lp)(jnp.ones(p+1)))

{ lambda a:f32[1000,5000] b:bool[1000]; c:f32[5001]. let
    d:f32[1] = dynamic_slice[slice_sizes=(1,)] c 0
    e:f32[] = squeeze[dimensions=(0,)] d
    f:f32[5000] = dynamic_slice[slice_sizes=(5000,)] c 1
    g:f32[1000] = dot_general[dimension_numbers=(([1], [0]), ([], []))] a f
    h:f32[1000] = add e g
    i:f32[1000] = logistic h
    j:f32[1000] = log i
    k:f32[1000] = convert_element_type[new_dtype=float32 weak_type=False] b
    l:f32[1000] = mul k j
    m:i32[1000] = convert_element_type[new_dtype=int32 weak_type=True] b
    n:i32[1000] = sub 1 m
    o:f32[1000] = neg i
    p:f32[1000] = log1p o
    q:f32[1000] = convert_element_type[new_dtype=float32 weak_type=False] n
    r:f32[1000] = mul q p
    s:f32[1000] = add l r
    t:f32[] = reduce_sum[axes=(0,)] s
    u:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] c c
    v:f32[] = mul 0.5 u
    w:f32[] = sub t v
  in (w,) }

This can be a bit tricky to read the first time you see it, but it’s waaaay easier that X86-Assembly or the LLVM-IR. Basically it says that to compute lp(jnp.ones(p+1)) you need to run through this program. The first line gives the inputs (with types and shapes). Then after the let statement, there are a the commands that need to be executed in order. A single execution looks like

d:f32[1] = dynamic_slice[slice_sizes=(1,)] c 0

This can be read as take a slice of vector c starting at 0 of shape (1,) and store it in d, which is a 1-dimensional 32bit float array. (The line after turns it into a scalar.)

All of the other lines can be read similarly. A good trick, if you don’t recognize the primitive¹⁶, is to look it up in the jax.lax sub-module.

Even a cursory read of this suggests that we could probably save a couple of tedious operations by passing in an integer y, rather than a Boolean y, but hey. That really shouldn’t cost much.

While the jaxpr is lovely, it’s a whole lot easier to reason about if you see it graphically. We can plot the expression graph using¹⁷ the haiku¹⁸ package from DeepMind.

from haiku.experimental import to_dot
import graphviz
import re
f = partial(log_posterior, X = X, y = y)
dot = to_dot(f)(jnp.ones(p+1))
#Strip out an obnoxious autogen title
dot = re.sub("<<.*>>;","\" \"", dot, count = 1, flags=re.DOTALL)
graphviz.Source(dot)

To understand this graph, the orange-y boxes represent the input for lp. In this case it’s an array of floating point digits with . The purple boxes are constants that are used in the function. Some of these are signed integers (s32), there’s a matrix (f32[1000, 5000]), and there is even a literal (0.5). The blue box is the output. That leaves the yellow boxes, which have all of the operations, with inward arrows indicating the inputs and outward arrows indicating the outputs.

Splitting the expression graph into linear and non-linear subgraphs

Looking at the graph, we can split it into three sub-graphs. The first sub-graph can be found by tracing an input value through the graph until it hits either a non-linear operation or the end of the graph. The sub-graph is created by making the penultimate node in that sequence an output node. This sub-graph represents a linear transformations.

Once we have reached the end of the linear portion, we can link the output from this operation to the input of the non-linear sub-graph.

Finally, we have one more trace of through the graph that is non-linear. We could couple this into the non-linear graph at the cost of having to reason about a bivariate Hessian (which will become complex).

The two non-linear portions of the graph are merged through a trivial linear combination.

Step right up to play the game of the year: Is it linear?

So we need to trace through these jaxprs and keep a record of which of the sub-graphs they are in (and we do not know how many sub-graphs there will be!). We also need to note if an operation is linear or not. This is not something that is automatically provided. We need to store this information ourselves.

The only way I can think to do this is to make a set of all of the JAX operations that I know to be linear. Many of them are just index or type stuff. Unfortunately, there is a more complex class of operation, which are only sometimes linear.

The first example we see of this is

g:f32[1000] = dot_general[
      dimension_numbers=(((1,), (0,)), ((), ()))
      precision=None
      preferred_element_type=None
    ] a f

This line represents the general tensor dot product between a and f. In this case, a is constant input (the matrix ) while f is a linear transformation of the input (beta[1:]), so the resulting step is linear. However, there is a second dot_general in the code, which occurs at

u:f32[] = dot_general[
      dimension_numbers=(((0,), (0,)), ((), ()))
      precision=None
      preferred_element_type=None
    ] c c

In this case, c is a linear transformation of the input (it’s just beta), but dot(c,c) is a quadratic function. Hence in this case, dot_general is not linear.

We are going to need to work out how to handle this case. In the folded code is a partial¹⁹ list of the jax.lax primitives that are linear or occasionally linear. All in all there are 69 linear or no-op primitives and 7 sometimes linear primitives.

jax_linear = {
  'add',
  'bitcast_convert_type',
  'broadcast',
  'broadcast_in_dim',
  'broadcast_shapes',
  'broadcast_to_rank',
  'clz',
  'collapse',
  'complex',
  'concatenate',
  'conj',
  'convert_element_type',
  'dtype',
  'dtypes',
  'dynamic_slice',
  'expand_dims',
  'full',
  'full_like',
  'imag',
  'neg',
  'pad',
  'padtype_to_pads',
  'real',
  'reduce',
  'reshape',
  'rev',
  'rng_bit_generator',
  'rng_uniform',
  'select',
  'select_n',
  'squeeze',
  'sub',
  'transpose',
  'zeros_like_array',
  'GatherDimensionNumbers',
  'GatherScatterMode',
  'ScatterDimensionNumbers',
  'dynamic_index_in_dim',
  'dynamic_slice',
  'dynamic_slice_in_dim',
  'dynamic_update_index_in_dim',
  'dynamic_update_slice',
  'dynamic_update_slice_in_dim',
  'gather',
  'index_in_dim',
  'index_take',
  'reduce_sum',
  'scatter',
  'scatter_add',
  'slice',
  'slice_in_dim',
  'conv',
  'conv_dimension_numbers',
  'conv_general_dilated',
  'conv_general_permutations',
  'conv_general_shape_tuple',
  'conv_shape_tuple',
  'conv_transpose',
  'conv_transpose_shape_tuple',
  'conv_with_general_padding',
  'cumsum',
  'fft',
  'all_gather',
  'all_to_all',
  'axis_index',
  'ppermute',
  'pshuffle',
  'psum',
  'psum_scatter',
  'pswapaxes',
  'xeinsum'
}

jax_sometimes_linear = { 
  'batch_matmul',
  'dot',
  'dot_general',
  'mul'
 }
jax_first_linear = {
  'div'
 }
jax_last_linear = {
  'custom_linear_solve',
  'triangular_solve',
  'tridiagonal_solve'
 }

All of the sometimes linear operations are linear as long as only one of their arguments depends on the function inputs. For both div and the various linear solves, the position of the input-dependent argument is restricted to one of the two positions.

Tracing through the jaxprs

In order to split our graph into appropriate sub-graphs we need to trace through the jaxpr and keep track of every variable and if it depends on linear or non-linear parts.

For simplicity, consider the following expression graph for computing lambda x, y: 0.5*(x+y).

This figure corresponds roughly to the jaxpr

{ lambda ; a:f32[] b:f32[]. let c:f32[] = add a b; d:f32[] = mul 0.5 c in (d,) }

For each node, the graph tells us

• its unique identifier (internally²⁰ JAX uses integers)
• which equation generated the value
• which nodes are its parents in the graph (the input(s) to the equation)
• whether or not this node depends on the inputs. This is useful for ignoring non-linearities that just apply to the constants bound to the jaxpr.

We can record this information in a dataclass.

import dataclasses as dc
@dc.dataclass
class Node:
  number: int = None
  eqn: int = None
  parents: List[int] = dc.field(default_factory=list)
  depends_on_input: bool = True

Now we can build up our graph with all of the side information we need. The format of a jaxpr places the constant inputs in the first node, followed by the non-constant inputs (which I’m calling the input variables). For simplicity, I am assuming that there is only one input variable.

import jax.core as jcore
from jax import make_jaxpr

jpr = make_jaxpr(lp)(jnp.ones(p+1))

node_list = {
  const.count: Node(
    number=const.count, 
    depends_on_input=False
  ) for const in jpr.jaxpr.constvars
}

node_list |= {
  inval.count: Node(number=inval.count) 
  for inval in jpr.jaxpr.invars
}

## For later, we need to know the node numbers that correspond
## to the constants and inputs

consts_and_inputs = {node.number for node in node_list.values()}

node_list |= {
  node.count: Node(
    number=node.count,
    eqn=j,
    parents=[
      invar.count for invar in eqn.invars if not isinstance(invar, jcore.Literal)
    ],
  )
  for j, eqn in enumerate(jpr.jaxpr.eqns)
  for node in eqn.outvars
}

for node in node_list.values():
  if len(node.parents) > 0:
    node.depends_on_input =  any(
      node_list[i].depends_on_input for i in node.parents
    )

node_list

{0: Node(number=0, eqn=None, parents=[], depends_on_input=False),
 1: Node(number=1, eqn=None, parents=[], depends_on_input=False),
 2: Node(number=2, eqn=None, parents=[], depends_on_input=True),
 3: Node(number=3, eqn=0, parents=[2], depends_on_input=True),
 4: Node(number=4, eqn=1, parents=[3], depends_on_input=True),
 5: Node(number=5, eqn=2, parents=[2], depends_on_input=True),
 6: Node(number=6, eqn=3, parents=[0, 5], depends_on_input=True),
 7: Node(number=7, eqn=4, parents=[4, 6], depends_on_input=True),
 8: Node(number=8, eqn=5, parents=[7], depends_on_input=True),
 9: Node(number=9, eqn=6, parents=[8], depends_on_input=True),
 10: Node(number=10, eqn=7, parents=[1], depends_on_input=False),
 11: Node(number=11, eqn=8, parents=[10, 9], depends_on_input=True),
 12: Node(number=12, eqn=9, parents=[1], depends_on_input=False),
 13: Node(number=13, eqn=10, parents=[12], depends_on_input=False),
 14: Node(number=14, eqn=11, parents=[8], depends_on_input=True),
 15: Node(number=15, eqn=12, parents=[14], depends_on_input=True),
 16: Node(number=16, eqn=13, parents=[13], depends_on_input=False),
 17: Node(number=17, eqn=14, parents=[16, 15], depends_on_input=True),
 18: Node(number=18, eqn=15, parents=[11, 17], depends_on_input=True),
 19: Node(number=19, eqn=16, parents=[18], depends_on_input=True),
 20: Node(number=20, eqn=17, parents=[2, 2], depends_on_input=True),
 21: Node(number=21, eqn=18, parents=[20], depends_on_input=True),
 22: Node(number=22, eqn=19, parents=[19, 21], depends_on_input=True)}

Now let’s identify which equations are linear and which aren’t.

linear_eqn =[False] * len(jpr.jaxpr.eqns)

for node in node_list.values():
  if node.eqn is None:
    continue

  prim = jpr.jaxpr.eqns[node.eqn].primitive.name
  
  if prim in jax_linear:
    linear_eqn[node.eqn] = True
  elif prim in jax_sometimes_linear:
    # this is a check for being called once
    linear_eqn[node.eqn] = (
      sum(
        node_list[i].depends_on_input for i in node.parents
      ) == 1
    )
  elif prim in jax_first_linear:
    linear_eqn[node.eqn] = (
      node_list[node.parents[0]].depends_on_input 
      and not any(node_list[pa].depends_on_input for pa in node.parents[1:])
    )
  elif prim in jax_last_linear:
    linear_eqn[node.eqn] = (
      node_list[node.parents[-1]].depends_on_input 
      and not any(node_list[pa].depends_on_input for pa in node.parents[:-1])
    )
  elif all(not node_list[i].depends_on_input for i in node.parents):
    linear_eqn[node.eqn] = True # Constants are linear

The only messy thing²¹ in here is dealing with the sometimes linear primitives. If I was sure that every JAX primitive was guaranteed to have only two inputs, this could be simplified, but sadly I don’t know that.

Partitioning the graph

Now it’s time for the fun: partitioning the problem into sub-graphs. To do this, we need to think about what rules we want to encode.

The first rule is that every input for an equation or sub-graph needs to be either a constant, the function input, or the output of some other sub-graph that has already been computed. This means that if we find an equation with an input that doesn’t satisfy these conditions, we need to split the sub-graph that it’s in into two sub-graphs.

The second rule is the only exception to the first rule. A sub-graph can have inputs from non-linear sub-graphs if an only if it contains a sequence of sum or sub terms and it finishes with the terminal node. This covers the common case where the function we are taking the Hessian of is a linear combination of independent functions. For instance, log_posterior(beta) = log_likelihood(beta) + log_prior(beta). In this case we can compute the Hessians for the non-linear sub-expressions and then combine them.

The third rule is that every independent use of the function input is the opportunity to start a new tree. (It may merge with a known tree.)

And that’s it. Should be simple enough to implement.

I’m feeling like running this bad boy backwards, so let’s do that. One of the assumption we have made is that the function we are tracing has a single output and that is always in the last node and defined in the last equation. So first off, lets get our terminal combination expressions.

## Find the terminal combination expressions
terminal_expressions = {"sum", "sub"}
comb_eqns = []
for eqn in jpr.jaxpr.eqns[::-1]:
  if any(
    node_list[a.count].depends_on_input 
    for a in eqn.invars 
    if not isinstance(a, jcore.Literal)
  )  and (
    eqn.primitive.name in terminal_expressions
  ):
    comb_eqns.append(eqn)
  else:
    break

print(comb_eqns)

[a:f32[] = sub b c]

Now for each of the terminal combination expressions, we will trace their parent back until we run out of tree. While we are doing this, we can also keep track of runs of linear operations. We also have to visit each equation once, so we need to keep track of our visited equations. This is, whether we like it or not, a depth-first search. It’s always a bloody depth-first search, isn’t it.

So what we are going to do is go through each of the combiner nodes and trace the graph down from it and note the path and it’s parent. If we run into a portion of the graph we have already traced, we will note that for later. These paths will either be merged or, if the ancestral path from that point is all linear, will be used as a linear sub-graph.

def dfs(visited, graph, subgraph, to_check, node):
  if node in visited:
    to_check.add(node)
  else:
    visited.add(node)
    subgraph.add(graph[node].eqn)
    for neighbour in graph[node].parents:
      dfs(visited, graph, subgraph, to_check, neighbour)
  

visited = consts_and_inputs
to_check = set()
subgraphs = []
for ce in comb_eqns:
  for v in (a for a in ce.invars if not isinstance(a, jcore.Literal)):
    if v.count not in visited:
      subgraphs.append(set())
      dfs(visited, node_list, subgraphs[-1], to_check, v.count)

to_check = to_check.difference(consts_and_inputs)
print(f"Subgraphs: {subgraphs}")
print(f"Danger nodes: {to_check}")

Subgraphs: [{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, {17, 18}]
Danger nodes: set()

The to_check nodes are only dangerous insofar as we need to make sure that if they are in one of the linear sub-graphs they are terminal nodes of a sub-graph. To that end, let’s make the linear sub-graphs.

linear_subgraph = []
nonlin_subgraph = []
n_eqns = len(jpr.jaxpr.eqns)
for subgraph in subgraphs:
  print(subgraph)
  split = next(
    (
      i for i, lin in enumerate(linear_eqn) 
      if not lin and i in subgraph
    )
  )
  if any(chk in subgraph for chk in to_check):
    split = min(
      split, 
      min(chk for chk in to_check if chk in subgraph)
    )

  linear_subgraph.append(list(subgraph.intersection(set(range(split)))))
  nonlin_subgraph.append(list(subgraph.intersection(set(range(split, n_eqns)))))

print(f"Linear subgraphs: {linear_subgraph}")
print(f"Nonlinear subgraphs: {nonlin_subgraph}")

{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}
{17, 18}
Linear subgraphs: [[0, 1, 2, 3, 4], []]
Nonlinear subgraphs: [[5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], [17, 18]]

The only interesting thing here is making sure that if there is a linear node in the graph that was visited twice, it is the terminal node of the linear graph. The better thing would be to actually split the linear graph, but I’m getting a little bit sick of this post and I don’t really want to deal with multiple linear sub-graphs. So I shan’t. But hopefully it’s relatively clear how you would do that.

In this case it’s pretty clear that we are ok.

any(linear_eqn[node_list[j].eqn] for j in to_check)

False

Putting it together

Well that’s a nice script that does what I want. Now let’s put it together in a function. I’m going to give it the very unspecific name transform_jaxpr because sometimes you’ve gotta annoy your future self.

def transform_jaxpr(
  jaxpr: jcore.ClosedJaxpr
) -> Tuple[List[Set[int]], List[Set[int]], List[jcore.JaxprEqn]]:
  assert len(jpr.in_avals) == 1
  assert len(jpr.out_avals) == 1

  from jax import core as jcore

  ## 1. Extract the tree and its relevant behavior
  node_list = {
    const.count: Node(
      number=const.count, 
      depends_on_input=False
    ) for const in jpr.jaxpr.constvars
  }

  node_list |= {
    inval.count: Node(number=inval.count) 
    for inval in jpr.jaxpr.invars
  }

  ## For later, we need to know the node numbers that correspond
  ## to the constants and inputs

  consts_and_inputs = {node.number for node in node_list.values()}

  node_list |= {
    node.count: Node(
      number=node.count,
      eqn=j,
      parents=[
        invar.count for invar in eqn.invars if not isinstance(invar, jcore.Literal)
      ],
    )
    for j, eqn in enumerate(jpr.jaxpr.eqns)
    for node in eqn.outvars
  }

  for node in node_list.values():
    if len(node.parents) > 0:
      node.depends_on_input =  any(
        node_list[i].depends_on_input for i in node.parents
      )

  ## 2. Identify which equations are linear_eqn

  linear_eqn =[False] * len(jpr.jaxpr.eqns)

  for node in node_list.values():
    if node.eqn is None:
      continue

    prim = jpr.jaxpr.eqns[node.eqn].primitive.name
    
    if prim in jax_linear:
      linear_eqn[node.eqn] = True
    elif prim in jax_sometimes_linear:
      # this is a check for being called once
      linear_eqn[node.eqn] = (
        sum(
          node_list[i].depends_on_input for i in node.parents
        ) == 1
      )
    elif prim in jax_first_linear:
      linear_eqn[node.eqn] = (
        node_list[node.parents[0]].depends_on_input 
        and not any(node_list[pa].depends_on_input for pa in node.parents[1:])
      )
    elif prim in jax_last_linear:
      linear_eqn[node.eqn] = (
        node_list[node.parents[-1]].depends_on_input 
        and not any(node_list[pa].depends_on_input for pa in node.parents[:-1])
      )
    elif all(not node_list[i].depends_on_input for i in node.parents):
      linear_eqn[node.eqn] = True # Constants are linear

  ##3. Find all the terminal expressions
  ## Find the terminal combination expressions
  terminal_expressions = {"sum", "sub"}
  comb_eqns = []
  for eqn in jpr.jaxpr.eqns[::-1]:
    if any(
      node_list[a.count].depends_on_input 
      for a in eqn.invars 
      if not isinstance(a, jcore.Literal)
    )  and (
      eqn.primitive.name in terminal_expressions
    ):
      comb_eqns.append(eqn)
    else:
      break
  
  ## 4. Identify the sub-graphs 
  def dfs(visited, graph, subgraph, to_check, node):
    if node in visited:
      to_check.add(node)
    else:
      visited.add(node)
      subgraph.add(graph[node].eqn)
      for neighbour in graph[node].parents:
        dfs(visited, graph, subgraph, to_check, neighbour)
    

  visited = consts_and_inputs
  to_check = set()
  subgraphs = []
  for ce in comb_eqns:
    for v in (a for a in ce.invars if not isinstance(a, jcore.Literal)):
      if v.count not in visited:
        subgraphs.append(set())
        dfs(visited, node_list, subgraphs[-1], to_check, v.count)

  to_check = to_check.difference(consts_and_inputs)

  ## 5. Find the linear sub-graphs
  linear_subgraph = []
  nonlin_subgraph = []
  n_eqns = len(jaxpr.eqns)
  for subgraph in subgraphs:
    split = next(
      (
        i for i, lin in enumerate(linear_eqn) 
        if not lin and i in subgraph
      )
    )
    if any(chk in subgraph for chk in to_check):
      split = min(
        split, 
        min(chk for chk in to_check if chk in subgraph)
      )

    linear_subgraph.append(list(subgraph.intersection(set(range(split)))))
    nonlin_subgraph.append(list(subgraph.intersection(set(range(split, n_eqns)))))
  
  return (linear_subgraph, nonlin_subgraph, comb_eqns)

For one final sense check, let’s compare these outputs to the original jaxpr.

for j, lin in enumerate(linear_subgraph):
  print(f"Linear: {j}")
  for i in lin:
    print(jpr.eqns[i])

for j, nlin in enumerate(nonlin_subgraph):
  print(f"Nonlinear: {j}")
  for i in nlin:
    print(jpr.eqns[i])

print("Combination equations")
for eqn in comb_eqns:
  print(eqn)

Linear: 0
a:f32[1] = dynamic_slice[slice_sizes=(1,)] b 0
a:f32[] = squeeze[dimensions=(0,)] b
a:f32[5000] = dynamic_slice[slice_sizes=(5000,)] b 1
a:f32[1000] = dot_general[dimension_numbers=(([1], [0]), ([], []))] b c
a:f32[1000] = add b c
Linear: 1
Nonlinear: 0
a:f32[1000] = logistic b
a:f32[1000] = log b
a:f32[1000] = convert_element_type[new_dtype=float32 weak_type=False] b
a:f32[1000] = mul b c
a:i32[1000] = convert_element_type[new_dtype=int32 weak_type=True] b
a:i32[1000] = sub 1 b
a:f32[1000] = neg b
a:f32[1000] = log1p b
a:f32[1000] = convert_element_type[new_dtype=float32 weak_type=False] b
a:f32[1000] = mul b c
a:f32[1000] = add b c
a:f32[] = reduce_sum[axes=(0,)] b
Nonlinear: 1
a:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] b b
a:f32[] = mul 0.5 b
Combination equations
a:f32[] = sub b c

Comparing to the original jaxpr, we see it has the same information (the formatting is a bit unfortunate, as the original __repr__ keeps track of the links between things, but what can you do?).

print(jpr)

{ lambda a:f32[1000,5000] b:bool[1000]; c:f32[5001]. let
    d:f32[1] = dynamic_slice[slice_sizes=(1,)] c 0
    e:f32[] = squeeze[dimensions=(0,)] d
    f:f32[5000] = dynamic_slice[slice_sizes=(5000,)] c 1
    g:f32[1000] = dot_general[dimension_numbers=(([1], [0]), ([], []))] a f
    h:f32[1000] = add e g
    i:f32[1000] = logistic h
    j:f32[1000] = log i
    k:f32[1000] = convert_element_type[new_dtype=float32 weak_type=False] b
    l:f32[1000] = mul k j
    m:i32[1000] = convert_element_type[new_dtype=int32 weak_type=True] b
    n:i32[1000] = sub 1 m
    o:f32[1000] = neg i
    p:f32[1000] = log1p o
    q:f32[1000] = convert_element_type[new_dtype=float32 weak_type=False] n
    r:f32[1000] = mul q p
    s:f32[1000] = add l r
    t:f32[] = reduce_sum[axes=(0,)] s
    u:f32[] = dot_general[dimension_numbers=(([0], [0]), ([], []))] c c
    v:f32[] = mul 0.5 u
    w:f32[] = sub t v
  in (w,) }

Making sub-functions

Now that we have the graph partitioned, let’s make our sub-functions. We do this by manipulating the jaxpr and then closing over the literals.

There are a few ways we can do this. We could build completely new JaxprEqn objects from the existing Jaxpr. But honestly, that is just annoying, so instead I’m just going to modify the basic, but incomplete, parser²².

The only modification from the standard eval_jaxpr is that we are explicitly specifying the invars in order to overwrite the standard ones. This relies on the lexicographic ordering of the jaxpr expression graph.

from typing import Callable

from jax import core as jcore
from jax import lax
from jax._src.util import safe_map

def eval_subjaxpr(
  *args,
  jaxpr: jcore.Jaxpr, 
  consts: List[jcore.Literal], 
  subgraph: List[int], 
  invars: List[jcore.Var]
):


  assert len(invars) == len(args)
  
  # Mapping from variable -> value
  env = {}
  
  def read(var):
    # Literals are values baked into the Jaxpr
    if type(var) is jcore.Literal:
      return var.val

    return env[var]

  def write(var, val):
    env[var] = val

  # We need to bind the input to the sub-function
  # to the environment.
  # We only need to write the consts that appear
  # in our sub-graph, but that's more bookkeeping
  safe_map(write, invars, args)
  safe_map(write, jaxpr.constvars, consts)

  # Loop through equations and evaluate primitives using `bind`
  outvars = []
  for j in subgraph:
    eqn = jaxpr.eqns[j]
    # Read inputs to equation from environment
    invals = safe_map(read, eqn.invars)  
    # `bind` is how a primitive is called
    outvals = eqn.primitive.bind(*invals, **eqn.params)
    # Primitives may return multiple outputs or not
    if not eqn.primitive.multiple_results: 
      outvals = [outvals]
    outvars += [eqn.outvars]
    safe_map(write, eqn.outvars, outvals) 
  
  return safe_map(read, outvars[-1])[0]

The final thing we should do is combine our transformation with this evaluation module to convert a function into a sequence of callable sub-functions. I am making liberal uses of lambdas to close over variables that the user should never see (like the sub-graph!). Jesus loves closures and so do I.

def decompose(fun: Callable, *args) -> Tuple[List[Callable], List[Callable], List[jcore.Var]]:
  from functools import partial
  from jax import make_jaxpr

  jpr = make_jaxpr(fun)(*args)
  linear_subgraph, nonlin_subgraph, comb_eqns = transform_jaxpr(jpr)

  assert len(linear_subgraph) == len(nonlin_subgraph)
  assert len(jpr.jaxpr.invars) == 1, "Functions must only have one input"

  def get_invars(sub: List[int]) -> List[jcore.Var]:
    # There is an implicit assumption everywhere in this post 
    # that each sub-function only has one non-constant input
    
    min_count = jpr.jaxpr.eqns[sub[0]].outvars[0].count
    literal_ceil = jpr.jaxpr.invars[0].count
    for j in sub:
      for v in jpr.jaxpr.eqns[j].invars:
        if (
          not isinstance(v, jcore.Literal) and
          v.count >= literal_ceil and 
          v.count < min_count
        ):
          return [v]
    raise Exception("Somehow you can't find any invars")
    

  lin_funs = []
  nlin_funs = []
  nlin_inputs = []
  lin_outputs = []
  nlin_outputs = []

  for lin in linear_subgraph:
    if len(lin) == 0:
      lin_funs += [None]
      lin_outputs += [jpr.jaxpr.invars[0].count]
    elif jpr.jaxpr.eqns[lin[-1]].primitive.multiple_results:
      raise Exception(f"This code doesn't deal with multiple outputs from subgraph {lin}")
    else:
      # find 
      lin_outputs += [jpr.jaxpr.eqns[lin[-1]].outvars[0].count]
      lin_funs += [
        partial(eval_subjaxpr,
          jaxpr = jpr.jaxpr, 
          consts = jpr.literals, 
          subgraph = lin, 
          invars = get_invars(lin)
        )
      ]
      
  for nlin in nonlin_subgraph:
    if len(nlin) == 0:
      nlin_funs += [None]
      nlin_inputs += [-1]
      nlin_outputs += [None]
    elif jpr.jaxpr.eqns[nlin[-1]].primitive.multiple_results:
      raise Exception(f"This code doesn't deal with multiple outputs from subgraph {nlin}")
    else:
      invar = get_invars(nlin)
      nlin_inputs += [lin_outputs.index(invar.count) if invar.count in lin_outputs else -1]
      nlin_outputs += [jpr.jaxpr.eqns[nlin[-1]].outvars[0].count]
      nlin_funs += [
        partial(eval_subjaxpr,
          jaxpr = jpr.jaxpr, 
          consts = jpr.literals, 
          subgraph = nlin, 
          invars = get_invars(nlin)
        )
      ]

  combine = [0.0] * len(linear_subgraph)
  # print(combine)
  for eqn in comb_eqns:
    combine[nlin_outputs.index(eqn.invars[0].count)] += 1.0
    if eqn.primitive.name == "sub":
      combine[nlin_outputs.index(eqn.invars[1].count)] += -1.0
    else:
      combine[nlin_outputs.index(eqn.invars[1].count)] += 1.0


  return lin_funs, nlin_funs, nlin_inputs, combine

But is it faster?

Now let’s take a look at whether we have actually saved any time.

import jax
times_hess = timeit.repeat(lambda: jax.hessian(partial(log_posterior, X = X, y = y))(mode_jax), number = 5, repeat = 5)
print(f"Full Hessian: The average time with p = {p} is {np.mean(times_hess): .3f}(+/-{np.std(times_hess): .3f})")

times_smarter = timeit.repeat(lambda: smarter_hessian(partial(log_posterior, X = X, y = y))(mode_jax), number = 5, repeat = 5)
print(f"Smarter Hessian: The average time with p = {p} is {np.mean(times_smarter): .3f}(+/-{np.std(times_smarter): .3f})")

Full Hessian: The average time with p = 5000 is  3.444(+/- 0.201)
Smarter Hessian: The average time with p = 5000 is  3.569(+/- 0.024)

Well that didn’t make much of a difference. If anything, it’s a little bit slower. This is likely due to compiler operations that can be improved if you just lower the whole thing.

But you forgot the diagonal trick

That said, the decomposition into linear and non-linear parts was not the real source of the savings. If we assume the Hessian of the likelihood is diagonal, then we can indeed do a lot better!

The problem here is that while smarter_hessian worked for any²³ JAX-traceable function, we are now making a structural assumption. In theory, we could go through the JAX primitives and mark all of the ones that would (conditionally) lead to diagonal Hessians, but honestly I kinda want this bit of the post to be done. So I will leave that as an exercise to the interested reader.

def smart_hessian(fun: Callable) -> Callable:
  from jax import jacfwd
  from jax import hessian
  from jax import numpy as jnp
  def hess(*args):
    assert len(args) == 1, "This only works for functions with one input"
    
    lin_funs, nlin_funs, nlin_inputs, combine = decompose(fun, *args)
    n_in = args[0].shape[0]
    part = jnp.zeros((n_in, n_in))

    for lin, nlin, nimp, comb in zip(lin_funs, nlin_funs, nlin_inputs, combine):
      
      if lin is not None:
        lin_val = lin(*args)
        jac = jacfwd(lin)(*args)
      

      h_args = (lin_val,) if lin is not None else args
      D = jvp(grad(nlin), h_args, (jnp.ones_like(h_args[0]),))[1] if nlin is not None else None

      if lin is not None and nlin is not None:

        part += comb * (jac.T @ (jac * D[:,None]))
      elif lin is not None:
        part += comb * jac.T @ jac
      elif nlin is not None:
        part += comb * jnp.diag(D)
      
    return part
  return hess



H_smart = -1.0 * smart_hessian(partial(log_posterior, X = X, y = y))(mode_jax)

print(f"The error is {jnp.linalg.norm(H_jax - H_smart).tolist()}!")

times_smart = timeit.repeat(lambda: smart_hessian(partial(log_posterior, X = X, y = y))(mode_jax), number = 5, repeat = 5)
print(f"Smart (diagonal-aware) Hessian: The average time with p = {p} is {np.mean(times_smart): .3f}(+/-{np.std(times_smart): .3f})")

The error is 2.3684690404479625e-06!
Smart (diagonal-aware) Hessian: The average time with p = 5000 is  2.269(+/- 0.031)

That is a proper saving!

The power of compiler optimizations

I think what I’ve shown here is that one of the really powerful things about compiled languages like JAX is that you can perform a pile of code optimizations that can greatly improve their performance.

In the ideal world, this type of optimization should be invisible to the end user. Were I to do this seriously²⁴, I would make sure that if the assumptions of the optimized code weren’t met, the behaviour would revert back to the standard jax.hessian.

Recognizing when to perform an optimization is, in reality, the whole art of this type of process. And it’s very hard. For this post, I was able to do that to automatically recognize the linear operation, but I didn’t try to find conditions that ensured the Hessian would be diagonal.

Sparsity detection and sparse autodiff

Would you believe that people have spent a lot of time studying the efficiency gains when you have things like sparse Hessians? There is, in fact, a massive literature on sparse autodiff and it is implemented in several autodiff libraries, including in Julia.

Sparsity exploiting autodiff uses symbolic analysis of the expression tree for a function to identify when certain derivatives are going to be zero. For Hessians, it needs to identify when two variables have at most linear dependencies.

Once you have worked out the sparsity pattern, you need to do something with it. In the logistic case, it is diagonal, but in a lot of cases it will depend on more than one element of the latent representation. That is the Hessian will be sparse²⁵, but it won’t be diagonal.

I guess the question is can we generalist the observation if the Hessian is diagonal we only need to compute a single Hessian-vector product to general sparsity structures.

In general, we won’t be able to get away with a single product and will instead need a specially constructed set of probing vectors, where is a number to be determined (that is hopefully much smaller than ). This set of vectors will have the special property that This means that the non-zero elements of each probing vector corresponds to a disjoint grouping of the variables.

To do this, we need to construct our set of probing vectors in a very special way. Each will be a vector containing zeros and ones. The set of indices with have color . The aim is to associate each index with a unique color in such a way that we can recover the algorithm. We can do this with a structurally symmetric orthogonal partition, which is detailed in Section 4 of this great review article.

Implementing²⁶ sparsity-aware autodiff Hessians does require some graph algorithms, and is frankly beyond the scope of my patience here. But it certainly is possible and you would get quite general performance from it.

Critically, because it reduces the computation of a dense Hessian matrix with Hessian-vector products, it is extremely well suited to modern GPU acceleration techniques!

Could we do more?

There are so many many many ways to improve the very simple symbolic reduction of the autodiff beyond the simple “identify $f(Ax)” strategy. For more complex cases, it might be necessary to relax the only one input and only one output assumption.

It also might be possible to chain multiple instances of this, although this would require a more complex Hessian chain rule. Nevertheless, the extra complexity might be balanced by savings form the applicable instances of sparse autodiff.

But probably the thing that actually annoys me in all of this is that we are constantly recomputing the Jacobian for the linear equation, which is fixed. A better implementation would consider implementing symbolic differentiation for linear sub-graphs, which should lead to even more savings.

But is JAX the right framework for this?

All of this was a fair bit of work so I’m tempted to throw myself at the sunk-cost fallacy and just declare it to be good. But there is a problem. Because JAX doesn’t do a symbolic transformation of the program (only a trace through paths associated with specific values), there is no guarantee that the sparsity pattern for remains the same at each step. And there is nothing wrong with that. It’s an expressive, exciting language.

But all of the code transformation to make a sparsity-exploiting Hessian doesn’t come for free. And the idea of having to do it again every time a Hessian is needed is … troubling. If we could guarantee that the sparsity pattern was static, then we could factor all of this complex parsing and coloring code away and just run it once for each problem.

Theoretically, we could do something like hashing on the jaxpr, but I’m not sure how much that would help.

Ideally, we could do this in a library that performs symbolic manipulations and can compile them into an expression graph. JAX is not quite²⁷ that language. An option for this type of symbolic manipulation would be Aesara. It may even be possible to do it in Stan, but even my wandering mind doesn’t want to work out how to do this in OCaml.

Continuous distributions in 1D

If and are both continuous, univariate distributions, it is pretty easy to construct a transport map. In particular, if is the cumulative distribution function of , then is a transport map. This works because, if , then . From this, we can use everyone’s favourite result that you can sample from a continuous univariate random variable by evaluating the quantile function at a uniform random value.

There are, of course, two problems with this: it only works in one dimension and we usually don’t know explicitly.

The second of these isn’t really a problem if we are willing to do something splendifferously dumb. And I am. Because I’m gay and frivolous⁷.

If I write then I can differentiate this to get This is a very non-linear differential equation. We can make it even more non-linear differential equation by repeating the procedure to get Noting that we get This is a rubbish differential equation, but it has the singular advantage that it doesn’t depend⁸ on the normalising constant for , which can be useful. The downside is that the boundary conditions are infinite on both ends.

Regardless of that particular challenge, we could use this to build a generic algorithm.

1. Sample

2. Use a numerical differential equation solver to solve the equation with boundary conditions for some sufficiently large number and return

This will sample from truncated to .

I was going to write some python code to do this, but honestly it hurts my soul. So I shan’t.

Transport maps: A less terrible method that works on general densities

Outside of one dimension, there is (to the best of my knowledge) no direct solution to the transport problem. That means that we need to solve our own. Thankfully, the glorious Youssef Marzouk and a bunch of his collaborators have spent some quality time mapping out this idea. A really nice survey of their results can be found in this paper.

Essentially the idea is that we can try to find the most convenient transport map available to us. In particular, it’s useful to minimise the Kullback-Leibler divergence between and its transport. After a little bit⁹ of maths, this is equivalent to minimising where is the Jacobian of . To finish the specification of the optimisation problem, it’s enough to consider triangular maps¹⁰ with the additional constraint that their Jacobians have positive determinants. Using a triangular map has two distinct advantages: it’s parsimonious and it makes the positive determinant constraint much easier to deal with. Triangular maps are also sufficient for the problem (my man Bogachev showed it in 2005).

That said, this can be a somewhat tricky optimisation problem. Youssef and his friends have spilt a lot of ink on this topic. And if you’re the sort of person who just fucking loves a weird optimisation problem, I’m sure you’ve got thoughts. With and without the triangular constraint, this can be parameterised as the composition of a sequence of simple functions, in which case you turn three times and scream neural net and a normalising¹¹ flow appears.

What if we only have samples from the target density

All of that is very lovely. And quite nice in its context. But what happens if you don’t actually have access of the (unnormalised) log density of the target? What if you only have samples?

The good news is that you’re not shit out of luck. But it’s a bit tricky. And once again, that lovely review paper by Youssef and friends will tell us how to do it.

In particular, they noticed that if you swap the direction of the KL divergence, you get the optimisation problem for the inverse mapping that aims to minimise where is once again a triangular map subject to the monotonicity constraints Because we have the freedom to choose the reference density , we can choose it to be iid standard normals, in which case we get the optimisation problem which is a convex, separable optimisation problem that can be solved¹² using, for instance, a stochastic gradient method. This can be turned into an unconstrained optimisation problem by explicitly parameterising the monotonicity constraint.

The monotonicity of makes the resulting nonlinear solve to compute relatively straightforward. In fact, if isn’t too big you can solve this sequentially dimension-by-dimension. But, of course, when you’ve got a lot of parameters this is a poor method and it would make more sense¹³ to attack it with some sort of gradient descent method. It might even be worth taking the time to learn the inverse function so that can be applied for, essentially, free.

So does it work?

To some extent, the answer is yes. This is very much normalising flows in its most embryonic form. They work to some extent. And this presentation makes some of the problems fairly obvious:

1. There’s no real guarantee that is going to be a nice smooth map, which means that we may have problems moving beyond the training sample.

2. The most natural way to organise the computations are naturally sequential involving sweeps across the parameters. This can be difficult to parallelise efficiently on modern architectures.

3. The complexity of the triangular map is going to depend on the order of variables. This is fine if you’re processing something that is inherently sequential, but if you’re working with image data, this can be challenging.

Of course, there are a pile of ways that these problems can be overcome in whole or in part. I’d point you to the last five years of ML conference papers. You’re welcome.

A very quick introduction to inverse problems

It turns out that learning parameters of differential equation (and other physical models) has a long and storied history in applied mathematics under the name of inverse problems. If that sounds like statistics, you’d be right. It’s statistics, except with no interest in measurement or, classically, uncertainty.

The classic inverse problem framing involves a forward map that takes as its input some parameters (often a function) and returns the full state of a system (often another function). For instance, the forwards map could be the solution of a partial differential equation like The thing that you should notice about this is that the forward map is a) possibly expensive to compute, b) not explicitly known, and c) extremely¹⁷ non-linear.

The problem is specified with data points and the aim is to find the value of that best fits the data. The traditional choice is to minimise the mean-square error

Now every single one of you will know immediately that this question is both vague and ill-posed. There are many functions that will fit the data. This means that we need to enforce¹⁸ some sort of complexity penalty on . This leads to the method known as Tikhonov regularisation¹⁹ where is some Banach space and is some tuning parameter.

As you can imagine, there’s a lot of maths under this about when there is a unique minimum, how the reconstruction behaves as and , and how the choice of effects the estimation of . There is also quite a lot of work²⁰ looking at how to actually solve these sorts of optimisation problems.

Eventually, the field evolved and people started to realise that it’s actually fairly important to quantify the uncertainty in the estimate. This is … tricky under the Tikhonov regularlisation framework, which became a big motivation for Bayesian inverse problems.

As with all Bayesianifications, we just need to turn the above into a likelihood and a prior. Easy. Well, the likelihood part, at least, is easy. If we want to line up with Tikhonov regularisation, we can choose a Gaussian likelihood

This is familiar to statisticians, the forward model is essentially working as a non-standard link function in a generalised linear model. There are two big practical differences. The first one is that is very non-linear and almost certainly not monotone. The second problem is that evaluations of are typically very²¹ expensive. For instance, you may need to solve a system of differential equations. This means that any computational method²² is going to need to minimise the number of likelihood evaluations.

The choice of prior on can, however, be a bit tricky. The problem is that in most traditional inverse problems is a function²³ and so we need to put a carefully specified prior on it. And there is a lot of really interesting work on what this means in a Bayesian setting. This is really the topic for another blogpost, but it’s certainly an area where you need to be aware of the limitations of different high-dimensional priors and how they perform in various contexts. For instance, if the function you are trying to reconstruct is likely to have a lot of sharp boundaries²⁴ then you need to make sure that your prior can support functions with sharp boundaries. My little soldier bois²⁵ don’t, so you need to get more²⁶ creative.

The likelihood for a normalising flow

Our aim now is to cast the normalising flow idea into the inverse problems framework. To do this, we remember that we begin our flow from a sample from and we then deform it until it becomes a sample from at some known time (which I’m going to choose as ). This means that if , then

We can now derive a relationship between and using the change of variables formula. In particular, which means that our log likelihood will be

The log-determinant term looks like it might cause some trouble. If is parameterised as a triangular map it can be written explicitly, but there is, of course, another route.

For notational ease, let’s consider , for some . Then We can differentiate this with respect to to get ²⁷ to get where I used one of those magical vector calculus identities to get that trace. Remembering that , the log-determinant of the Jacobian at zero is zero and so we get the initial condition

The likelihood can be evaluated²⁸ by solving the system of differential equations and the log likelihood is evaluated as

It turns out that you can take gradients of the log-likelihood efficiently by solving an augmented system of differential equations that’s twice the size of the original. This allows for all kinds of gradient-driven inferential shenanigans.

But oh that complexity

One big problem with normalising flows as written is that we only have two pieces of information about the entire trajectory :

1. we know that , and

2. we know that .

We know absolutely nothing about outside of those boundary conditions. This means that our model for basically gets to freestyle in those areas.

We can avoid this to some extent by choosing appropriate neural network architectures and/or appropriate penalties in the classical case or priors in the Bayesian case. There’s a whole mini-literature on choosing appropriate penalties.

Just to show how complex it is, let me quickly sketch what Finlay etc suggest as a way to keep the dynamics as boring as possible in the information desert. They lean into the literature on optimal transport theory to come up with the double penalty where the first term minimises the kinetic energy and, essentially, finds the least exciting path from to , while the second term ensures that the Jacobian of doesn’t get too big²⁹, which means that the mapping doesn’t have many sharp changes. Both of these penalty terms are designed to both aid generalisation and to make sure the differential equation isn’t unnecessarily difficult for a ODE solver.

A slightly odd feature of these penalties is that they are both data dependent. That suggests that a prior would, probably, require an amount of work. This is work that I don’t feel like doing today. Especially because this blog post isn’t about bloody normalising flows.

Diffusions and stochastic differential equations

Diffusions are to applied mathematicians what gaffer tape is to³³ a roadie. They are a ubiquitous, convenient, and they hold down the fort when nothing else works.

There are a number of diffusions that are familiar in statistics and machine learning. The most famous one is probably the Langevin diffusion which is asymptotically distributed according to . This forms the basis of a bunch of MCMC methods as well as some faster, less adjusted methods.

But that is not the only diffusion. Today’s friend is the Ornstein-Uhlenbeck (OU) process, which is a Gaussian process that The OU process can be thought of as a mean-reverting Brownian motion. As such, it has continuous but nowhere differentiable sample paths

The stationary distribution of is , where is the identity matrix. In fact, if we start the diffusion at stationarity by setting then X_t is a stationary Gaussian process with covariance function

More interestingly in our context, however, is what happens if we start the diffusion from a fixed point , that will eventually be a sample from . In that case, we can solve the linear stochastic differential equation exactly to get where the integral on the right hand side can be interpreted³⁴ as a white noise integral and so and the variance is From these equations, we see that the mean of the diffusion hurtles exponentially fast towards zero and the variance moves at the same speed towards .

More importantly, this means that, given a starting point , we can generate data from any part of the diffusion ! If we want a sequence of observations from the same trajectory, we can generate them sequentially using the fact that and OU process is a Markov³⁵ process. This means that we are no longer limited to information at just two points along the trajectory.

Reversing the diffusion

So far, there is nothing to learn here. The OU process has a known drift and variance, so everything is splendid. It’s even easy to simulate from. The challenge pops up when we try to reverse the diffusion, that is, when we try to remove noise from a sample rather than add noise to it.

In some sense, this shouldn’t be too disgusting. A diffusion is a Markov process and, if we run the Markov process back in time, we still get a Markov process. In fact, we are going to get another diffusion process.

The twist is that the new diffusion process is going to be quite a bit more complex than the original one. The problem is that unless comes from a Gaussian distribution, this process will be non-Gaussian, and thus somewhat tricky to find the reverse trajectory of.

To see this, consider and recall that and The first two terms in that integrand are Gaussian densities and thus their product is a bivariate Gaussian density Unfortunately, as is not Gaussian, the marginal distribution will be non-Gaussian. This means that our reverse time transition density is also going to be very non-linear.

In order to work out a stochastic differential equation that runs backwards in time and generates the same trajectory, we need a little bit of theory on how the unconditional density and the transition density evolves in time (here and everywhere st). These are related through the Kolmogorov equations.

To introduce these, we need to briefly consider the more general diffusion for nice³⁶ vector/matrix-valued functions and . Kolmogorov showed that the unconditional density evolves according the the partial differential equation subject to the initial condition This is known as Kolmogorov’s forward equation or the Fokker-Planck equation.

The other key result is about the density of conditioned on some future value , . We write this density as and it satisfies the partial differential equation subject to the terminal condition This is known as the Kolmogorov backward equation. Great names. Beautiful names.

Let’s consider a differential equation for the joint density Going ham with the product rule gives The first-order derivatives simplify, using the product rule, to

Staring at this for a moment, we notice that this looks has the same structure as the first-order term on the forward equation. In that case, the second-order term would be

If we notice that and we can re-write the second-order derivative terms in Equation 1 as

This is almost, but not quite, what we want. We are a single minus sign away. Remembering that we probably don’t want it to turn up in any derivatives³⁷. To this end, let’s make the substitution With this substitution the second order terms are where

If we write we get the joint PDE

In order to identify the reverse time diffusion, we are going to find the reverse time backward equation, which confusingly, is for As is a constant in both and , we can divide both sides of Equation 2 by it to get where again and and are known.

This is the forward Kolmogorov equation for the time-reversed³⁸ diffusion where is another white nose. Anderson (1982) shows how to connect the white noise that’s driving the forward dynamics with the white noise that’s driving the reverse dynamics , but that’s overkill for our present situation.

In the context of an OU process, we get the reverse equation where time runs backwards and I’ve used the formula for the logarithmic derivative.

Unlike the forward process, the reverse process is the solution to a non-linear stochastic differential equation. In general, this cannot be solved in closed form and we need to use a numerical SDE solver to generate a sample.

It’s worth noting that the OU process is an overly simple cartoon of a diffusion model. In practice, is usually an increasing function of time so the system injects more noise as the diffusion moves along. This changes some of the exact equations slightly, but you can still sample analytically for any (as long as you choose a fairly simple function for ). There is a large literature on these choices and, to be honest, I can’t be bothered going through them here. But obviously if you want to implement a diffusion model yourself you should look this stuff up.

Estimating the score

The reverse dynamics are driven by the score function Typically, we do not know the density and while we could solve the forward equation in order to estimate it, that is wildly inefficient in high dimensions.

If we can assume that for each , is approximately , then the resulting reverse diffusion is linear In this case is Gaussian with a mean and covariance that has closed form solution in terms of and (perhaps after some numerical quadrature and matrix exponentials).

Unfortunately, as discussed above this is not true. A better approximation would be a mixture of Gaussians but, in general, we can use any method to approximate There are no particular constraints on it, except we expect it to be fairly smooth³⁹ in both and . Hence, we can just learn the score.

As we are going to solve the SDE numerically, we only need to estimate the score at a finite set of locations. In every application that I’ve seen, these are pre-specified, however it would also be possible to use a basis function expansion to interpolate to arbitrary time points. But, to be honest, I think every single example I’ve seen just uses a regularly spaced grid.

So how do we estimate ? Well just like every other situation, we need to define a likelihood (or, I guess, an optimisation criterion). One way to think about this would be to note that you’ll never perfectly recover the initial signal. This is because we need to solve a non-linear stochastic partial differential equation and there will, inherently, be noise in that solution. So instead, assume that we have an initial sample and that after solving the backward equation we have an unbiased estimator of with standard deviation , where is the number of time steps. We know a lot about how the error of SDE solvers scale with and so we can use that to set an appropriate scale for . For instance, if you’re using the Euler–Maruyama method, then it has strong order and would likely be an appropriate scaling.

This strongly suggests a likelihood that looks like where is the estimate of you get by running the reverse diffusion conditioned on , where is an exact sample at time from the forward diffusion started at .

This is the key to the success of diffusion models: given our training sample , we generate new data and we can generate as much of that data as we want. Furthermore, we can choose any set of s we want. We can sample a single pair multiple times or we can look at a diversity of sampling data.

We can even try to recover an intermediate state from information about a future state , . This gives us quite the opportunity to target our learning to areas of the space where we have relatively poor estimates of the score function.

Of course, that’s not what people do. They do stochastic gradient descent to minimise possibly subject to some penalties on . In fact, the distribution on is usually a discrete uniform. As with any sufficiently complex task, there is a lot of detailed work on exactly how to best parameterise, solve, and evaluate this optimisation procedure.

Generating samples

Once the model is trained and we have an estimate of the score function, we can generate new samples by first sampling and running the reverse diffusion starting from for some sufficiently large . One of the advantages of using a variant of the OU process with a non-constant is that we can choose to be smaller. Nevertheless, there will always be a little bit of error introduced by the fact that is only approximately . But really, in the context of all of the other errors, this one is pretty small.

Anyway, run the diffusion backwards and if you’ve estiamted well for the entire trajectory, you will get something that looks a lot like a new sample from .

White noise and its associated things

Thankfully I’ve covered this in a previous blog. The engineering definition of white noise as a GP such that for every , is an iid random variable is not good enough for our purposes. Such a process is hauntingly irregular⁷ and it’s fairly difficult to actually do anything with it. Instead, we consider white noise as a random function defined on the subsets of our domain. This feels like it’s just needless technicality, but it turns out to actually be very very useful.

This doesn’t feel like we are helping very much because how on earth am I going to define the product ? Well the answer, you may be shocked to discover, requires a little bit more maths. We need to define an integral, which turns out to not be shockingly difficult to do. The trick is to realise that if I have an indicator function then¹² In that calculation, I just treated like I would any other measure. (If you’re more of a probability type of girl, it’s the same thing as noticing .)

We can extend the above by taking the sum of two indicator function where and are disjoint and and are any real numbers. By the same reasoning above, and using the linearity of the integral, we get that where the last line follows by doing the ordinary¹³ integral of .

It turns out that every interesting function can be written as the limit of piecewise constant functions¹⁴ and we can therefore define for any function¹⁵

With this notion in hand, we can finally define the action of an operator on white noise.

The generalised Gaussian process

One of those inconvenient things that you may have noticed from above is that is not going to be a function. It is going to be a measure or, as it is more commonly known, a generalised Gaussian process. This is the GP analogue of a generalised function and, as such, only gives an actual value when you integrate it against some sufficiently smooth function.

In order to separate this out from the ordinary GP , we will write it as These two ideas coincide in the special case where which will occur when smooths the white noise sufficiently. In all of the cases we really care about today, this happens. But there are plenty of Gaussian processes that can only be considered as generalised GPs¹⁶

Approximating GPs when is a differential operator

This type of construction for is used in two different situations: kernel convolution methods directly use the representation, and the SPDE methods of Lindgren, Lindström and Rue use it indirectly.

I’m interested in the SPDE method, as it ties into today’s topic. Also because it works really well. This method uses a slightly modified version of the above equation where is the (left) inverse of . I have covered this method in a previous post, but to remind you the SDPE method in its simplest form involves three steps:

1. Approximate for some set of weights and a set of deterministic functions that we are going to use to approximate the GP

2. Approximate¹⁷ the test function for some set of deterministic weights

3. Plug these approximations into the equation to get the equation

As this has to be true for every vector , this is equivalent to the linear system where and .

Obviously this method is only going to be useful if it’s possible to compute the elements of and efficiently. In the special case where is a differential operator¹⁸ and the basis functions are chosen to have compact support¹⁹, these calculations form the basis of the finite element method for solving partial differential equations.

The most important thing, however, is that if is a differential operator and the basis functions have compact support, the matrix is sparse and the matrix can be made²⁰ diagonal, which means that has a sparse precision matrix. This can be used to make inference with these GPs very efficient and is the basis for GPs in the INLA software.

A natural question to ask is when will we end up with a sparse precision matrix? The answer is not quite when is a differential operator. Although that will lead to a sparse precision matrix (and a Markov process), it is not required. So the purpose of the rest of this post is to quantify all of the cases where a GP has the Markov property and we can make use of the resulting computational savings.

Rewriting the Markov property I: Splitting spaces

The Markov property defined above is great and everything, but in order to manipulate it, we need to think carefully about the how the domains , and can be used to divide up the space . To do this, we need to basically localise the Markov property to one set of , , . This concept is called a splitting³⁰ of and by .

This emphasizes that we can split our space into three separate components: inside , outside and on the boundary of and the ability to do that for any³³ domain is the key part of the Markov³⁴ property.

A slightly more convenient way to deal with splitting spaces is the case where the we have overlapping sets , that cover the domain (ie ) and the splitting set is their intersection . In this case, the splitting equation becomes I shan’t lie: that looks wild. But it makes sense when you take and , in which case and .

The final thing to add before we can get to business is a way to get rid of all of the annoying s. The idea is to take the intersection of all of the as the splitting space. If we define we can re-write³⁵ the splitting equation as

This gives the following statement of the Markov property.

Rewriting the Markov property II: The dual random field

We are going to fall further down the abstraction rabbit hole in the hope of ending up somewhere useful. In this case, we are going to invent an object that has no reason to exist and we will show that it can be used to compactly restate the Markov property. It will turn out in the next section that it is actually a useful characterization that will lead (finally) to an operational characterisation of a Markovian Gaussian process.

This definition looks frankly a bit wild, but I promise you, we will use it.

The reason for its structure is that it directly relates to the Markov property. In particular, the existence of a dual field implies that, if we have any , then That’s the first thing we need to show to demonstrate the Markov property.

The second part is much easier. If we note that , it follows that

This gives us our third (and final) characterisation of the (second-order) Markov property.

There is probably more to say about dual fields. For instance, the dual of the dual field is the original field. Neat, huh. But really, all we need to do is know that an orthogonal dual field implies a the Markov property. Because next we are going to construct a dual field, which will give us an actually useful characterisation of Markovian GPs.

Building out our toolset with the conjugate GP

In this section, our job is to construct a dual random field. To do this, we are going to exploit the notion of a conjugate³⁹ Gaussian process, which is a generalised⁴⁰ GP such that⁴¹ It is going to turn out that is the random field generated by . The condition that can be assumed a fortiori. What we need to show is that the existence of a conjugate Gaussian process implies that, for all , .

We will return to the issue of whether or not actually exists later, but assuming it does let’s see how it’s associated random field relates to for . While it is not always true that these things are equal, it is always true that We will consider when equality holds in the next section. But first let’s show the inclusion.

The space contains all random variables of the form , where the support of is compact in , which means that it is a positive distance from . That means that, for some , the support of is outside⁴² of . So if we fix that and consider any smooth with support in⁴³ , then, from the definition of the conjugate GP, we have⁴⁴ This means that is perpendicularity to and, therefore, . Now, is defined as the intersection of these spaces, but it turns out that⁴⁵ for any spaces and , This is because and so every function that’s orthogonal to functions in is also orthogonal to functions in . The same goes for . We have shown that and every is in for some . This gives the inclusion

To give conditions for when it’s an actual equality is a bit more difficult. It, maybe surprisingly, involves thinking carefully about the reproducing kernel Hilbert space of . We are going to take this journey together in two steps. First we will give a condition on the RKHS that guarantees that exists. Then we will look at when .

When does exits? or, A surprising time with the reproducing kernel Hilbert space

First off, though, we need to make sure that exists. Obviously⁴⁶ if it exists then it is unique and .

But does it exist? The answer turns out to be sometimes. But also usually. To show this, we need to do something that is, frankly, just a little bit fancy. We need to deal with the reproducing kernel Hilbert space⁴⁷. This feels somewhat surprising, but it turns out that it is a fundamental object⁴⁸ and intrinsically tied to the space .

The reproducing kernel space, which we will now⁴⁹ call because we are using for something else in this section, is a set of deterministic generalised functions , that can be evaluated at functions⁵⁰ as A generalised function if there is a corresponding random variable in that satisfies It can be shown⁵¹ that there is a one-to-one correspondence between and , in the sense that for every there is a unique .

We can use this correspondence to endow with an inner product

So far, so abstract. The point of the conjugate GP is that it gives us an explicit construction of the⁵² mapping . And, importantly for the discussion of existence, if there is a conjugate GP then the RKHS has a particular relationship with .

To see this, let’s assume exists. Then, for each , the generalised function is in because, by the definition of we have that Hence, the embedding is given by .

Now, if we do a bit of mathematical trickery and equate things that are isomorphic, . On its face, that doesn’t make much sense because on the left we have a space of actual functions and on the right we have a space of generalised functions. To make it work, we associate each smooth function with the generalised function defined above.

This make the closure⁵³ of under the norm and hence we have showed that if there is a conjugate GP, then It turns out that if is dense in then that implies that there exists a conjugate function defined through the isomorphism . This is because and is continuous. Hence if we choose then .

We have shown the following.

This is our first step towards making statements about the stochastic process into statements about the RKHS. We shall continue along this road.

You might, at this point, be wondering if that condition ever actually holds. The answer is yes. It does fairly often. For instance, if is a stationary GP with spectral density , the biorthogonal function exists if and only if there is some such that This basically says that the theory we are developing doesn’t work for GPs with extremely smooth sample paths (like a GP with the square-exponential covariance function). This is not a restriction that bothers me at all.

For non-stationary GPs that aren’t too smooth, this will also hold as long as nothing too bizarre is happening at infinity.

But when does ?

We have shown already⁵⁴ that (that last bit with all the complements can be read as “the support of is inside and always more than from the boundary.”). It follows then that This is nice because it shows that is related to the space that is if is a function that is the limit of a sequence of functions with for some , then and every such random variable has an associated .

So, in the sense⁵⁵ of isomorphisms these are equivalent, that is

This means that if we can show that , then we have two spaces that are isomorphic to the same space and use the same isomorphism . This would mean that the spaces are equivalent.

This can also be placed in the language of function spaces. Recall that Hence will be isomorphic to if and only if that is, if and only if every is the limit of a sequence of smooth functions compactly supported within .

This turns out to not always be true, but it’s true in the situations that we most care about. In particular, we get the following theorem, which I am certainly not going to prove.

The second condition is particularly important because it always holds for stationary GPs with as their covariance structure is shift invariant. It’s not impossible to come up with examples of generalised GPs that don’t satisfy this condition, but they’re all a bit weird (eg the “derivative” of white noise). So as long as your GP is not too weird, you should be fine.

At long last, an RKHS characterisation of the Markov property

And with that, we are finally here! We have that is the dual random field to , and we have a lovely characterisation of in terms of the RKHS . We can combine this with our definition of a Markov property for GPs with a dual random field and get that a GP is Markovian if and only if We can use the isomorphism to say that if , , then there is a such that Moreover, this isomorphism is unitary (aka it preserves the inner product) and so Hence, has the Markov property if and only if

Let’s memorialise this as a theorem.

This result is particularly nice because it entirely characterises the RHKS inner product of a Markovian GP. The reason for this is a deep result from functional analysis called Peetre’s Theorem, which states, in our context, that locality implies that the inner product has the form where⁵⁷ are integrable functions and only a finite number of them are non-zero at any point .

This connection between the RKHS and the dual space also gives the following result for stationary GPs.

This follows from the characterisation of the RKHS as having the inner product as where is the Fourier transform of and the fact that a differential operator can is transformed to a polynomial in Fourier space.

Putting this all in terms of

Waaaay back near the top of the post I described a way to write a (generalised) GP in terms of its covariance operator and the white noise process From the discussions above, it follows that the corresponding conjugate GP is given by This means that the RKHS inner product is given by From the discussion above, if is Markovian, then is⁵⁸ a differential⁵⁹ operator.

Building the expression graph

The elimination tree⁹ is a (forest of) rooted tree(s) that compactly represent the non-zero pattern of the Cholesky factor . In particular, the elimination tree has the property that, for any , if and only if there is a path from to in the tree. Or, in the language of trees, if and only if is a descendant of in the tree .

We can describe¹⁰ by listing the parent of each node. The parent node of in the tree is the smallest with .

We can turn this into an algorithm. An efficient version, which is described in Tim Davies book takes about operations. But I’m going to program up a slower one that takes operations, but has the added benefit¹¹ of giving me the column counts for free.

To do this, we are going to walk the tree and dynamically add up the column counts as we go.

To start off, let’s do this in standard python so that we can see what the algorithm look like. The key concept is that if we write as the elimination tree encoding the structure of¹² L[:j, :j], then we can ask about how this tree connects with node j.

A theorem gives a very simple answer to this.

This tells us that the connection between and node is that for each non-zero elements of the th row of , we can walk $ must have a path in from and we will eventually get to a node that has no parent in . Because there must be a path from to in , it means that the parent of this terminal node must be .

As with everything Cholesky related, this works because the algorithm proceeds from left to right, which in this case means that the node label associated with any descendant of is always less than .

The algorithm is then a fairly run-of-the-mill¹³ tree traversal, where we keep track of where we have been so we don’t double count our columns.

Probably the most important thing here is that I am using the full sparse matrix rather than just its lower triangle. This is, basically, convenience. I need access to the left half of the th row of , which is conveniently the same as the top half of the th column. And sometimes you just don’t want to be dicking around with swapping between row- and column-based representations.

import numpy as np

def etree_base(A_indices, A_indptr):
  n = len(A_indptr) - 1
  parent = [-1] * n
  mark = [-1] * n
  col_count = [1] * n
  for j in range(n):
    mark[j] = j
    for indptr in range(A_indptr[j], A_indptr[j+1]):
      node = A_indices[indptr]
      while node < j and mark[node] != j:
        if parent[node] == -1:
          parent[node] = j
        mark[node] = j
        col_count[node] += 1
        node = parent[node]
  return (parent, col_count)

To convince ourselves this works, let’s run an example and compare the column counts we get to our previous method.

from scipy import sparse
import scipy as sp
    

def make_matrix(n):
  one_d = sparse.diags([[-1.]*(n-2), [2.]*n, [-1.]*(n-2)], [-2,0,2])
  A = (sparse.kronsum(one_d, one_d) + sparse.eye(n*n))
  A_csc = A.tocsc()
  A_csc.eliminate_zeros()
  A_lower = sparse.tril(A_csc, format = "csc")
  A_index = A_lower.indices
  A_indptr = A_lower.indptr
  A_x = A_lower.data
  return (A_index, A_indptr, A_x, A_csc)

def _symbolic_factor(A_indices, A_indptr):
  # Assumes A_indices and A_indptr index the lower triangle of $A$ ONLY.
  n = len(A_indptr) - 1
  L_sym = [np.array([], dtype=int) for j in range(n)]
  children = [np.array([], dtype=int) for j in range(n)]
  
  for j in range(n):
    L_sym[j] = A_indices[A_indptr[j]:A_indptr[j + 1]]
    for child in children[j]:
      tmp = L_sym[child][L_sym[child] > j]
      L_sym[j] = np.unique(np.append(L_sym[j], tmp))
    if len(L_sym[j]) > 1:
      p = L_sym[j][1]
      children[p] = np.append(children[p], j)
        
  L_indptr = np.zeros(n+1, dtype=int)
  L_indptr[1:] = np.cumsum([len(x) for x in L_sym])
  L_indices = np.concatenate(L_sym)
  
  return L_indices, L_indptr

# A_indices/A_indptr are the lower triangle, A is the entire matrix
A_indices, A_indptr, A_x, A = make_matrix(37)
parent, col_count = etree_base(A.indices, A.indptr)
L_indices, L_indptr = _symbolic_factor(A_indices, A_indptr)

true_parent = L_indices[L_indptr[:-2] + 1]
true_parent[np.where(np.diff(L_indptr[:-1]) == 1)] = -1
print(all(x == y for (x,y) in zip(parent[:-1], true_parent)))

true_col_count  = np.diff(L_indptr)
print(all(true_col_count == col_count))

True
True

Excellent. Now we just need to convert it to JAX.

Or do we?

To be honest, this is a little pointless. This function is only run once per matrix so we won’t really get much speedup¹⁴ from compilation.

Nevertheless, we might try.

@jit
def etree(A_indices, A_indptr):
 # print("(Re-)compiling etree(A_indices, A_indptr)")
  ## innermost while loop
  def body_while(val):
  #  print(val)
    j, node, parent, col_count, mark = val
    update_parent = lambda x: x[0].at[x[1]].set(x[2])
    parent = lax.cond(lax.eq(parent[node], -1), update_parent, lambda x: x[0], (parent, node, j))
    mark = mark.at[node].set(j)
    col_count = col_count.at[node].add(1)
    return (j, parent[node], parent, col_count, mark)

  def cond_while(val):
    j, node, parent, col_count, mark = val
    return lax.bitwise_and(lax.lt(node, j), lax.ne(mark[node], j))

  ## Inner for loop
  def body_inner_for(indptr, val):
    j, A_indices, A_indptr, parent, col_count, mark = val
    node = A_indices[indptr]
    j, node, parent, col_count, mark = lax.while_loop(cond_while, body_while, (j, node, parent, col_count, mark))
    return (j, A_indices, A_indptr, parent, col_count, mark)
  
  ## Outer for loop
  def body_out_for(j, val):
     A_indices, A_indptr, parent, col_count, mark = val
     mark = mark.at[j].set(j)
     j, A_indices, A_indptr, parent, col_count, mark = lax.fori_loop(A_indptr[j], A_indptr[j+1], body_inner_for, (j, A_indices, A_indptr, parent, col_count, mark))
     return (A_indices, A_indptr, parent, col_count, mark)

  ## Body of code
  n = len(A_indptr) - 1
  parent = jnp.repeat(-1, n)
  mark = jnp.repeat(-1, n)
  col_count = jnp.repeat(1,  n)
  init = (A_indices, A_indptr, parent, col_count, mark)
  A_indices, A_indptr, parent, col_count, mark = lax.fori_loop(0, n, body_out_for, init)
  return parent, col_count

Wow. That is ugly. But let’s see¹⁵ if it works!

parent_jax, col_count_jax = etree(A.indices, A.indptr)

print(all(x == y for (x,y) in zip(parent_jax[:-1], true_parent)))
print(all(true_col_count == col_count_jax))

True

True

Success!

I guess we could ask ourselves if we gained any speed.

Here is the pure python code.

import timeit
A_indices, A_indptr, A_x, A = make_matrix(20)

times = timeit.repeat(lambda: etree_base(A.indices, A.indptr),number = 1)
print(f"n = {A.shape[0]}: {[round(t,2) for t in times]}")

A_indices, A_indptr, A_x, A = make_matrix(50)
times = timeit.repeat(lambda: etree_base(A.indices, A.indptr),number = 1)
print(f"n = {A.shape[0]}: {[round(t,2) for t in times]}")


A_indices, A_indptr, A_x, A = make_matrix(200)
times = timeit.repeat(lambda: etree_base(A.indices, A.indptr),number = 1)
print(f"n = {A.shape[0]}: {[round(t,2) for t in times]}")

n = 400: [0.0, 0.0, 0.0, 0.0, 0.0]
n = 2500: [0.03, 0.03, 0.03, 0.03, 0.03]

n = 40000: [0.83, 0.82, 0.82, 0.82, 0.82]

And here is our JAX’d and JIT’d code.

A_indices, A_indptr, A_x, A = make_matrix(20)
times = timeit.repeat(lambda: etree(A.indices, A.indptr),number = 1)
print(f"n = {A.shape[0]}: {[round(t,2) for t in times]}")

A_indices, A_indptr, A_x, A = make_matrix(50)
times = timeit.repeat(lambda: etree(A.indices, A.indptr),number = 1)
print(f"n = {A.shape[0]}: {[round(t,2) for t in times]}")


A_indices, A_indptr, A_x, A = make_matrix(200)
times = timeit.repeat(lambda: etree(A.indices, A.indptr),number = 1)
print(f"n = {A.shape[0]}: {[round(t,2) for t in times]}")

parent, col_count= etree(A.indices, A.indptr)
L_indices, L_indptr = _symbolic_factor(A_indices, A_indptr)

A_indices, A_indptr, A_x, A = make_matrix(1000)
times = timeit.repeat(lambda: etree(A.indices, A.indptr),number = 1)
print(f"n = {A.shape[0]}: {[round(t,2) for t in times]}")

n = 400: [0.13, 0.0, 0.0, 0.0, 0.0]

n = 2500: [0.12, 0.0, 0.0, 0.0, 0.0]

n = 40000: [0.14, 0.02, 0.02, 0.02, 0.02]

n = 1000000: [2.24, 2.11, 2.12, 2.12, 2.12]

You can see that there is some decent speedup. For the first three examples, the computation time is dominated by the compilation time, but we see when the matrix has a million unknowns the compilation time is negligible. At this scale it would probably be worth using the fancy algorithm. That said, it is probably not worth sweating a three second that is only done once when your problem is that big!

The non-zero pattern of

Now that we know how many non-zeros there are, it’s time to populate them. Last time, I used some dynamic memory allocation to make this work, but JAX is certainly not going to allow me to do that. So instead I’m going to have to do the worst thing possible: think.

The way that we went about it last time was, to be honest, a bit arse-backwards. The main reason for this is that I did not have access to the elimination tree. But now we do, we can actually use it.

The trick is to slightly rearrange¹⁶ the order of operations to get something that is more convenient for working out the structure.

Recall from last time that we used the left-looking Cholesky factorisation, which can be written in the dense case as

def dense_left_cholesky(A):
  n = A.shape[0]
  L = np.zeros_like(A)
  for j in range(n):
    L[j,j] = np.sqrt(A[j,j] - np.inner(L[j, :j], L[j, :j]))
    L[(j+1):, j] = (A[(j+1):, j] - L[(j+1):, :j] @ L[j, :j].transpose()) / L[j,j]
  return L

This is not the only way to organise those operations. An alternative is the up-looking Cholesky factorisation, which can be implemented in the dense case as

def dense_up_cholesky(A):
  n = A.shape[0]
  L = np.zeros_like(A)
  L[0,0] = np.sqrt(A[0,0])
  for i in range(1,n):
    #if i > 0:
    L[i, :i] = (np.linalg.solve(L[:i, :i], A[:i,i])).transpose()
    L[i, i] = np.sqrt(A[i,i] - np.inner(L[i, :i], L[i, :i]))
  return L

This is quite a different looking beast! It scans row by row rather than column by column. And while the left-looking algorithm is based on matrix-vector multiplies, the up-looking algorithm is based on triangular solves. So maybe we should pause for a moment to check that these are the same algorithm!

A = np.random.rand(15, 15)
A = A + A.transpose()
A = A.transpose() @ A + 1*np.eye(15)

L_left = dense_left_cholesky(A)
L_up = dense_up_cholesky(A)

print(round(sum(sum(abs((L_left - L_up)[:])))),2)

0 2

They are the same!!

The reason for considering the up-looking algorithm is that it gives a slightly nicer description of the non-zeros of row i, which will let us find the location of the non-zeros in the whole matrix. In particular, the non-zeros to the left of the diagonal on row i correspond to the non-zero indices of the solution to the lower triangular linear system¹⁷ Because is sparse, this is a system of linear equations, rather than equations that we would have in the dense case. That means that the sparsity pattern of will be the union of the sparsity patterns of the columns of that correspond to the non-zero entries of .

This means two things. Firstly, if , then . Secondly, if $x^{(i)}_j $ and , then . These two facts give us a way of finding the non-zero set of if we remember just one more fact: a definition of the elimination tree is that if is a descendant of in the elimination tree.

This reduces the problem of finding the non-zero elements of to the problem of finding all of the descendants of in the subtree . And if there is one thing that people who are ok at programming are excellent at it is walking down a damn tree.