Manual Dense Ensemble

Average several independently initialized layers

Initialize K Dense layers with separate random keys, apply all of them to the same input vector and return their mean output. The layers have separate parameter values but the same architecture.

This exercise uses ordinary jax.vmap over initialization and an explicit loop over application. It illustrates how to handle a stacked variable tree; it does not implement a general Linen lifting API or train an ensemble.

Map initialization over keys

Construct one nn.Dense(features=F) module definition and split jax.random.PRNGKey(int(seed)) into K keys. Then call:

stacked = jax.vmap(model.init, in_axes=(0, None))(keys, x)

Axis 0 belongs to the keys. None keeps the whole input vector x available to each initialization call. The input is used for shape inference and the forward computation during initialization; it is not a sequence of per-member input slices.

The variable tree contains:

params
  kernel  (K, D, F)
  bias    (K, F)

There are K * (D * F + F) scalar parameters. Split keys let the random kernels differ, while Dense's zero-initialized biases remain equal. Reusing one key would produce equal initial values in separate parameter slices, not diverse initialized members.

Select one complete tree at a time

For member i, apply jax.tree_util.tree_map to select index i from every leaf. Keep the outer params collection in the resulting dictionary so it remains a valid argument to model.apply.

single = tree_map(lambda p, i=i: p[i], stacked)
output_i = model.apply(single, x)

Each selected kernel has shape (D, F) and each bias has shape (F,). Start an accumulator at jnp.zeros((F,)), add member outputs in index order, then divide by K.

i=i binds the current index as a default argument. It is useful when saving callbacks for later. In this loop, tree_map invokes its callback immediately, so lambda p: p[i] would also select the current member. JAX tracing does not make tree_map store that callback for a later index.

A two-member example

These known parameters show both the per-member outputs and their mean:

import jax
import jax.numpy as jnp
import flax.linen as nn

x = jnp.array([1., 2.])
stacked = {"params": {
    "kernel": jnp.array([[[1., 0.], [0., 1.]],
                         [[2., 0.], [0., -1.]]]),
    "bias": jnp.array([[0., 0.], [1., 1.]]),
}}
model = nn.Dense(features=2)
outputs = []
for i in range(2):
    single = jax.tree_util.tree_map(lambda p, i=i: p[i], stacked)
    outputs.append(model.apply(single, x))
members = jnp.stack(outputs)
mean = members.mean(axis=0)

assert jnp.array_equal(members, jnp.array([[1., 2.], [3., -1.]]))
assert jnp.array_equal(mean, jnp.array([2., .5]))
vectorized = jax.vmap(model.apply, in_axes=(0, None))(stacked, x)
assert jnp.allclose(vectorized, members)

The vectorized apply expresses the member axis directly. It is an alternative to the explicit loop, not a guarantee of a particular parallel schedule or speedup. A Python loop with static K is unrolled when the whole function is JIT-compiled; changing K can require another compilation.

What this mean represents

There is no activation or prediction head here. Because each Dense is affine, averaging its outputs is mathematically equivalent to using the average kernel and average bias:

mean_i(x @ W_i + b_i) = x @ mean_i(W_i) + mean_i(b_i)

Floating-point reduction order can introduce small differences. This equality does not generally hold for networks with nonlinear layers. This exercise is about parameter-tree handling, not evidence that averaging randomly initialized models improves prediction quality.