Why would backprop be slower in these two scenarios?

autograd
1

I am trying to implement an ensemble, and for my uses I only need the uncertainty measurements over the final layer’s outputs. To that end, I first implemented a network which was some Linear layers (with relu non-linearity), and then for the output I had N layers. For a single input the model would then output N predictions, as in a normal ensemble.

I then implemented an ensemble similar to this except where I essentially just had N copies of the network (but with a single output in each layer). When I trained this, on the exact same data, backprop takes almost twice as long in the first scenario, despite have an order of magnitude less parameters. Does anyone have an idea why this would be? it seems counterintuitive to me that, given less parameters and the same amount of data, that backprop would take longer.

Thanks!

2

Could you post the model implementations you are using (a proxy model which shows the general idea might be sufficient)?

3

Sure!

class Ensemble1(nn.Module):
    def __init__(self, state_dim, hidden_dim, num_actions, num_heads, ensemble_size):
        super(Ensemble1, self).__init__()
        self.input_layer = nn.Linear(state_dim, hidden_dim)
        self.h1 = nn.Linear(hidden_dim, hidden_dim)
        self.h2 = nn.Linear(hidden_dim, hidden_dim)
        self.ensemble = nn.ModuleList()
        for j in range(ensemble_size):
            self.ensemble.append(nn.ModuleList())
            for i in range(num_heads):
                self.ensemble[j].append(nn.Linear(hidden_dim, num_actions))
        self.num_heads = num_heads

    def batched_forward(self, x):
        x = torch.relu(self.input_layer.forward(x))
        x = torch.relu(self.h1.forward(x))
        x = torch.relu(self.h2.forward(x))
        ensembles = []
        for idx, heads in enumerate(self.ensemble):
            vals = []
            for head in heads:
                vals.append(head.forward(x[:, idx, :]))
            ensembles.append(torch.stack(vals, dim=1))
        return torch.stack(ensembles, dim=1)


class Ensemble2(nn.Module):
    def __init__(self, state_dim, hidden_dim, num_actions, num_heads, ensemble_size):
        super(Ensembl2, self).__init__()
        self.ensemble = nn.ModuleList()
        for e in range(ensemble_size):
            self.ensemble.append(MLP(state_dim, hidden_dim, num_actions, num_heads))

    def batched_forward(self, x):
        vals = []
        for idx, net in enumerate(self.ensemble):
            vals.append(net.forward(x[:, idx, :]).transpose(0, 1))
        return torch.stack(vals, dim=1)

Ensemble1 is the method with the slower backprop, even though the model is more lightweight than Ensemble2. In Ensemble2, the MLP model is essentially the same architecture as Ensemble1 but without the extra ensemble_size - 1 extra heads. batched_forward is how I run a forward pass for the ensemble. My inputs are all vectors of size state_dim, so when I draw a batch the tensor shape is (batch_size * ensemble_size, state_dim), which (after shuffling the rows) I reshape to be (batch_size, ensemble_size, state_dim).

4

The code is not executable as MLP is undefined.

5 
class MLP(nn.Module):
    def __init__(self, state_dim, hidden_dim, num_actions, num_heads):
        super(MLP, self).__init__()
        self.input_layer = nn.Linear(state_dim, hidden_dim)
        self.h1 = nn.Linear(hidden_dim, hidden_dim)
        self.h2 = nn.Linear(hidden_dim, hidden_dim)
        self.output_heads = nn.ModuleList()
        for i in range(num_heads):
            self.output_heads.append(nn.Linear(hidden_dim, num_actions))
        self.num_heads = num_heads

    def forward(self, x):
        x = torch.relu(self.input_layer.forward(x))
        x = torch.relu(self.h1.forward(x))
        x = torch.relu(self.h2.forward(x))
        vals = []
        for i in range(self.num_heads):
            vals.append(self.output_heads[i].forward(x))
        return torch.stack(vals)

here you go

6

@ptrblck were you able to look into this?