Deep Ensemble with torch.func

Run every member of a deep ensemble in one vectorised call, then measure how much the members disagree.

A deep ensemble is a set of networks with the same architecture, trained from different random starts. Its prediction is the mean of the members' probabilities. A loop over members is slow on a GPU. torch.func lets you stack the members' weights and run them as one batched model:

params, buffers = torch.func.stack_module_state(models)
one = lambda p, b, x: torch.func.functional_call(base, (p, b), (x,))
logits = torch.func.vmap(one, in_dims=(0, 0, None))(params, buffers, X)

logits has shape (M, n, C): one row of outputs for each member. Turn them into probabilities and return two things:

mean_probs  = mean over members of softmax(logits)               (n, C)
mutual_info = H(mean_probs) - mean over members of H(member)    (n,)
H(p)        = -sum_c p_c * log(p_c)

mutual_info is zero when all members agree and grows when they disagree. It measures the uncertainty that comes from the model, not from noise in the data.

Function signature

def ensemble_predict(models, X):

Example

Input

one input, two members with probabilities
member 0: [0.9, 0.1]
member 1: [0.1, 0.9]

Output

mean_probs   ->  tensor([[0.5, 0.5]])
mutual_info  ->  tensor([0.3681])

Reasoning

The mean is [0.5, 0.5], with entropy log 2 = 0.6931. Each member has entropy -0.9 log 0.9 - 0.1 log 0.1 = 0.3251. The difference is 0.3681: each member is confident, but they do not agree.

Constraints and notes

• models is a list of M >= 1 modules with the same architecture, each mapping (n, d) to logits (n, C). X has shape (n, d).
• Use torch.func.vmap, not a Python loop over the members.
• The outputs must not track gradients, and the members must not change.
• Guard log(0): a probability can underflow to zero.