Parallel or efficient computation of value with momentum

Question

I have a matrix $A$ with dimensions $M \times N$ and I want to compute $A'$ such that:

$$ A'_{i,j} = \alpha A'_{i,j-1} + (1 - \alpha) A_{i,j} \\ 1 \leq i \leq M, 1 \leq j \leq N, \alpha \in [0, 1] $$

Where $A'_{i,0}$ is some given constant.

I want to perform this computation as part of a machine learning training process on the GPU (I am using TensorFlow), but the only way to do it I can think of is with a loop over $j$, which makes the training tremendously slower (even if it's a TensorFlow loop, not a regular Python loop).

I know that being each value dependant on the previous one this is not a parallel-friendly computation, but I was wondering if there is some trick or reformulation that I am missing to make this in a smarter way.

Answer 1

I have noted that the above expression can be rewritten as:

$$ A'_{i,j} = \alpha^j A'_{i,0} + \sum_{1 \leq k \leq N} (1 - \alpha) \alpha^{j-k}A_{i,k} $$

Which makes me think I can define a $M \times N$ matrix $P$ as:

$$ P_{i,j} = \alpha^j A'_{i,0} $$

And a $N \times N$ matrix $R$ as:

$$ R_{i,j} = \begin{cases} \alpha^{j-i} & \text{if}~i \leq j \\ 0 & \text{otherwise} \end{cases} $$

And so I think that I would have:

$$ A' = P + (1 - \alpha) A R $$

Which would be a more parallelizable computation. The only challenges would be to efficiently construct $P$ and $R$ (which is probably not that hard) and, more importantly, the memory usage if $N$ is big (which unfortunately in my case it kinda is).