Dictionaries and tensor chain positions inputs

[AdrienCorenflos]

Hi,

I have been thinking about it and I don't believe allowing users to pass unordered data structures to the state is a good idea. The reason is fairly simple: how is the program supposed to know which coordinate of the inverse mass matrix corresponds to which element of the position (this even makes the linear regression sampling test non-deterministic as the hashing may differ from run to run)? Similarly, if the position is a non-flat tensor, the internals of the unravel logic are not clear (to the outer logic) so as to which order the individual items of the positions state will be in.

I am not certain of what a good solution to this would be yet, but I thought it was worth raising the concern. Am I just fantasising or does this seem like a legit problem to you too?

Adrien

[AdrienCorenflos]

@rlouf dunno if you monitor conversations ☝️

[rlouf]

Does seem like a legit concern now that I understand 😅

[mathDR]

I will post an example wherein I don't understand how to initialize the mass matrix for sampling. (Note, I got another way to work, which I will paste as well).

In my example, I am solving a simple Gaussian Process Regression (but making it a bit modular so you can swap our means and covariances).

The issue is how I generate the hyperparameters for the GP. There are 3 "categories" of hyperparameters: those for the mean, those for the kernel, and those for the likelihood noise.

I instantiate it with:

def init_GP_HyperParameters(
    rng_key=None,
    meanfunc=None,
    covfunc=None,
    X_Train=None,
    init_mean=None,
    init_cov=None,
    init_lik=None
    ):
    params={}
    if rng_key is None:
        params['mean'] = init_mean
        params['cov'] = init_cov
        params['lik'] = init_lik
    else:
        mean_key, cov_key, lik_key = split(rng_key, 3)
        params['mean'] = distrax.Normal(0, 1).sample(
            seed=mean_key, sample_shape=meanfunc()
        )
        params['cov'] = distrax.Normal(0, 1).sample(
            seed=mean_key, sample_shape=covfunc(None,X_Train)
        )
        params['lik'] = distrax.Normal(0, 1).sample(
            seed=lik_key, sample_shape=(1,)
        )
    return params

which allows for either passing in a meanfunc and covfunc and randomizing the initial hyperparameters, or passing in initial values.

An example meanfunc is:

def meanZero(meanhyper=None, x=None):
    ''' 
    Zero mean function. The mean function does not have any parameters
      m(x) = 0
    '''
    if meanhyper == None: # report number of parameters
        return 0
    n, D = x.shape
    return jnp.zeros((n,1)) 

and covfunc

def covSEard(hypers=None, x=None, z=None):

    '''
    Squared Exponential covariance function with Automatic Relevance Detemination
    (ARD) distance measure. The covariance function is parameterized as:
    k(x^p,x^q) = sf2 * exp(-(x^p - x^q)'*inv(P)*(x^p - x^q)/2)
    
    where the P matrix is diagonal with ARD parameters ell_1^2,...,ell_D^2, where
    D is the dimension of the input space and sf2 is the signal variance. The
    hyperparameters are:
    
    hyp = [ log(ell_1)
            log(ell_2)
               .
            log(ell_D)
            log(sqrt(sf2)) ]
    '''
    
    [n, D] = x.shape
    if hypers == None:                # report number of parameters
        return x.shape[1] + 1 
    
    ell = 1/jnp.exp(hypers[0:D])       # characteristic length scale

    sf2 = jnp.exp(2.*hypers[D])      # signal variance

    if z == 'diag':
        A = jnp.zeros((n,1))
    elif z == None:
        A = sq_dist(jnp.dot(jnp.diag(ell),x.T).T)
    else:                                           # compute covariance between data sets x and z
        A = sq_dist(jnp.dot(jnp.diag(ell),x.T).T,jnp.dot(jnp.diag(ell),z.T).T)       # cross covariances
 
    A = sf2*jnp.exp(-0.5*A)
                
    return A

with

def sq_dist(a, b=None):
    '''
      Compute a matrix of all pairwise squared distances
      between two sets of vectors, stored in the row of the two matrices:
      a (of size n by D) and b (of size m by D).
    '''
    tmp = a.shape
    n = tmp[0]
    if len(tmp) < 2:
        a = jnp.reshape(a,(n,1))
        D = 1
    else:
        D = a.shape[1]
    m = n    

    if b == None:
        b = a.T
    else:
        m = b.shape[0]
        b = b.T

    C = jnp.zeros((n,m))

    for d in range(D):
        tt = a[:,d]
        tt = tt.reshape(n,1)

        tem = jnp.kron(np.ones((1,m)), tt)
        tem = tem - jnp.kron(np.ones((n,1)), b[d,:])
        C = C + tem * tem  
    return C

(for completeness).

In either case, what is returned is a dict of (key, values), where the keys are strings and the values are jax DeviceArrays.

I was following along with the Introduction notebook

and wanted to replicate the following for my example:

inv_mass_matrix = np.array([0.5, 0.5])
num_integration_steps = 60
step_size = 1e-3

i.e. initialize an inv_mass_matrix for passing into hmc.kernel() but I found no way to do so. Would the "correct" way to be passing in 3 matrices: one of 1x1, one 2x2 and the last 1x1? (for my example)?

---------------------------------------------------------------------------
SOLUTION
What I got working was to extend the HierarchicalBNN Example. where my fit_and_eval function looks like:

def fit_and_eval(
    rng_key, potential_fn, meanfunc, covfunc, X_train, Y_train, X_test
):
    init_key, warmup_key, inference_key = split(
        rng_key, 3
    )

    # initialization
    potential = potential_fn(
        meanfunc, covfunc, X_train, Y_train
    )
    # Make this able to take init_key to randomize (pass in mean and cov funcs)
    initial_hyp = init_GP_HyperParameters(init_key,meanfunc,covfunc,X_train)

    initial_state = nuts.new_state(initial_hyp, potential)

    kernel_generator = lambda step_size, inverse_mass_matrix: nuts.kernel(
        potential, step_size, inverse_mass_matrix
    )

    # warm up
    final_state, (step_size, inverse_mass_matrix), info = stan_warmup.run(
        warmup_key, kernel_generator, initial_state, num_warmup
    )

    # inference
    nuts_kernel = jax.jit(nuts.kernel(potential, step_size, inverse_mass_matrix))
    states = inference_loop(inference_key, nuts_kernel, final_state, num_samples)
    samples = states.position
    
    return samples

So apparently here, stan_warmup "does the right thing" with the mass matrix?

[rlouf]

You might have a better understanding of what happens after having had a look at how the mass matrix is used:

https://github.com/blackjax-devs/blackjax/blob/main/blackjax/inference/hmc/metrics.py

You can see that we unravel the chain position and multiply the resulting vector by the matrix. The issue, if you want to manually specify and input it, is that you have to understand how JAX proceeds to unravel what it calls pytrees (list, tuples, dicts). However we shouldn't expect our users to know how JAX unravels pytrees. But as you now understand reading @AdrienCorenflos's messages in the main thread, this is not an easy design decision.

@mathDR Maybe you have an idea from a user perspective? What would be the most natural way to specify it for you?

[AdrienCorenflos]

For the sake of openness and also to keep records of the discussion, @rlouf and I had the following discussion offline

log-chat-tensor-input.log

The gist of the discussion is essentially contained in the following remarks:

1. Having a np.array as an input for the mass_matrix when the input for the position is a pytree is ambiguous:

If the position is a dictionary dict(x=1, y=0), and the mass_matrix is an array [[1, 2],[2,3]] (values are irrelevant) then which element of the matrix corresponds to x? Or y? Problem becomes even worse when pytrees become more complicated, and even when the order is well defined in the pytree, it can still be very confusing as an API for the user.

2. While this is a problem for manual input of mass_matrix, adaptation meta-algorithms are insensitive to ordering up to a certain extent (initialisation may still be a hurdle).

3. Most users will probably not set their matrices manually aside from illustrative/research purposes. Meta-algorithms however will possibly directly interact at the mass_matrix level.

4. A naive solution would be to enforce that the mass_matrix have the same pytree structure as the position, i.e. if position = (x, y), then mass_matrix = ([[1.]], [[2.]]) would be acceptable. However, this implicitely assumes independence across the pytree leaves which in most cases is a wrong assumption.

[AdrienCorenflos]

I tried to report the above with some objectivity. Now this is my personal opinion: in order to allow for different level of complexity in setting mass_matrices, several levels of inputs should be supported

1. mass_matrix is a float : a.k.a. the following returns True

def case_1(_position, mass_matrix):
    return np.issubdtype(np.dtype(mass_matrix), np.number)

then the ravelled mass matrix is the float times identity. Formally we would do

ravelled_position, _ = ravel_pytree(position)
ravelled_mass_matrix = mass_matrix * jnp.eye(ravelled_position.shape[0])

2. mass_matrix has the same pytree structure as position with leaves of the same shape: a.k.a. it verifies

def case_2(position, mass_matrix):
    type_ok = tree_structure(position) == tree_structure(mass_matrix)  # tree_util.tree_structure
    if not type_ok: 
        return False
    flattened_position, _ = tree_flatten(position)
    flattened_mass_matrix, _ = tree_flatten(mass_matrix)
    all_shapes_the_same = all(u.shape == v.shape for u, v in zip(flattened_position, flattened_mass_matrix)
    return all_shapes_the_same

then the ravelled mass matrix is a diagonal formed of all the arrays in mass_matrix leaves. Formally we would do

ravelled_position, _ = ravel_pytree(position)
ravelled_mass_matrix, _ = ravel_pytree(mass_matrix)

3. More complicated, I have not yet thought about the case of non-1D arrays. mass_matrix has the same pytree structure as position with leaves squared shape: a.k.a. it verifies

def case_3(position, mass_matrix):
    type_ok = tree_structure(position) == tree_structure(mass_matrix)
    if not type_ok: 
        return False
    flattened_position, _ = tree_flatten(position)
    flattened_mass_matrix, _ = tree_flatten(mass_matrix)
    shape_2d = all(len(v.shape) == 2 for v in flattened_mass_matrix)
    if not shape_2d: 
        return False
    return all(u.shape == v.shape[0] == v.shape[1] for u, v in zip(flattened_position, flattened_mass_matrix)

then the ravelled mass matrix is a diagonal formed of all the arrays in mass_matrix leaves. Formally we would do

flattened_position, _ = tree_flatten(position)
flattened_mass_matrix, _ = tree_flatten(mass_matrix)
ravelled_mass_matrix = scipy.linalg.block_diag(flattened_mass_matrix)  # hopefully ravel and flatten keep the order coherent...

4. The most general case. mass_matrix has a nested pytree structure corresponding to a position pytree where leaf would be replace by a position pytree. Moreover all the leaves of the final tree are compatible 2D arrays: a.k.a. it verifies

def case_4(position, mass_matrix):
    position_tree_structure = tree_structure(position)
    flattened_position, position_unflatten  = tree_flatten(position)
    #  I think this actually raises if no end case if found but need to check
    partially_flattened_mass_matrix, patrial_unflatten = tree_flatten(mass_matrix, is_leaf=lambda leaf: tree_structure(leaf) == position_tree_structure)
    if tree_structure(patrial_unflatten(flattened_position)) != position_tree_structure
        return False
    #  check shapes properly here
    return  ...

then the ravelled mass matrix is a diagonal formed of all the arrays in mass_matrix leaves. Formally we would do

flattened_position, _ = tree_flatten(position)
flattened_mass_matrix, _ = tree_flatten(mass_matrix)
ravelled_mass_matrix = np.block(...)  # organise the matrices coherently here.

The fourth item needs some ironing but hopefully you get the gist. Of course the goal is to make all this intuitive to the user. On the other hand having this kind of possible inputs has a clear advantage: it allows for low level optimisation down the line: e.g. matrix multiplications would become that more efficient is done block-wise, and even better if they reduce to component-wise multiplication.

[rlouf]

IN PROGRESS

Let's think in terms of user interface before implementation. The following cases are fairly straightforward:

state = 1.0
inv_mass_mat = 1.0

state = jnp.array([1., 2.])
inv_mass_mat = jnp.ones(2)  # inv_mass_mat = jnp.ones(2, 2)

Now let us us consider the case where the position of the chain is a dict:

state = {"x": 1.0, "y": 1.0}
inv_mass_mat_diag = jnp.array([1., 3.])

Ok so the main problem I see here, as a user, is that I am not sure what variable 3. and 1. correspond to. That's a problem. So maybe the next thing I want to try is:

state = {"x": 1.0, "y": 1.0}
inv_mass_mat = {"x": 1.0, "y": 3.0}

Yeah ok, I wrote something but I'm not sure what I'm really doing here. The issue is, this is a diagonal mass matrix. What if I want non-diagonal? Well:

state = {"x": 1.0, "y": 1.0}
inv_mass_mat = {"x": {"x": 1, "y": 4}, "y": {"x": 2, "y": 3}}

Which kind of works. I assume that when you're here you know what you're doing and the meaning of these individual coefficients?

What if I don't know what I'm doing?

Well that's what adaptation algorithms are for!

Should we provide a simple mass-matrix adaptation scheme that we can use if we know the step-size we want to use?

Conclusion

There are two separate concerns here:

• Users who know what they're doing, know what coefficients they want in their matrix. Whatever Adrien proposes is fine I believe;
• Users see the mass matrix as a necessary evil, they do not know or care what to use. But they know what step size to use. In this case we should provide a simple mass-matrix adaptation algorthm.

[albcab]

This seems like a good solution to me.

Might be overthinking this, but we could also provide support for a square matrix in the case of a pytree state/position (as long as dimensions match):

state = {"x": 1.0, "y": 1.0}
inv_mass_mat = jnp.arange(4, dtype=jnp.float32).reshape(2, 2)

and raise a warning telling the user how the matrix is interpreted:

flat_state, unravel_fn = ravel_pytree(state)
structured_inv_mass_matrix = {k: unravel_fn(cov) for k, cov in zip(state.keys(), inv_mass_matrix)}
warning(f"Covariances are set to {structured_inv_mass_matrix}\nif this is not intended...")

or in the case of a one dimensional vector:

state = {"x": 1.0, "y": 1.0}
inv_mass_mat_diag = jnp.arange(2, dtype=jnp.float32)
flat_state, unravel_fn = ravel_pytree(state)
structured_inv_mass_matrix_diag = unravel_fn(inv_mass_mat_diag)
warning(f"Variances are set to {structured_inv_mass_matrix} (covariances are all 0)\nif this is not intended...")