Mind-Blowing: Can Non-Linear RNNs Actually Be Parallelized Too?
In recent years, linear RNNs have attracted attention from researchers due to their properties such as parallelizable training and constant inference cost (see, for example, my earlier post 《Google's New Work Tries to "Revive" RNNs: Can RNNs Shine Again?》), which has let RNNs keep "a seat at the table" amid the Transformer's takeover. However, as things currently stand, that "seat" seems to belong exclusively to linear RNNs, since non-linear RNNs cannot be trained efficiently in parallel — leaving them "willing but unable" in the architecture wars.
However, a paper titled 《Parallelizing Non-Linear Sequential Models over the Sequence Length》 takes a different view. It proposes an iterative algorithm that, allegedly, enables parallel training of non-linear RNNs! Is this really possible? Let's dig in and find out.
Finding a Fixed Point
The original paper presents its method in a very general way, with an emphasis on PDEs and ODEs, but here we'll approach it directly through RNNs. Consider the common simple non-linear RNN:
\begin{equation}x_t = \tanh(Ax_{t-1} + u_t)\label{eq:rnn}\end{equation}more
Because of the presence of $\tanh$, this can only be computed serially. Now let's subtract $Ax_{t-1}$ from both sides:
\begin{equation}x_t - Ax_{t-1} = \tanh(Ax_{t-1} + u_t) - Ax_{t-1}\end{equation}
Of course, this doesn't change the fact that it's still fundamentally a non-linear RNN. However, we can notice that if the $x_{t-1}$ on the right-hand side were replaced by a given vector like $u_t$, this would become a linear RNN, and according to the results in 《Google's New Work Tries to "Revive" RNNs: Can RNNs Shine Again?》, it could be computed in parallel. At this point, sharp readers may have already guessed what comes next — iterative solving!
First, rewrite the RNN above as
\begin{equation}x_t^{(n)} - Ax_{t-1}^{(n)} = \tanh(Ax_{t-1}^{(n-1)} + u_t) - Ax_{t-1}^{(n-1)}\label{eq:rnn-iter}\end{equation}
Starting from a given $x_t^{(0)}$, we repeatedly iterate the above equation. Ideally, it converges to a fixed point $x_t^*$, which is exactly the result of the original non-linear RNN computation. Of course, in theory the total computational cost of iterating through equation $\eqref{eq:rnn-iter}$ is greater than directly computing recursively via equation $\eqref{eq:rnn}$. But since each iteration step is a parallelizable linear RNN, and if convergence is reasonably fast so that the number of iterations needed isn't too large, the total time cost will usually end up faster than direct non-linear RNN recursion (especially when the sequence length is large).
A Simplified Form
In fact, the reason non-linear RNNs are slow isn't primarily about the inability to parallelize — the real crux of the matter is that they involve a large number of non-element-wise operations, such as the matrix multiplication $Ax_{t-1}$ inside $\tanh$ in equation $\eqref{eq:rnn}$. And the reason linear RNNs are fast isn't just that they allow parallel training — more crucially, they can be diagonalized to turn matrix multiplication into element-wise multiplication, and element-wise multiplication isn't too slow even when computed serially.
Once we've turned the non-linear RNN into an iteration of linear RNNs via equation $\eqref{eq:rnn-iter}$, we can likewise enjoy the "privilege" of diagonalizability that linear RNNs have, thereby improving computational speed. Specifically, diagonalizing $A$ over the complex field as $P\Lambda P^{-1}$, equation $\eqref{eq:rnn-iter}$ becomes
\begin{equation}x_t^{(n)} - P\Lambda P^{-1} x_{t-1}^{(n)} = \tanh(P\Lambda P^{-1} x_{t-1}^{(n-1)} + u_t) - P\Lambda P^{-1} x_{t-1}^{(n-1)}\end{equation}
Left-multiplying both sides by $P^{-1}$:
\begin{equation}P^{-1} x_t^{(n)} - \Lambda P^{-1} x_{t-1}^{(n)} = P^{-1}\tanh(P\Lambda P^{-1} x_{t-1}^{(n-1)} + u_t) - \Lambda P^{-1} x_{t-1}^{(n-1)}\end{equation}
Letting $y_t = P^{-1} x_t$, the above equation simplifies to
\begin{equation}y_t^{(n)} - \Lambda y_{t-1}^{(n)} = P^{-1}\tanh(P\Lambda y_{t-1}^{(n-1)} + u_t) - \Lambda y_{t-1}^{(n-1)}\end{equation}
Since an RNN is typically followed by a projection layer anyway, the $P$ in $x_t = P y_t$ can in principle be absorbed into that external projection layer. In other words, the equation above theoretically has the same expressive power as the original $\eqref{eq:rnn}$, but since $\Lambda$ is diagonal, the computational cost of the recursion is significantly reduced. The equation above also involves the inverse matrix $P^{-1}$, which is not only expensive to compute but also unfavorable for optimization, so we might as well replace $P^{-1}$ and $P\Lambda$ with two unrelated parameter matrices:
\begin{equation}y_t^{(n)} - \Lambda y_{t-1}^{(n)} = P\tanh(Q y_{t-1}^{(n-1)} + u_t) - \Lambda y_{t-1}^{(n-1)}\end{equation}
as long as they're initialized as $PQ=\Lambda$.
The Perturbation Idea
Suppose $x_t^{(0)}=0$; then equation $\eqref{eq:rnn-iter}$ decomposes the original non-linear RNN into a sequence of linear RNNs:
\begin{equation}\begin{array}{c} x_t^{(1)} - Ax_{t-1}^{(1)} = \tanh(u_t)\\ x_t^{(2)} - Ax_{t-1}^{(2)} = \tanh(Ax_{t-1}^{(1)} + u_t) - Ax_{t-1}^{(1)} \\ \vdots \\ x_t^{(n)} - Ax_{t-1}^{(n)} = \tanh(Ax_{t-1}^{(n-1)} + u_t) - Ax_{t-1}^{(n-1)} \\ \vdots \\ \end{array}\label{eq:rnns}\end{equation}
And if we assume that $x_{t-1},u_t$ are all small quantities, then applying $\tanh x \approx x$ to the right-hand side of equation $\eqref{eq:rnn}$ gives:
\begin{equation}x_t = \tanh(Ax_{t-1} + u_t) \approx Ax_{t-1} + u_t \approx Ax_{t-1} + \tanh(u_t)\label{eq:rnn-approx}\end{equation}
This is exactly the first equation in $\eqref{eq:rnns}$, so if the assumption holds, then $x_t^{(1)}$ might already be close enough to the ideal $x_t^*$, with each subsequent iteration step rapidly converging toward it. From this we can see that "subtracting $Ax_{t-1}$ from both sides" is the key move: it makes the first iteration step of $\eqref{eq:rnn-iter}$ approximate the first-order linear approximation of the original non-linear RNN, which speeds up convergence. This is a classic technique in mathematical physics known as "perturbation."
Speeding Up Convergence
Following the logic of perturbation methods, the key to improving convergence speed is to improve the precision of the approximate expansion. A relatively simple improvement, for example, is to assume only that $x_{t-1}$ is small, in which case a first-order Taylor expansion gives (treating $u_t$ as a column vector, where $\circ$ denotes the Hadamard product)
\begin{equation}x_t = \tanh(Ax_{t-1} + u_t) \approx \tanh(u_t) + (\text{sech}^2 u_t\circ A)x_{t-1}\end{equation}
So the improved result turns equation $\eqref{eq:rnn-iter}$ into
\begin{equation}x_t^{(n)} - A_t x_{t-1}^{(n)} = \tanh(Ax_{t-1}^{(n-1)} + u_t) - A_t x_{t-1}^{(n-1)}\label{eq:iter-plus1}\end{equation}
where $A_t = \text{sech}^2 u_t\circ A$. A more refined improvement is to expand around the result of the previous iteration at each step:
\begin{equation}\begin{aligned} x_t =&\, \tanh(Ax_{t-1} + u_t) \\ \approx&\, \tanh(Ax_{t-1}^{(n-1)} + u_t) + (\text{sech}^2 (Ax_{t-1}^{(n-1)} + u_t)\circ A)(x_{t-1} - x_{t-1}^{(n-1)}) \end{aligned}\end{equation}
so that equation $\eqref{eq:rnn-iter}$ becomes
\begin{equation}x_t^{(n)} - A_t^{(n)} x_{t-1}^{(n)} = \tanh(Ax_{t-1}^{(n-1)} + u_t) - A_t^{(n)} x_{t-1}^{(n-1)}\label{eq:iter-plus2}\end{equation}
where $A_t^{(n)}=\text{sech}^2 (Ax_{t-1}^{(n-1)} + u_t)\circ A$. This final iterative scheme is, in fact, exactly "Newton's method" for finding numerical solutions to equations, which has quadratic convergence.
Why Bother With Convergence
In theory, both improvements $\eqref{eq:iter-plus1}$ and $\eqref{eq:iter-plus2}$ genuinely do improve convergence speed. However, they make the matrix $A$ in each linear recursion step depend on $t$ or even $n$, which greatly increases the complexity of parallelization and also prevents us from using the diagonalization trick from the "Simplified Form" section to speed things up. On the other hand, if we keep the iterative scheme $\eqref{eq:rnn-iter}$ as is, while it offers many efficiency benefits, convergence really can't be well guaranteed.
Is there really no way to reconcile this contradiction? Actually, in my view, the most straightforward approach is to simply "not worry about it" — once we've used the non-linear RNN to derive $\eqref{eq:rnn-iter}$, we can just forget about the original non-linear RNN and take equation $\eqref{eq:rnn-iter}$ as the base model. In other words, why worry about whether equation $\eqref{eq:rnn-iter}$ converges to the original non-linear RNN? Why not just take it as a new starting point directly? Whatever result gradient descent learns is the result — and if what gradient descent learns is a model that does not converge to the original non-linear RNN, that simply means not converging to the original RNN is the more suitable outcome.
Once we set aside this mental constraint, many things become much clearer. First, even though equation $\eqref{eq:iter-plus2}$ theoretically has excellent convergence speed, it's still conditional, and in the context of deep learning, guaranteeing those conditions would be a rather extravagant demand. In other words, even the convergence of equation $\eqref{eq:iter-plus2}$ is not absolutely guaranteed, so why should we be too harsh on equation $\eqref{eq:rnn-iter}$, as if the pot were calling the kettle black? Second, once we treat equation $\eqref{eq:rnn-iter}$ as a new starting point, we can simply understand it as a new way of using linear RNNs — or as an approach to addressing certain shortcomings of linear RNNs (such as the fact that linear RNNs are not Turing-complete). This gives us much more room to work with.
Overall, letting go of concerns about convergence seems to be a better way to break out of this mental deadlock and explore more general results.
The General Case
Everything discussed at length above has revolved around the simple non-linear RNN, i.e., equation $\eqref{eq:rnn}$. What about the more commonly used LSTM and GRU?
Take GRU as an example. Its original form is
\begin{equation}\begin{aligned} z_{t} & = \sigma \left( W_{z} x_{t} + U_{z} h_{t - 1} + b_{z} \right) \\ r_{t} & = \sigma \left( W_{r} x_{t} + U_{r} h_{t - 1} + b_{r} \right) \\ \hat{h}_t & = \tanh \left( W_{h} x_{t} + U_{h} (r_t \circ h_{t - 1}) + b_{c} \right)\\ h_{t} & = \left(1 - z_{t}\right) \circ h_{t - 1} + z_{t} \circ \hat{h}_t \end{aligned}\end{equation}
In the initial stage, all the gates can be approximately treated as $\frac{1}{2}$, so following the same approach as equation $\eqref{eq:rnn-approx}$ we get
\begin{equation}\begin{aligned} h_{t} &\, = \left(1 - z_{t}\right) \circ h_{t - 1} + z_{t} \circ \hat{h}_t \\ &\, \approx \frac{1}{2} h_{t - 1} + \frac{1}{2} \hat{h}_t \\ &\, \approx \frac{1}{2} h_{t - 1} + \frac{1}{2} \left(\tanh ( W_{h} x_{t} + b_{c} ) + \frac{1}{2}U_{h} h_{t - 1}\right) \\ &\, = \frac{1}{2} \left(I + \frac{1}{2}U_{h}\right)h_{t - 1} + \frac{1}{2} \tanh ( W_{h} x_{t} + b_{c} ) \\ \end{aligned}\end{equation}
So we can choose $A=\frac{1}{2} \left(I + \frac{1}{2}U_{h}\right)$ and rewrite the GRU as the iteration
\begin{equation}\begin{aligned} z_{t}^{(n)} & = \sigma \left( W_{z} x_{t} + U_{z} h_{t - 1}^{(n-1)} + b_{z} \right) \\ r_{t}^{(n)} & = \sigma \left( W_{r} x_{t} + U_{r} h_{t - 1}^{(n-1)} + b_{r} \right) \\ \hat{h}_t^{(n)} & = \tanh \left( W_{h} x_{t} + U_{h} (r_t^{(n)} \circ h_{t - 1}^{(n-1)}) + b_{c} \right)\\ h_{t}^{(n)} & = Ah_{t-1}^{(n)} - Ah_{t-1}^{(n - 1)} + \left(1 - z_{t}^{(n)}\right) \circ h_{t - 1}^{(n-1)} + z_{t}^{(n)} \circ \hat{h}_t^{(n)} \end{aligned}\end{equation}
Overall, this transformation of non-linear RNNs into iterations of linear RNNs, from a practical standpoint, uses the non-linear RNN merely as a guide to derive a method for parameter sharing and composition across multiple layers of linear RNNs — however many iteration steps there are, that's how many layers' worth of computation you get. This naturally raises a question: unless one can prove that non-linear RNNs like GRU and LSTM have an absolute advantage, why not just directly stack several layers of "linear RNN + MLP"?
Summary
This post has briefly explored the question of parallelizing non-linear RNNs. Through the "perturbation" idea from mathematical physics, we can transform a non-linear RNN into an iteration of linear RNNs, thereby exploiting the parallelizability of linear RNNs to achieve parallel computation for non-linear RNNs as well.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.