Best way to handle variable-length data?

[jzluo]

Hi, very new to jax but so far I've managed to get a ton of mileage out of just using it as a supercharged numpy. I'm trying to run logistic regressions on a series of traits for a group of subjects. Simple example:

import jax.numpy as jnp
import numpy as np
from jax.scipy.optimize import minimize
from jax.scipy.special import expit


def predict(X, b):
    return expit(jnp.matmul(X, b))

def loglikelihood(b, X, y):
    return -1 * (jnp.sum(y * jnp.log(predict(X, b)) + (1 - y) * jnp.log(1 - predict(X, b))))

@jax.jit
def fit(X, y):
    return minimize(loglikelihood,
               x0=jnp.zeros(X.shape[1]),
               args=(X, y),
               method="BFGS",
               )

Say y is an N x Traits array where -9 indicates a sample being excluded from the test for that trait:

array([[ 1,  1, -9],
       [-9,  1,  1],
       [ 0,  0,  0],
       [ 0, -9,  0]])

I would love to vmap(fit, in_axes = (None, 1)), but the shapes aren't static since different numbers of samples are excluded in each test. Instead, I've done an ugly:

def multi_fit(X, y):
    res = []
    for i in range(y.shape[1]):
        idx = np.flatnonzero(y[:, i] > -9)
        res.append(fit(X[idx], y[idx, i]))
    return res

Is there a better way to handle this situation?
Thanks!

[jakevdp]

Hi, thanks for the question! In this case the best way to proceed is probably to continue operating on fixed-size vectors while accounting for the missing data entries, without constructing intermediate values that have data-dependent shapes. So, for example, you could construct a log-likelihood that skips the missing values, something like this:

def loglikelihood(b, X, y):
  valid_mask = (y != -9)
  return -1 * (jnp.sum(valid_mask * (y * jnp.log(predict(X, b)) + (1 - y) * jnp.log(1 - predict(X, b)))))

Then you can vmap over your minimization without having to construct variable shaped arrays.

[jzluo]

Thank you, this is great! However, I'm noticing a large decrease in performance with vmap vs a manual for loop:

from sklearn.datasets import make_classification
X, y = make_classification(n_samples=100000, n_informative=4)
y = np.column_stack([y]*20)

for i in range(j):
    index = np.random.choice(y.shape[0], size=y.shape[0]//3)
    y[index, i] = -9

vmapped = jax.vmap(fit, in_axes=(None, 1))

%%timeit
res = vmapped(X, y)

33.1 s ± 351 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
res_manual = []
for i in range(y.shape[1]):
    res_manual.append(fit(X, y[:, i]))

2.5 s ± 35.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Am I doing something wrong?

[jakevdp]

If you're interested in accurate microbenchmarks of JAX code, you should start with FAQ: Benchmarking JAX Code to avoid inadvertently measuring the wrong thing (e.g., I suspect asynchronous dispatch may be confounding your results here). I'd give it a try myself but it's difficult without a complete reproduction that I can run (for example, what is j? What is fit in your benchmark?)

[jzluo]

Oops, sorry... Here's a complete example:

import jax
import jax.numpy as jnp
import numpy as np
from jax.scipy.special import expit
from jax.scipy.optimize import minimize
from sklearn.datasets import make_classification
jax.config.update("jax_enable_x64", True)

def predict(X, b):
    return expit(jnp.matmul(X, b))

def loglikelihood(b, X, y):
    valid_mask = (y != -9)
    return -1 * (jnp.sum(valid_mask * (y * jnp.log(predict(X, b)) + (1 - y) * jnp.log(1 - predict(X, b)))))

@jax.jit
def fit(X, y):
    res = minimize(loglikelihood,
               x0=jnp.zeros(X.shape[1]),
               args=(X, y),
               method="BFGS",
               )
    return res

def manual_fit(X, y):
    res = []
    for j in range(y.shape[1]):
        res.append(fit(X, y[:, j]))
    return res

vmapped = jax.vmap(fit, in_axes=(None, 1))

n_tests = 20
X, y = make_classification(n_samples=10000, n_informative=4)
y = np.column_stack([y]*n_tests)

for i in range(n_tests):
    index = np.random.choice(y.shape[0], size=y.shape[0]//3)
    y[index, i] = -9

%timeit jax.block_until_ready(manual_fit(X, y))
293 ms ± 6.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit jax.block_until_ready(vmapped(X, y))
1.9 s ± 3.24 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

I wasn't sure how jax.block_until_ready works with manual_fit(), so I also tried printing the result instead just to make sure the computation is done and got similar timings:

%%timeit
res = manual_fit(X, y)
print(res[0][0])
...
289 ms ± 862 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

[jakevdp]

I see, thanks. Indeed the vmapped version is slower. I think the reason for this is that the way optimize is coded, when minimization is done over a batch, it must perform the same number of iterations for every sample. So I suspect that there are some samples that due to the initialization, require many more iterations than others, so the vmapped case computes many more iterations for samples which converge quickly.

[jzluo]

Ah I see. Thank you, I'm learning a lot!