# Parallel or efficient computation of value with momentum

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.

• It is parallel-friendly, since you can do the computation for each $i$ in parallel. – D.W. Apr 10 '17 at 16:32
• @D.W. Well, in that way, I guess, and in fact my answer below is basically that (in matrix form), but I was wondering if there is some way in which I don't have to repeat all the computations for every element. – jdehesa Apr 10 '17 at 16:43

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).