Why would backprop be slower in these two scenarios?

username: David_Ireland
cooked: <p>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.</p>
<p>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.</p>
<p>Thanks!</p>

username: ptrblck
cooked: <p>Could you post the model implementations you are using (a proxy model which shows the general idea might be sufficient)?</p>

username: David_Ireland
cooked: <p>Sure!</p>
<pre><code class="lang-auto">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)
</code></pre>
<p><code>Ensemble1</code> is the method with the slower backprop, even though the model is more lightweight than <code>Ensemble2</code>. In Ensemble2, the <code>MLP</code> model is essentially the same architecture as <code>Ensemble1</code> but without the extra <code>ensemble_size - 1</code> extra heads. <code>batched_forward</code> is how I run a forward pass for the ensemble. My inputs are all vectors of size <code>state_dim</code>, so when I draw a batch the tensor shape is <code>(batch_size * ensemble_size, state_dim)</code>, which (after shuffling the rows) I reshape to be <code>(batch_size, ensemble_size, state_dim)</code>.</p>

username: ptrblck
cooked: <p>The code is not executable as <code>MLP</code> is undefined.</p>

username: David_Ireland
cooked: <pre><code class="lang-auto">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)
</code></pre>
<p>here you go</p>

username: David_Ireland
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> were you able to look into this?</p>