Reversible rnns (`RevGRUCell` and `RevLSTMCell`)

[MartinuzziFrancesco]

Paper

RevGRU

$$\begin{align} \begin{bmatrix} \mathbf{z}_1^{(t)} \\\ \mathbf{r}_1^{(t)} \end{bmatrix} &= \sigma\left( \mathbf{W}_1^z \mathbf{x}^{(t)} + \mathbf{U}_1^z \mathbf{h}_2^{(t-1)} \right), \\\ \mathbf{g}_1^{(t)} &= \tanh\left( \mathbf{W}_1^g \mathbf{x}^{(t)} + \mathbf{U}_1^g \left( \mathbf{r}_1^{(t)} \circ \mathbf{h}_2^{(t-1)} \right) \right), \\\ \mathbf{h}_1^{(t)} &= \mathbf{z}_1^{(t)} \circ \mathbf{h}_1^{(t-1)} + \left( 1 - \mathbf{z}_1^{(t)} \right) \circ \mathbf{g}_1^{(t)}, \\\ \\\ \begin{bmatrix} \mathbf{z}_2^{(t)} \\\ \mathbf{r}_2^{(t)} \end{bmatrix} &= \sigma\left( \mathbf{W}_2^z \mathbf{x}^{(t)} + \mathbf{U}_2^z \mathbf{h}_1^{(t)} \right), \\\ \mathbf{g}_2^{(t)} &= \tanh\left( \mathbf{W}_2^g \mathbf{x}^{(t)} + \mathbf{U}_2^g \left( \mathbf{r}_2^{(t)} \circ \mathbf{h}_1^{(t)} \right) \right), \\\ \mathbf{h}_2^{(t)} &= \mathbf{z}_2^{(t)} \circ \mathbf{h}_2^{(t-1)} + \left( 1 - \mathbf{z}_2^{(t)} \right) \circ \mathbf{g}_2^{(t)}. \end{align}$$

RevLSTM

$$\begin{align} \begin{bmatrix} \mathbf{f}_1^{(t)} \\\ \mathbf{i}_1^{(t)} \\\ \mathbf{o}_1^{(t)} \\\ \mathbf{p}_1^{(t)} \end{bmatrix} &= \sigma\left( \mathbf{W}_1 \mathbf{x}^{(t)} + \mathbf{U}_1 \mathbf{h}_2^{(t-1)} \right), \\\ \mathbf{g}_1^{(t)} &= \tanh\left( \mathbf{W}_g \mathbf{x}^{(t)} + \mathbf{U}_g \mathbf{h}_2^{(t-1)} \right), \\\ \mathbf{c}_1^{(t)} &= \mathbf{f}_1^{(t)} \circ \mathbf{c}_1^{(t-1)} + \mathbf{i}_1^{(t)} \circ \mathbf{g}_1^{(t)}, \\\ \mathbf{h}_1^{(t)} &= \mathbf{p}_1^{(t)} \circ \mathbf{h}_1^{(t-1)} + \mathbf{o}_1^{(t)} \circ \tanh\left( \mathbf{c}_1^{(t)} \right). \end{align}$$

where h2 anc c2 are computed in the same way

Official implementation

https://github.com/matthewmackay/reversible-rnn