How to calculate multiple linear layer in one pass

username: ChrisLiu2
cooked: <p>I’ve been experimenting with a transformer-like architecture where one input vector x is passed through h different linear layers, one for each head, and let’s call them y_1, y_2, … y_h. I want to pass the results for each head (y_1,y_2,…y_h) into an additional linear layer independently of each other, something like y_1-&gt;lin_1-&gt;y’1, y_2-&gt;lin_2-&gt;y’_2,…</p>
<p>The first step, namely converting x to y_1, y_2, …y_h is simple. I can just use a linear layer of hidden_size as input size, and h*hidden_size as output size, and divide evenly the output into each head.</p>
<p>But the second step is hard to implement with GPU parallelization in mind. Currently, I’m iterating through the number of heads and doing each of the linear layers one by one. The pseudocode is:</p>
<pre><code class="lang-auto"># lins, y, y_prime are lists
for head in h:
    y_prime[head] = lins[head](y[head])
</code></pre>
<p>Is there a way supported by pytorch that can calculate the results for all linear layers simultaneously? Thanks.</p>

username: richard
cooked: <p>If I am understanding the problem correctly, this should be implementable using a single matrix multiply.</p>
<p>Assume that <code>y</code> is a tensor of shape <code>(n_head, hidden_size)</code>, and that each linear layer takes in <code>hidden_size</code> in-dims and returns <code>output_size</code> out-dims. This means that each <code>y'_i</code> has shape <code>(output_size,)</code>.</p>
<p>Instead of using <code>n_head</code> number of linear layers, we can store a single large weight of shape <code>(n_head, output_size, hidden_size)</code> and use do a batched matrix multiply between the weight and the <code>y</code>:</p>
<pre><code class="lang-auto">weight = torch.randn(n_head, output_size, hidden_size, requires_grad=True)
y = torch.randn(n_head, hidden_size)

y_prime = torch.matmul(y.unsqueeze(-1), weight).squeeze(-1)
</code></pre>
<p>This gets a little weirder but can still be extended to when you have a batch dimension in addition to your head dimension and if you care about the bias in a linear layer.</p>
<p>If you’re adventurous and on a nightly build of PyTorch, <a href="https://pytorch.org/docs/master/generated/torch.vmap.html">torch.vmap</a> can do this without the dimension squeezing business:</p>
<pre><code class="lang-auto">from torch._vmap_internals import vmap
weight = torch.randn(n_head, output_size, hidden_size, requires_grad=True)
y = torch.randn(n_head, hidden_size)
y_prime = vmap(torch.matmul)(weight, y)
</code></pre>

username: ChrisLiu2
cooked: <p>Thanks a lot for this elegant solution. I have one more question about batched matul: If instead of (n_head, hidden), the y_i has shape (batch, windows, n_head, hidden), do I need to manually unsqueeze the first two dimensions of the weight tensor and repeat them to match the size of y_i?</p>

username: richard
cooked: <p><code>torch.matmul</code> does broadcasting for batch dimensions (it does batched matrix-matrix multiply), so no, you don’t need to manually unsqueeze the first two dimensions of the weight tensor. However, we do have to be careful to make sure that the <code>n_head</code> dimension lines up correctly. So if <code>y</code> has shape <code>(batch, windows, n_head, hidden)</code> and <code>weight</code> has shape <code>(n_head, output, hidden)</code>, then we would want to line up the n_head dimensions by:</p>
<ul>
<li>unsqueezing <code>y</code> to have shape <code>(batch, windows, n_head, hidden, 1)</code>
</li>
<li>doing <code>matmul(weight, unsqueezed_y)</code> to get an output of shape <code>(batch, windows, n_head, output, 1)</code>
</li>
<li>squeezing away the last dimension of size 1.</li>
</ul>