Implementing the Pool-adjacent-violators algorithm (PAVA) in jax

[bantin]

Hi jax folks. I'd like to implement a variant of the pool-adjacent-violators algorithm (PAVA) in Jax. For context, PAVA solves the following simple optimization problem, which is typically called isotonic regression:

min_x |x - y|_2^2 s.t x_1 >= x_2 >= ... >= x_N

Here x and y are both vectors. In my application, there is an additional parameter \gamma, since we'd like to specify that x is not just decaying, but exponentially decaying

min_x |x - y|_2^2 s.t x_1 * gamma >= x_2, x_2 * gamma >= x_3, etc

There are efficient algorithms for solving these in O(N) time where N is the length of x and y. In my case, I'll be happy with a naive O(T^2) implementation. It's easy to code up this naive implementation in a few lines of python:

def _pava_naive(y, gamma=1.0):
  # initialize x to y
  x = y.copy()

  for i in range(len(y)-1): 
    j = i
    # backtrack until constraint is satisfied
    while j >= 0 and x[j+1] > gamma * x[j]: # current block is from j to i+1
      data_block = x[j:i+2]
      numerator = np.sum(data_block * gamma**(np.arange(len(data_block))))
      denominator = np.sum(gamma**(np.arange(len(data_block)) + 1))
      block_start_val = numerator / denominator
      val = block_start_val * gamma**(np.arange(len(data_block)))
      x[j:i+2] = val
      j = j - 1
  return x

This works by sliding an index, i over the array, then sliding another index j backwards over the array starting from i. We then backtrack along the array (by decrementing j) until there are no constraint violations. It's easy to prove that this solves the problem above.

However, it's not clear to me how to implement this backtracking approach in jax. For instance, the following code won't work because we can't slice dynamically inside the body_fun of jax.lax.while_loop:

def _pava_naive_jax(y, gamma=1.0):
 x = y.copy()
 x = jnp.array(x)

 def _backtrack_until_satisfied(state):
   i, j, x = state
   data_block = x[j:i+2]
   numerator = np.sum(data_block * gamma**(jnp.arange(len(data_block))))
   denominator = np.sum(gamma**(jnp.arange(len(data_block)) + 1))
   block_start_val = numerator / denominator
   val = block_start_val * gamma**(jnp.arange(len(data_block)))
   x = x.at[j:i+2].set(val)
   j = j - 1
   return (i, j, x)

 def _outer_loop_fun(i, x):
   j = i
   state = (i, j, x)
   (i, j, x) = jax.lax.while_loop(
     jnp.all(jnp.array([
       j >= 0,
       x[j+1] > gamma * x[j]
     ])),
     _backtrack_until_satisfied(state),
     state)
   return x
 
 return jax.lax.fori_loop(0, len(y)-1, _outer_loop_fun, x)

It's not clear to me that there's a nice way of expressing this "backtrack until constraints are satisfied" idea in a JAX friendly way. I need to be able to JIT and VMAP this entire computation (otherwise I'd just use Julia!). Any thoughts here?

[jakevdp]

This does not look like an algorithm that is well-suited to JAX's computation model, both due to the dynamic array shapes, and due to the use of nested control flow, which will not perform well on accelerators.

That said, you can probably make it work by employing some tricks to express your dynamically-shaped computations in terms of statically-shaped arrays.

For example, you could replace this:

data_block = x[j:i+2]
numerator = np.sum(data_block * gamma**(np.arange(len(data_block))))

with something like this:

ind = jnp.arange(len(x))
data_block = jnp.where((ind >= j) && (ind < i + 2), x, 0
numerator = jnp.sum(data_block * gamma ** (ind - i))

That should compute the same value (unless you run into overflow/underflow with gamma) and do so in a way that is compatible with JAX transforms, because all array shapes are static. Similar tricks should work in the rest of the algorithm.

[bantin]

I see, thanks @jakevdp! I'm curious about nested control flow not working well on accelerators. This makes a bit of intuitive sense, but are there general principles about what to avoid?

In my case, solving this constrained problem happens inside of ADMM steps for another optimization problem. My original hope was to eventually write out the O(N) algorithm (which is similar to what I wrote above, but with some dynamic programming to avoid recomputation). I'm wondering if using a more general interior point method might be better here. The complexity guarantees are better for the dynamic programming solution, but maybe that won't be borne out in practice.

[jakevdp]

The general principle would be "avoid control flow on GPU" 😁. I'm not sure where it's documented, but essentially in the current GPU implementation every control flow decision requires a round-trip to the host. The control flow situation is better on TPU, but since TPU doesn't have a scalar core, scalar-heavy computations like this will not perform well.

[bantin]

every control flow decision requires a round-trip to the host

Yikes! Thanks for clarifying this. If there is a good place to read up on this, I'd be keen to know. Marking this as answered in the meantime.