Does it matter at which stage to average?

Mehran_Ziadloo: <p>This is a rather fundamental question and if I want to be honest, not really a PyTorch question.</p>
<p>Let’s say that I have two linear layers:</p>
<pre><code class="lang-python">layer1 = nn.Linear(input_size, hidden_size)
layer2 = nn.Linear(hidden_size, output_size)
</code></pre>
<p>I’m going to populate the input tensor like this:</p>
<pre><code class="lang-python">input = torch.randn(batch_size, arbitrary_size, input_size) 
</code></pre>
<p>Now, consider two different ways of applying the two layers to the input:</p>
<pre><code class="lang-python">hidden = layer1(input)
output = layer2(hidden).mean(dim=1)
</code></pre>
<p>and</p>
<pre><code class="lang-python">hidden = layer1(input).mean(dim=1)
output = layer2(hidden)
</code></pre>
<p>In both cases the output will be of shape <code>[batch_size, output_size]</code>. But in the first case, we’ll have:</p>
<pre><code class="lang-auto">[batch_size, arbitrary_size, input_size]
-&gt; [batch_size, arbitrary_size, output_size]
-&gt; [batch_size, output_size]
</code></pre>
<p>While in the second case, it’s going to be:</p>
<pre><code class="lang-auto">[batch_size, arbitrary_size, input_size]
-&gt; [batch_size, output_size]
-&gt; [batch_size, output_size]
</code></pre>
<p>My question is: What are the differences (if any) between the two cases? Are the learning capacity the same? How about the computational demand?</p>

KFrank: <p>Hi Mehran!</p>
<aside class="quote no-group" data-username="Mehran_Ziadloo" data-post="1" data-topic="188836" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mehran_ziadloo/48/17822_2.png" class="avatar"> Mehran_Ziadloo:</div>
<blockquote>
<pre data-code-wrap="python"><code class="lang-python">layer1 = nn.Linear(input_size, hidden_size)
layer2 = nn.Linear(hidden_size, output_size)
</code></pre>
</blockquote>
</aside>
<p>If you have two <code>Linear</code>s, one after the other <em>without</em> any intervening<br>
nonlinear activations (such as <code>relu()</code> or <code>sigmoid()</code>), then they collapse,<br>
in effect, into a single <code>Linear (input_size, output_size)</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>What are the differences (if any) between the two cases?</p>
</blockquote>
</aside>
<p>The <code>output</code> in the two cases will be numerically equal, up to some floating-point<br>
round-off error.</p>
<aside class="quote no-group">
<blockquote>
<p>Are the learning capacity the same?</p>
</blockquote>
</aside>
<p>The number of parameters in the two cases will be the same, which is an<br>
argument that they will have the same learning capacity.  The gradients<br>
produced by backpropagating will be different in the two cases (because<br>
you inject the <code>.mean()</code> in a different place), but my intuition suggests that<br>
training will progress in more or less the same way in the two cases.</p>
<aside class="quote no-group">
<blockquote>
<p>How about the computational demand?</p>
</blockquote>
</aside>
<p>The computational demand will be less in the second case (where <code>.mean()</code><br>
is injected after the first layer, rather than the second), because the size and<br>
dimensionality of the computation is reduced earlier in the forward pass.<br>
(However, based on how pytorch, especially with gpus, operates on entire<br>
tensors, the realized time savings is likely to be significantly less than the<br>
actual reduction in the number of individual floating-point operations.)</p>
<p>But again, there is no real point in what you are doing.  If this is what you<br>
want, just use a single layer of <code>Linear (input_size, output_size)</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

Mehran_Ziadloo: <p>Thanks, Frank.</p>
<p>This was my intuition too but I started doubting myself when I saw the <a href="https://github.com/huggingface/transformers/blob/5936c8c57ccb2bda3b3f28856a7ef992c5c9f451/src/transformers/models/whisper/modeling_whisper.py#L1962C54-L1962C54" rel="noopener nofollow ugc">OpenAI’s Whisper model (on HuggingFace)</a> using the first approach! I thought there might be something to it that I don’t understand. But now that it’s not just me, I feel more confident that my knowledge is not lacking.</p>