Difference in the order of applying linear decoder and average pooling for sequence models
I am working with sequence modelling in pytorch and trying to determine if the order of the pooling and linear decoding layer matters. Given that I have a sequence with the shape (Batch, Seqlen, dim_model) and I want to transform it into (Batch, dim_output) I will need a pooling layer for reducing the second dimension (SeqLen) and an affine transformation that maps dim_model to dim_output. Assume Batch = 16, SeqLen = 6000, dim_model = 32, dim_output = 5, we have the following input:
import torch

pooler = lambda x: x.mean(dim=1)
decoder = torch.nn.Linear(32, 5)

x = torch.randn(16, 6000, 32)

Would this:
y = decoder(pooler(x))

Be the same as:
y = pooler(decoder(x))

The normalized difference between both outputs suggest that they are close:
torch.norm(decoder(pooler(x)) - pooler(decoder(x)))

output:
tensor(6.5412e-08, grad_fn=&lt;CopyBackwards&gt;)

But can one say they are equivalent? Are the gradients computed in the same way?
I am interesting the case of using arbitrary pooling layer, this includes for instance the &quot;last&quot; pooler:
pooler = lambda x: x[:,-1]
torch.norm(decoder(pooler(x)) - pooler(decoder(x)))

output:
tensor(0., grad_fn=&lt;CopyBackwards&gt;)


A linear layer does x -&gt; Ax+b for some matrix A and vector b.
If you have a bunch of x (x1, x2, x3, ..., xn) then A[(x1+...+xn)/n] = (Ax1 +... +Axn)/n, so for mean pooling, applying pooling first and then doing the linear layer results (up to floating point errors) in the same value as applying the linear layer first and then doing the pooling.
For &quot;last pooling&quot;, the result is the same because it doesn't matter whether you apply A to every element and then afterwards only pick the final one, or if you pick the final one, and apply A to it.
However, for plenty of  other operations,  the result  would not be the same. E.g. for max pooling, the result would in general not be the same.
e.g. if x1 = (1, 0, 0), x2 = (0, 1, 0), x3 = (0, 0, 1), and A = ((1, 1, 1)) then Ax1 = Ax2 =Ax3 = (1), so applying max pooling after the linear layer just gives you (1),
but max pooling applied to x1, x2, x3 gives you (1, 1, 1) and A(1, 1, 1) = 3.
