xLSTM: Extended Long Short-Term Memory — part 28/30

Cleaned-text code points [407558, 423526). Bibliography omitted here; source.md includes it.

δ 𝒉 ~ t \displaystyle\delta_{\tilde{\bm{h}}_{t}} = 𝐨 t ⊙ δ 𝒉 t ext \displaystyle=\mathbf{o}_{t}\odot\delta_{\bm{h}_{t}}^{\text{ext}} (59) δ 𝑪 t − 1 ⊤ \displaystyle\delta_{\bm{C}_{t-1}}^{\top} = f t ​ δ 𝑪 t ⊤ + 𝒒 t − 1 ​ δ 𝒉 ~ t − 1 ⊤ max ⁡ { | 𝒏 t − 1 ⊤ ​ 𝒒 t − 1 | , 1 } \displaystyle=\mathrm{f}_{t}\delta_{\bm{C}_{t}}^{\top}+\frac{\bm{q}_{t-1}\delta_{\tilde{\bm{h}}_{t-1}}^{\top}}{\max\left\{{{\left|\bm{n}_{t-1}^{\top}\bm{q}_{t-1}\right|}},1\right\}} (60) δ 𝒏 t − 1 \displaystyle\delta_{\bm{n}_{t-1}} = f t ​ δ 𝒏 t − 𝒒 t − 1 ⊤ ​ 𝑪 t − 1 ⊤ ​ δ 𝐡 ~ t − 1 max ⁡ { | 𝒏 t − 1 ⊤ ​ 𝒒 t − 1 | , 1 } 2 ​ Ω ​ ( 𝒏 t − 1 ⊤ ​ 𝒒 t − 1 ) ​ 𝒒 t − 1 \displaystyle=\mathrm{f}_{t}\delta_{\bm{n}_{t}}-\frac{\bm{q}_{t-1}^{\top}\bm{C}_{t-1}^{\top}\delta_{\tilde{\mathbf{h}}_{t-1}}}{\max\left\{{{\left|\bm{n}_{t-1}^{\top}\bm{q}_{t-1}\right|}},1\right\}^{2}}\Omega\left(\bm{n}_{t-1}^{\top}\bm{q}_{t-1}\right)\bm{q}_{t-1} (61) δ 𝒗 t ⊤ \displaystyle\delta_{\bm{v}_{t}}^{\top} = i t ​ 𝒌 t ⊤ ​ δ 𝑪 t ⊤ \displaystyle=\mathrm{i}_{t}\bm{k}_{t}^{\top}\delta_{\bm{C}_{t}}^{\top} (62) δ 𝒌 t ⊤ \displaystyle\delta_{\bm{k}_{t}}^{\top} = i t ​ ( 𝒗 t ⊤ ​ δ 𝑪 t + δ 𝒏 t ⊤ ) \displaystyle=\mathrm{i}_{t}\left(\bm{v}_{t}^{\top}\delta_{\bm{C}_{t}}+\delta_{\bm{n}_{t}}^{\top}\right) (63) δ 𝒒 t \displaystyle\delta_{\bm{q}_{t}} = 𝑪 t ⊤ ​ δ 𝒉 ~ t max ⁡ { | 𝒏 t ⊤ ​ 𝒒 t | , 1 } − 𝒒 t ⊤ ​ 𝑪 t ⊤ ​ δ 𝒉 ~ t max ⁡ { | 𝒏 t ⊤ ​ 𝒒 t | , 1 } 2 ​ Ω ​ ( 𝒏 t ⊤ ​ 𝒒 t ) ​ 𝒏 t \displaystyle=\frac{\bm{C}_{t}^{\top}\delta_{\tilde{\bm{h}}_{t}}}{\max\left\{{{\left|\bm{n}_{t}^{\top}\bm{q}_{t}\right|}},1\right\}}-\frac{\bm{q}_{t}^{\top}\bm{C}_{t}^{\top}\delta_{\tilde{\bm{h}}_{t}}}{\max\left\{{{\left|\bm{n}_{t}^{\top}\bm{q}_{t}\right|}},1\right\}^{2}}\Omega\left(\bm{n}_{t}^{\top}\bm{q}_{t}\right)\bm{n}_{t} (64) δ 𝒙 t \displaystyle\delta_{\bm{x}_{t}} = ∑ g ∈ { q , k , v } 𝑾 g ⊤ ​ δ 𝒈 t \displaystyle=\sum_{g\in\{q,k,v\}}\bm{W}^{\top}_{g}\delta_{\bm{g}_{t}} (65) δ 𝑾 g ⊤ \displaystyle\delta_{\bm{W}_{g}}^{\top} = ∑ t 𝒙 t ​ δ 𝒈 t ⊤ , \displaystyle=\sum_{t}\bm{x}_{t}\delta_{\bm{g}_{t}}^{\top}\ , g ∈ { q , k , v } \displaystyle g\in\{q,k,v\} (66) δ 𝒃 g \displaystyle\delta_{\bm{b}_{g}} = ∑ t δ 𝒈 t , \displaystyle=\sum_{t}\delta_{\bm{g}_{t}}\ , g ∈ { q , k , v } \displaystyle g\in\{q,k,v\} (67) δ f ~ t \displaystyle\delta_{\tilde{\mathrm{f}}_{t}} = ( 𝟏 ⊤ ​ ( 𝑪 t − 1 ⊙ δ 𝑪 t ) ​ 𝟏 + 𝟏 ⊤ ​ ( 𝒏 t − 1 ⊙ δ 𝒏 t ) ) ​ γ ​ ( f ~ t ) \displaystyle=\left(\bm{1}^{\top}\left(\bm{C}_{t-1}\odot\delta_{\bm{C}_{t}}\right)\bm{1}+\bm{1}^{\top}\left(\bm{n}_{t-1}\odot\delta_{\bm{n}_{t}}\right)\right)\gamma\left(\tilde{\mathrm{f}}_{t}\right) (68) δ 𝐢 ~ t \displaystyle\delta_{\tilde{\mathbf{i}}_{t}} = ( 𝟏 ⊤ ​ ( ( 𝒗 t ​ 𝒌 t ⊤ ) ⊙ δ 𝑪 t ) ​ 𝟏 + 𝟏 ⊤ ​ ( 𝒌 t ⊙ δ 𝒏 t ) ) ​ exp ⁡ ( i ~ t ) \displaystyle=\left(\bm{1}^{\top}\left(\left(\bm{v}_{t}\bm{k}^{\top}_{t}\right)\odot\delta_{\bm{C}_{t}}\right)\bm{1}+\bm{1}^{\top}\left(\bm{k}_{t}\odot\delta_{\bm{n}_{t}}\right)\right)\exp\left(\tilde{\mathrm{i}}_{t}\right) (69) δ 𝐨 ~ t \displaystyle\delta_{\tilde{\mathbf{o}}_{t}} = 𝒉 ~ t ⊙ σ ′ ​ ( o ~ t ) ⊙ δ 𝒉 t \displaystyle=\tilde{\bm{h}}_{t}\odot\sigma^{\prime}\left(\tilde{\mathrm{o}}_{t}\right)\odot\delta_{\bm{h}_{t}} (70)

and Ω ⁡ ( z ) = Θ ⁡ ( z − 1 ) − Θ ⁡ ( − z − 1 ) \Omega\left(z\right)=\Theta\left(z-1\right)-\Theta\left(-z-1\right) , Θ ⁡ ( z ) \Theta\left(z\right) being the Heaviside step function. γ ⁡ ( z ) \gamma\left(z\right) is either σ ′ ​ ( z ) \sigma^{\prime}\left(z\right) or exp ⁡ ( z ) \exp\left(z\right) , depending on the forget gate activation.

##### Parallel mLSTM Forward Pass.

The mLSTM recurrence in Equations ( 19 - 27 ) can be reformulated in a parallel form, which is used to speed up training. After training we can still use the recurrent formulation for fast text generation.

Instead of processing each input 𝒙 t ∈ ℝ d \bm{x}_{t}\in\mathbb{R}^{d} at time step t t sequentially, the parallel version processes all timesteps of a full sequence 𝑿 ∈ ℝ T × d \bm{X}\in\mathbb{R}^{T\times d} at once, where T T is the sequence length and d d is the head dimension. We present the forward pass of the mLSTM for a single head and drop the head dimension for simplicity.

Let 𝐟 ~ ∈ ℝ T \tilde{\bm{\mathrm{f}}}\in\mathbb{R}^{T} be the forget gate pre-activations and 𝐢 ~ ∈ ℝ T \tilde{\bm{\mathrm{i}}}\in\mathbb{R}^{T} be the input gate pre-activations for a full sequence. We construct the forget gate activation matrix 𝐅 ∈ ℝ T × T \bm{\mathrm{F}}\in\mathbb{R}^{T\times T} by 𝐅 i ​ j = { 0 for ​ i < j 1 for ​ i = j ∏ k = j + 1 i σ ⁡ ( f ~ k ) for ​ i > j , \bm{\mathrm{F}}_{ij}=\begin{cases}0&\text{for}\;i<j\\ 1&\text{for}\;i=j\\ \prod_{k=j+1}^{i}\sigma\left(\tilde{\mathrm{f}}_{k}\right)&\text{for}\;i>j\end{cases}\ , (71) and the input gate pre-activation matrix 𝐈 ~ ∈ ℝ T × T \tilde{\bm{\mathrm{I}}}\in\mathbb{R}^{T\times T} by 𝐈 ~ i ​ j = { 0 for ​ i < j i ~ j for ​ i ⩾ j . \tilde{\bm{\mathrm{I}}}_{ij}=\begin{cases}0&\text{for}\;i<j\\ \tilde{\mathrm{i}}_{j}&\text{for}\;i\geqslant j\end{cases}\ . (72)

By applying the elementwise exponential input gate activation function naively, we obtain the unstabilized gate activation matrix 𝐃 ∈ ℝ T × T \bm{\mathrm{D}}\in\mathbb{R}^{T\times T} as 𝐃 = 𝐅 ⊙ exp ⁡ ( 𝐈 ~ ) . \bm{\mathrm{D}}=\bm{\mathrm{F}}\odot\exp(\tilde{\bm{\mathrm{I}}})\ . (73)

In order to avoid overflow due to the exponential function we apply the same stabilization as in the recurrent sLSTM, see Equation 15 . In the parallel formulation of the mLSTM we get a numerically stable gate activation matrix 𝐃 ′ ∈ ℝ T × T \bm{\mathrm{D}}^{\prime}\in\mathbb{R}^{T\times T} by taking the logarithm of 𝐃 \bm{\mathrm{D}} element-wise and subtracting the row-wise maximum value of 𝐃 \bm{\mathrm{D}} from each element: 𝐃 ~ \displaystyle\widetilde{\bm{\mathrm{D}}} = log ⁡ 𝐃 = log ⁡ ( 𝐅 ⊙ exp ⁡ ( 𝐈 ~ ) ) = log ⁡ 𝐅 + 𝐈 ~ \displaystyle=\log\bm{\mathrm{D}}=\log\left(\bm{\mathrm{F}}\odot\exp(\tilde{\bm{\mathrm{I}}})\right)=\log\bm{\mathrm{F}}+\tilde{\bm{\mathrm{I}}} (74) 𝐃 ′ \displaystyle\bm{\mathrm{D}}^{\prime} = exp ⁡ ( 𝐃 ~ − max ⁡ 𝐃 ~ ) \displaystyle=\exp(\widetilde{\bm{\mathrm{D}}}-\max\widetilde{\bm{\mathrm{D}}}) (75)

Given the queries, keys and values 𝑸 , 𝑲 , 𝑽 ∈ ℝ T × d \bm{Q},\bm{K},\bm{V}\in\mathbb{R}^{T\times d} , for a full sequence we can compute all hidden pre-activation states 𝐇 ~ ∈ ℝ T × d \widetilde{\bm{\mathrm{H}}}\in\mathbb{R}^{T\times d} in parallel for the un-stabilized version by 𝐇 ~ = 𝑪 ​ 𝑽 , with ​ 𝑪 = 𝑪 ~ max ⁡ ( | ∑ j = 1 T 𝑪 ~ i ​ j | , 1 ) , and ​ 𝑪 ~ = 𝑸 ​ 𝑲 ⊤ d ⊙ 𝐃 . \widetilde{\bm{\mathrm{H}}}=\bm{C}\,\bm{V}\ ,\quad\text{with}\,\,\,\bm{C}=\frac{\widetilde{\bm{C}}}{\max\left(|\sum_{j=1}^{T}\widetilde{\bm{C}}_{ij}|,1\right)}\ ,\quad\text{and}\,\,\,\widetilde{\bm{C}}=\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{d}}\odot\bm{\mathrm{D}}\ . (76) Note that we extract the 1 d \frac{1}{\sqrt{d}} factor for 𝑲 \bm{K} explicitly here and further on. For the stabilized version this yields 𝐇 ~ = 𝑪 ​ 𝑽 , with ​ 𝑪 = 𝑪 ~ ′ max ⁡ ( | ∑ j = 1 T 𝑪 ~ i ​ j ′ | , exp ⁡ ( − max ⁡ 𝐃 ~ ) ) , and ​ 𝑪 ~ ′ = 𝑸 ​ 𝑲 ⊤ d ⊙ 𝐃 ′ , \widetilde{\bm{\mathrm{H}}}=\bm{C}\,\bm{V}\ ,\quad\text{with}\,\,\,\bm{C}=\frac{\widetilde{\bm{C}}^{\prime}}{\max\left(|\sum_{j=1}^{T}\widetilde{\bm{C}}^{\prime}_{ij}|,\exp(-\max\widetilde{\bm{\mathrm{D}}})\right)}\ ,\quad\text{and}\,\,\,\widetilde{\bm{C}}^{\prime}=\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{d}}\odot\bm{\mathrm{D}}^{\prime}\ , (77) where for both versions the hidden pre-activation states 𝐇 ~ \widetilde{\bm{\mathrm{H}}} are identical.

With the output gate pre-activations 𝐎 ~ ∈ ℝ T × d \widetilde{\bm{\mathrm{O}}}\in\mathbb{R}^{T\times d} we can compute the hidden states 𝑯 ∈ ℝ T × d \bm{H}\in\mathbb{R}^{T\times d} for all timesteps by applying the output gate in parallel for each timestep element-wise: 𝐇 = σ ⁡ ( 𝐎 ~ ) ⊙ 𝐇 ~ . \bm{\mathrm{H}}=\sigma(\widetilde{\bm{\mathrm{O}}})\odot\widetilde{\bm{\mathrm{H}}}\ . (78)

This gives the parallel forward pass of the mLSTM for a full input sequence 𝑿 ∈ ℝ T × d \bm{X}\in\mathbb{R}^{T\times d} .

##### Parallel mLSTM Backward Pass.

We present the backward pass of the mLSTM for the stabilized version only. For completeness we summarize the forward pass in the stabilized version before we present the backward pass.

Given the forget gate matrix 𝐅 ∈ ℝ T × T \bm{\mathrm{F}}\in\mathbb{R}^{T\times T} , the logarithm of the forget gate matrix 𝐅 ¯ = log ⁡ 𝐅 ∈ ℝ T × T \overline{\bm{\mathrm{F}}}=\log\bm{\mathrm{F}}\in\mathbb{R}^{T\times T} , and the input gate matrix 𝐈 ∈ ℝ T × T \bm{\mathrm{I}}\in\mathbb{R}^{T\times T} as introduced above, together with the queries, keys and values 𝑸 , 𝑲 , 𝑽 ∈ ℝ T × d \bm{Q},\bm{K},\bm{V}\in\mathbb{R}^{T\times d} , we can write the forward pass of the mLSTM in the stabilized version as: 𝐃 ~ \displaystyle\widetilde{\bm{\mathrm{D}}} = 𝐅 ¯ + 𝐈 ~ \displaystyle=\overline{\bm{\mathrm{F}}}+\tilde{\bm{\mathrm{I}}} (79) 𝒎 \displaystyle\bm{m} = max j ⁡ 𝐃 ~ i ​ j , \displaystyle=\max_{j}\widetilde{\bm{\mathrm{D}}}_{ij}\ ,\quad row-wise maximum (80) 𝐃 ′ \displaystyle\bm{\mathrm{D}}^{\prime} = exp ⁡ ( 𝐃 ~ − 𝒎 ​ 1 ⊤ ) \displaystyle=\exp(\widetilde{\bm{\mathrm{D}}}-\bm{m}\,\bm{1}^{\top}) (81) 𝑪 ~ ′ \displaystyle\widetilde{\bm{C}}^{\prime} = 𝑸 ​ 𝑲 ⊤ d ⊙ 𝐃 ′ \displaystyle=\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{d}}\odot\bm{\mathrm{D}}^{\prime} (82) 𝒃 \displaystyle\bm{b} = ∑ j = 1 T 𝑪 ~ i ​ j ′ = 𝑪 ~ ′ ​ 1 , \displaystyle=\sum_{j=1}^{T}\widetilde{\bm{C}}^{\prime}_{ij}=\widetilde{\bm{C}}^{\prime}\ \bm{1}\ ,\quad row-wise sum (83) 𝒏 \displaystyle\bm{n} = max ⁡ ( | 𝒃 | , exp ⁡ ( − 𝒎 ) ) \displaystyle=\max\left(|\bm{b}|,\exp(-\bm{m})\right) (84) 𝑪 \displaystyle\bm{C} = 𝑪 ~ ′ ⊙ ( 𝒏 − 1 ​ 1 ⊤ ) \displaystyle=\widetilde{\bm{C}}^{\prime}\odot\left(\bm{n}^{-1}\ \bm{1}^{\top}\right) (85) 𝐇 ~ \displaystyle\widetilde{\bm{\mathrm{H}}} = 𝑪 ​ 𝑽 \displaystyle=\bm{C}\ \bm{V} (86)

With this forward pass we can compute the gradients δ a \delta_{a} for all intermediate and input variables to the mLSTM forward pass in the backward pass. We denote the gradient with respect to variable a a as δ a \delta_{a} .

Given the output gradient δ 𝐇 ~ ∈ ℝ T × d \delta_{\widetilde{\bm{\mathrm{H}}}}\in\mathbb{R}^{T\times d} we can compute the backward pass for the intermediate gradients as: δ 𝑪 ⊤ \displaystyle\delta_{\bm{C}}^{\top} = 𝑽 ​ δ 𝐇 ~ ⊤ \displaystyle=\bm{V}\delta_{\widetilde{\bm{\mathrm{H}}}}^{\top} (87) δ 𝒏 \displaystyle\delta_{\bm{n}} = − ( 𝑪 ~ ′ ⊙ ( 𝒏 − 2 ​ 1 ⊤ ) ⊙ δ 𝑪 ) ​ 1 \displaystyle=-\left(\widetilde{\bm{C}}^{\prime}\odot\left(\bm{n}^{-2}\ \bm{1}^{\top}\right)\odot\delta_{\bm{C}}\right)\ \bm{1} (88) = − ( ( 𝑪 ~ ′ ⊙ δ 𝑪 ) 𝟏 ) ⊙ 𝒏 − 2 \displaystyle=-\left(\left(\widetilde{\bm{C}}^{\prime}\odot\delta_{\bm{C}}\right)\bm{1}\right)\odot\bm{n}^{-2} (89) δ 𝒃 \displaystyle\delta_{\bm{b}} = sign ​ ( 𝒏 ) ⊙ δ 𝒏 ⊙ { 1 if ​ | 𝒃 | > exp ⁡ ( − 𝐦 ) 0 otherwise \displaystyle=\mbox{sign}\left(\bm{n}\right)\odot\delta_{\bm{n}}\odot\begin{cases}1&\text{if}\;|\bm{b}|>\exp(-\bm{\mathrm{m}})\\ 0&\text{otherwise}\end{cases} (90) δ 𝑪 ~ ′ , 𝑪 \displaystyle\delta_{\widetilde{\bm{C}}^{\prime},\bm{C}} = ( 𝒏 − 1 ​ 1 ⊤ ) ⊙ δ 𝑪 , \displaystyle=\left(\bm{n}^{-1}\ \bm{1}^{\top}\right)\odot\delta_{\bm{C}}\ ,\quad column-wise broadcast (91) δ 𝑪 ~ ′ , 𝒃 ⊤ \displaystyle\delta_{\widetilde{\bm{C}}^{\prime},\bm{b}}^{\top} = 𝟏 ​ δ 𝒃 ⊤ , \displaystyle=\bm{1}\ \delta_{\bm{b}}^{\top}\ ,\quad column-wise broadcast (92) δ 𝑪 ~ ′ \displaystyle\delta_{\widetilde{\bm{C}}^{\prime}} = δ 𝑪 ~ ′ , 𝑪 + δ 𝑪 ~ ′ , 𝑩 \displaystyle=\delta_{\widetilde{\bm{C}}^{\prime},\bm{C}}+\delta_{\widetilde{\bm{C}}^{\prime},\bm{B}} (93) δ 𝐃 ′ \displaystyle\delta_{\bm{\mathrm{D}}^{\prime}} = 𝑸 ​ 𝑲 ⊤ d ⊙ δ 𝑪 ~ ′ \displaystyle=\frac{\bm{Q}\bm{K}^{\top}}{\sqrt{d}}\odot\delta_{\widetilde{\bm{C}}^{\prime}} (94) δ 𝐃 ~ \displaystyle\delta_{\widetilde{\bm{\mathrm{D}}}} = exp ⁡ ( 𝐃 ~ − 𝒎 ) ⊙ δ 𝐃 ′ = 𝐃 ′ ⊙ δ 𝐃 ′ \displaystyle=\exp(\widetilde{\bm{\mathrm{D}}}-\bm{m})\odot\delta_{\bm{\mathrm{D}}^{\prime}}=\bm{\mathrm{D}}^{\prime}\odot\delta_{\bm{\mathrm{D}}^{\prime}} (95) We do not compute the gradients for 𝒎 \bm{m} as they cancel out (see the proof in the recurrent sLSTM).

With these intermediate gradients the gradients for the logarithmic forget gate matrix δ 𝐅 ¯ ∈ ℝ T × T \delta_{\overline{\bm{\mathrm{F}}}}\in\mathbb{R}^{T\times T} , the input gate matrix δ 𝐈 ∈ ℝ T × T \delta_{\bm{\mathrm{I}}}\in\mathbb{R}^{T\times T} , and the queries, keys and values δ 𝑸 , δ 𝑲 , δ 𝑽 ∈ ℝ T × d \delta_{\bm{Q}},\delta_{\bm{K}},\delta_{\bm{V}}\in\mathbb{R}^{T\times d} are given by δ 𝐅 ¯ \displaystyle\delta_{\overline{\bm{\mathrm{F}}}} = δ 𝐃 ~ \displaystyle=\delta_{\widetilde{\bm{\mathrm{D}}}} (96) δ 𝐈 \displaystyle\delta_{\bm{\mathrm{I}}} = δ 𝐃 ~ \displaystyle=\delta_{\widetilde{\bm{\mathrm{D}}}} (97) δ 𝑸 \displaystyle\delta_{\bm{Q}} = ( 𝐃 ′ ⊙ δ 𝑪 ~ ′ ) ​ 𝑲 d \displaystyle=\left(\bm{\mathrm{D}}^{\prime}\odot\delta_{\widetilde{\bm{C}}^{\prime}}\right)\frac{\bm{K}}{\sqrt{d}} (98) δ 𝑲 \displaystyle\delta_{\bm{K}} = ( 𝐃 ′ ⊙ δ 𝑪 ~ ′ ) ⊤ ​ 𝑸 d \displaystyle=\left(\bm{\mathrm{D}}^{\prime}\odot\delta_{\widetilde{\bm{C}}^{\prime}}\right)^{\top}\frac{\bm{Q}}{\sqrt{d}} (99) δ 𝑽 \displaystyle\delta_{\bm{V}} = 𝑪 ⊤ ​ δ 𝐇 ~ \displaystyle=\bm{C}^{\top}\delta_{\widetilde{\bm{\mathrm{H}}}} (100)

Having computed the gradients for the logarithmic forget gate matrix δ 𝐅 ¯ \delta_{\overline{\bm{\mathrm{F}}}} , we can compute the gradients for the forget gate pre-activations δ 𝐟 ~ = [ δ f ~ 1 , δ f ~ 2 , … , δ f ~ T ] ⊤ ∈ ℝ T \delta_{\tilde{\bm{\mathrm{f}}}}=\left[\delta_{\tilde{\mathrm{f}}_{1}},\delta_{\tilde{\mathrm{f}}_{2}},...,\delta_{\tilde{\mathrm{f}}_{T}}\right]^{\top}\in\mathbb{R}^{T} .

Recall the logarithmic forget gate matrix 𝐅 ¯ = log ⁡ 𝐅 \overline{\bm{\mathrm{F}}}=\log\bm{\mathrm{F}} is computed by 𝐅 ¯ i ​ j = log 𝐅 i ​ j = { − ∞ for ​ i < j 0 for ​ i = j ∑ k = j + 1 i log ⁡ σ ⁡ ( f ~ k ) ⏟ = : f ¯ k = ∑ k = j + 1 i f ¯ k for ​ i > j . \overline{\bm{\mathrm{F}}}_{ij}=\log\bm{\mathrm{F}}_{ij}=\begin{cases}-\infty&\text{for}\;i<j\\ 0&\text{for}\;i=j\\ \sum_{k=j+1}^{i}\underbrace{\log\sigma\left(\tilde{\mathrm{f}}_{k}\right)}_{=:\overline{\mathrm{f}}_{k}}=\sum_{k=j+1}^{i}\overline{\mathrm{f}}_{k}&\text{for}\;i>j\end{cases}\ . (101)

With the substitution 𝐟 ¯ = log ⁡ σ ⁡ ( 𝐟 ~ ) \overline{\bm{\mathrm{f}}}=\log\sigma(\tilde{\bm{\mathrm{f}}}) we compute the gradients for the logarithmic forget gate activations δ 𝐟 ¯ = [ δ f ¯ 1 , δ f ¯ 2 , … , δ f ¯ T ] ⊤ ∈ ℝ T \delta_{\overline{\bm{\mathrm{f}}}}=\left[\delta_{\overline{\mathrm{f}}_{1}},\delta_{\overline{\mathrm{f}}_{2}},...,\delta_{\overline{\mathrm{f}}_{T}}\right]^{\top}\in\mathbb{R}^{T} as δ f ¯ k \displaystyle\delta_{\overline{\mathrm{f}}_{k}} = ∑ j = 1 k − 1 ∑ i = k T ( δ 𝐅 ¯ ) i ​ j , \displaystyle=\sum_{j=1}^{k-1}\sum_{i=k}^{T}\left(\delta_{\overline{\bm{\mathrm{F}}}}\right)_{ij}\ , (102) δ f ~ k \displaystyle\delta_{\tilde{\mathrm{f}}_{k}} = σ ⁡ ( − f ~ k ) ⋅ δ f ¯ k , \displaystyle=\sigma(-\ \tilde{\mathrm{f}}_{k})\cdot\delta_{\overline{\mathrm{f}}_{k}}\ , (103)

where the last equation makes use of the following: d d ​ x ​ ( log ⁡ σ ​ ( x ) ) = − ( 1 + exp ( − x ) ) − 1 ⋅ exp ( − x ) ⋅ ( − 1 ) = exp ⁡ ( − x ) 1 + exp ⁡ ( − x ) = 1 1 + exp ⁡ ( x ) = σ ⁡ ( − x ) \displaystyle\begin{split}\frac{\mathrm{d}}{\mathrm{d}x}\left(\log\sigma(x)\right)&=-\left(1+\exp(-x)\right)^{-1}\cdot\exp(-x)\cdot(-1)\\ &=\frac{\exp(-x)}{1+\exp(-x)}=\frac{1}{1+\exp(x)}\\ &=\sigma(-x)\end{split} (104)

Finally, we compute the input gate pre-activations’ gradients δ 𝐢 ~ = [ δ i ~ 1 , δ i ~ 2 , … , δ i ~ S ] ⊤ ∈ ℝ T \delta_{\tilde{\bm{\mathrm{i}}}}=\left[\delta_{\tilde{\mathrm{i}}_{1}},\delta_{\tilde{\mathrm{i}}_{2}},...,\delta_{\tilde{\mathrm{i}}_{S}}\right]^{\top}\in\mathbb{R}^{T} as the column-wise sum over the rows of the input gate matrix δ 𝐈 \delta_{\bm{\mathrm{I}}} : δ i ~ k \displaystyle\delta_{\tilde{\mathrm{i}}_{k}} = ∑ i = k T ( δ 𝐈 ) i ​ k \displaystyle=\sum_{i=k}^{T}\left(\delta_{\bm{\mathrm{I}}}\right)_{ik} (105)

This completes the backward pass of the parallel mLSTM for a full input sequence 𝑿 ∈ ℝ T × d \bm{X}\in\mathbb{R}^{T\times d} .