agt:

Can you point me to the pytorch docs that implies the gradients of the first layers (before aggregation) will be averaged?


Don’t know if such a thing exists. That’s just how gradients work with tensor broadcasting (if scalar from one of inputs is used multiple times, you sum relevent derivatives).



 agt:

Do you have any superior idea for doing this kind of aggregation?


That’s a bit too abstract. Without positional info, you’re basically limited to mean() and exponential smoothing. Self-attention maybe. If you allow input position to affect output, you’re in the realm of usual sequence handlers: rnn, transformer/attention, causal convolutions etc.