LSTM - divide gradients by number of timesteps IMMEDIATELY or in the end?

Question

From this answer I know that

the gradient of an average of many functions, is equal to the average of the gradients of those functions taken separately.

The error gradient that you want to calculate for gradient descent is $\nabla_{\theta} C(X, \theta)$, which you can therefore write as:

$$\nabla_{\theta} C(X, \theta) = \nabla_{\theta}(\frac{1}{|X|}\sum_{x \in X} L(x, \theta))$$

The derivative of the sum of any two functions is the sum of the derivatives, i.e.

$$\frac{d}{dx}(y+z) = \frac{dy}{dx} + \frac{dz}{dx}$$

In addition, any fixed multiplier that doesn't depend on the parameters you are taking the gradient with (in this case, the size of the dataset) can just be treated as an external factor:

$$\nabla_{\theta} C(X, \theta) = \frac{1}{|X|}\sum_{x \in X} \nabla_{\theta} L(x, \theta)$$

Question:

When working with LSTM specifically, am I indeed allowed to apply this $\frac{1}{|X|}$ at the end of the backprop once I've summed-up gradients of weights across all timesteps?

Or should I always apply $\frac{1}{|X|}$ straight away, to any gradient flowing into the top layers of my network?

---

This bothers me because LSTM is not merely multiplying things like a basic RNN would. Instead, LSTM has summation, which then gets multiplied with other things as we descend to earlier timesteps.

For example, this happens when we have to add some gradient flowing from all the 4 gates the LSTM into grad_wrt_resultAtPreviosTimesep (It will be required to compute grad during previous timestep once we get there)

Using an overly-simplified algebra example:

I feel that multiplying immediately by $\frac{1}{X}$ would represent the left side of this expression: $(\frac{1}{X}a + 40)*100 \neq \frac{1}{X} (a+40)*100$. On the other hand, multiplying by $\frac{1}{X}$ in the end of backprop would represent the right side.

Edit:

Just as a reminder, here is what LSTM looks like at each timestep, taken from this blog:

Answer 1

Answer to your question is it doesn't matter. The gradient is just a product of Jacobians (because of the chain rule), so it doesn't matter if you multiply the result or the intermediate multiplicand, these give the same result. Or, better yet, is to multiply the function you're differentiating, since $\tfrac{d(\alpha f(x))}{dx} = \alpha f'(x)$. So instead of manipulating the gradients you should just average your loss over whatever you want it to be averaged over.

Note that I didn't assume anything about the function above. It's because these are fairly general basic rules, and they apply to all kinds of functions, RNNs and LSTMs included.

---

That said, you might not want to actually average your loss (remembers, it's the same as averaging the gradients) over the timesteps. Average over inputs x – yeah, sure. Over timesteps? Not necessarily a good idea.

If you do average over timesteps, then you essentially decrease importance of longer sequences, and thus your model will care less of making more mistakes in them, whereas if you do not average over timesteps (or if they are all of the same length), then every timestep-loss has the same weight, and the model cares about all of them equally.