Linear combinations of parameters across models

[pharringtonp19]

Is there a simple way to take a linear combination of the parameters of a model?

Below is my function which initializes parameters for a set of models

def init_fn(seed, n_features, n_models): 
      rngs = random.split(random.PRNGKey(seed), n_models) 
      params = vmap(model.init, in_axes=(0, None, None))(rngs, jnp.ones((1, n_features))) 
      opt_states = vmap(tx.init)(params)    
      return params, opt_states

[avital]

Can you elaborate on the question? I don't quite see how the code relates to the question about a linear combination of the parameters of the model

[pharringtonp19]

@avital Apologies for the lack of clarity. I meant to say, "is there a simple way to take a linear combination of parameters across models?"

The above code allows me to initialize the parameters for a set of models. That is, the above function returns params which contains parameters for a set of models.

I think it's relatively straightforward to average the outputs across models -- I first just need to vmap over model.apply as in jax.vmap(model.apply, in_axes=(0, None))(params,x) to get the outputs from each model, and then I can average.

What I would like to do in my case, though, is to first average the weights, and then use the averaged weights to predict, as in
model.apply(avg_params,x). How would I go about doing this? -- Thanks

[jheek]

To do this you can make use the jax.tree_map util to apply a function (mean in this case) over each param as follows:

avg_params = jax.tree_map(lambda p: jnp.mean(p, axis=0), params)

[pharringtonp19]

@jheek That is very clean! Thanks