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

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

VG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} {}{ {{}}\lx@inpgf@ignorespaces\hbox{\hbox{{\lxSVG@begingroup@{_scopebegin=1} {{}{}{{ {}{}}}{ {}{}} {{}{{\lx@inpgf@ignorespaces}}}{{}{\lx@inpgf@ignorespaces}}{}{{}{\lx@inpgf@ignorespaces}} {\lx@inpgf@ignorespaces }{{{{\lx@inpgf@ignorespaces}}\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{1.0}{0.0}{0.0}{1.0}{-12.2988pt}{0.0pt}\lxSVG@begingroup@{transform=matrix(1.0 0.0 0.0 1.0 -17.02 0)} \pgfsys@hbox{65}\lxSVG@closescope }}} \lxSVG@closescope }}} } \lxSVG@closescope \hbox to0.0pt{}{{ {}{}{}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}}}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}\left(\tilde{\mathbf{i}}_{t}\right)\ , 𝐢 ~ t \displaystyle\tilde{\mathbf{i}}_{t}\ = 𝑾 𝐢 ​ 𝒙 t + 𝑹 𝐢 ​ 𝒉 t − 1 + 𝒃 𝐢 \displaystyle=\ \bm{W}_{\mathbf{i}}\ \bm{x}_{t}\ +\ \bm{R}_{\mathbf{i}}\ \bm{h}_{t-1}\ +\ \bm{b}_{\mathbf{i}}\ \ input gate (38) = ​ ( 𝐟 ~ t ) ​ OR ​ σ ​ ( 𝐟 ~ t ) , \displaystyle=\ \mathchoice{\hbox to42.45pt{\vbox to6.25pt{\pgfpicture\makeatletter\hbox{\hskip 21.22247pt\lower-1.94443pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} {}{ {{}}\lx@inpgf@ignorespaces\hbox{\hbox{{\lxSVG@begingroup@{_scopebegin=1} {{}{}{{ {}{}}}{ {}{}} {{}{{\lx@inpgf@ignorespaces}}}{{}{\lx@inpgf@ignorespaces}}{}{{}{\lx@inpgf@ignorespaces}} {\lx@inpgf@ignorespaces }{{{{\lx@inpgf@ignorespaces}}\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{1.0}{0.0}{0.0}{1.0}{-21.22247pt}{0.0pt}\lxSVG@begingroup@{transform=matrix(1.0 0.0 0.0 1.0 -29.37 0)} \pgfsys@hbox{65}\lxSVG@closescope }}} \lxSVG@closescope }}} } \lxSVG@closescope \hbox to0.0pt{}{{ {}{}{}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}}}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}{\hbox to42.45pt{\vbox to6.25pt{\pgfpicture\makeatletter\hbox{\hskip 21.22247pt\lower-1.94443pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} {}{ {{}}\lx@inpgf@ignorespaces\hbox{\hbox{{\lxSVG@begingroup@{_scopebegin=1} {{}{}{{ {}{}}}{ {}{}} {{}{{\lx@inpgf@ignorespaces}}}{{}{\lx@inpgf@ignorespaces}}{}{{}{\lx@inpgf@ignorespaces}} {\lx@inpgf@ignorespaces }{{{{\lx@inpgf@ignorespaces}}\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{1.0}{0.0}{0.0}{1.0}{-21.22247pt}{0.0pt}\lxSVG@begingroup@{transform=matrix(1.0 0.0 0.0 1.0 -29.37 0)} \pgfsys@hbox{65}\lxSVG@closescope }}} \lxSVG@closescope }}} } \lxSVG@closescope \hbox to0.0pt{}{{ {}{}{}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}}}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}{\hbox to29.71pt{\vbox to4.38pt{\pgfpicture\makeatletter\hbox{\hskip 14.85706pt\lower-1.36108pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} {}{ {{}}\lx@inpgf@ignorespaces\hbox{\hbox{{\lxSVG@begingroup@{_scopebegin=1} {{}{}{{ {}{}}}{ {}{}} {{}{{\lx@inpgf@ignorespaces}}}{{}{\lx@inpgf@ignorespaces}}{}{{}{\lx@inpgf@ignorespaces}} {\lx@inpgf@ignorespaces }{{{{\lx@inpgf@ignorespaces}}\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{1.0}{0.0}{0.0}{1.0}{-14.85706pt}{0.0pt}\lxSVG@begingroup@{transform=matrix(1.0 0.0 0.0 1.0 -20.56 0)} \pgfsys@hbox{65}\lxSVG@closescope }}} \lxSVG@closescope }}} } \lxSVG@closescope \hbox to0.0pt{}{{ {}{}{}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}}}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}{\hbox to24.6pt{\vbox to3.13pt{\pgfpicture\makeatletter\hbox{\hskip 12.2988pt\lower-0.9722pt\hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} \lxSVG@begingroup@{stroke=#000000} \lxSVG@begingroup@{fill=#000000} \lxSVG@setlinewidth{\the\pgflinewidth}\lxSVG@begingroup@{stroke-width=0.4pt} \lx@inpgf@ignorespaces\nullfont\lxSVG@begingroup@{_scopebegin=1} \lxSVG@closescope \hbox to0.0pt{\lxSVG@begingroup@{_scopebegin=1} {}{ {{}}\lx@inpgf@ignorespaces\hbox{\hbox{{\lxSVG@begingroup@{_scopebegin=1} {{}{}{{ {}{}}}{ {}{}} {{}{{\lx@inpgf@ignorespaces}}}{{}{\lx@inpgf@ignorespaces}}{}{{}{\lx@inpgf@ignorespaces}} {\lx@inpgf@ignorespaces }{{{{\lx@inpgf@ignorespaces}}\lxSVG@begingroup@{_scopebegin=1} \lxSVG@transformcm{1.0}{0.0}{0.0}{1.0}{-12.2988pt}{0.0pt}\lxSVG@begingroup@{transform=matrix(1.0 0.0 0.0 1.0 -17.02 0)} \pgfsys@hbox{65}\lxSVG@closescope }}} \lxSVG@closescope }}} } \lxSVG@closescope \hbox to0.0pt{}{{ {}{}{}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}}}{\lx@inpgf@ignorespaces}{\lx@inpgf@ignorespaces}\hss}\lxSVG@discardpath\lxSVG@closescope \hss}}\lxSVG@closescope\endpgfpicture}}}\left(\tilde{\mathbf{f}}_{t}\right)\ \text{OR}\ \sigma\left(\tilde{\mathbf{f}}_{t}\right)\ , 𝐟 ~ t \displaystyle\tilde{\mathbf{f}}_{t}\ = 𝑾 𝐟 ​ 𝒙 t + 𝑹 𝐟 ​ 𝒉 t − 1 + 𝒃 𝐟 \displaystyle=\ \bm{W}_{\mathbf{f}}\ \bm{x}_{t}\ +\ \bm{R}_{\mathbf{f}}\ \bm{h}_{t-1}\ +\ \bm{b}_{\mathbf{f}}\ \ forget gate (39) = σ ⁡ ( 𝐨 ~ t ) , \displaystyle=\ \sigma\left(\tilde{\mathbf{o}}_{t}\right)\ , 𝐨 ~ t \displaystyle\tilde{\mathbf{o}}_{t}\ = 𝑾 𝐨 ​ 𝒙 t + 𝑹 𝐨 ​ 𝒉 t − 1 + 𝒃 𝐨 \displaystyle=\ \bm{W}_{\mathbf{o}}\ \bm{x}_{t}\ +\ \bm{R}_{\mathbf{o}}\ \bm{h}_{t-1}\ +\ \bm{b}_{\mathbf{o}}\ \ output gate (40)

Here, the cell input activation function φ \varphi is tanh \tanh , the hidden state activation function is the identity. φ \varphi helps stabilizing the recurrence.

Considering external gradient contribution δ 𝒉 t ext \delta^{\text{ext}}_{\bm{h}_{t}} from subsequent layers and recurrent gradient contribution δ 𝒉 t 𝑹 \delta^{\bm{R}}_{\bm{h}_{t}} from gradients from future states flowing over the cell interaction matrix 𝑹 \bm{R} , we obtain the recursive backward pass of sLSTM, where δ a \delta_{a} indicates gradients with respect to parameter / internal variable a a :

δ 𝒉 t \displaystyle\delta_{\bm{h}_{t}}\ = δ 𝒉 t e ​ x ​ t + δ 𝒉 t 𝑹 \displaystyle=\ \delta^{ext}_{\bm{h}_{t}}\ +\ \delta^{\bm{R}}_{\bm{h}_{t}} (41) δ 𝒄 t − 1 \displaystyle\delta_{\bm{c}_{t-1}}\ = 𝐟 t ⊙ δ 𝒄 t + 𝐨 t − 1 ⊙ 𝒏 t − 1 − 1 ⊙ δ 𝒉 t − 1 \displaystyle=\ \mathbf{f}_{t}\odot\delta_{\bm{c}_{t}}\ +\ \mathbf{o}_{t-1}\odot{\bm{n}_{t-1}}^{-1}\odot\delta_{\bm{h}_{t-1}} (42) δ 𝒏 t − 1 \displaystyle\delta_{\bm{n}_{t-1}}\ = 𝐟 t ⊙ δ 𝒏 t − 𝐨 t − 1 ⊙ 𝒄 t − 1 ⊙ 𝒏 t − 1 − 2 ⊙ δ 𝒉 t − 1 \displaystyle=\ \mathbf{f}_{t}\odot\delta_{\bm{n}_{t}}\ -\ \mathbf{o}_{t-1}\odot\bm{c}_{t-1}\odot\bm{n}_{t-1}^{-2}\odot\delta_{\bm{h}_{t-1}} (43) δ 𝐟 ~ t \displaystyle\delta_{\tilde{\mathbf{f}}_{t}}\ = 𝐟 t ′ ⊙ 𝒄 t − 1 ⊙ δ 𝒄 t + 𝐟 t ′ ⊙ 𝒏 t − 1 ⊙ δ 𝒏 t \displaystyle=\ \mathbf{f}^{\prime}_{t}\odot\bm{c}_{t-1}\odot\delta_{\bm{c}_{t}}\ +\ \mathbf{f}^{\prime}_{t}\odot\bm{n}_{t-1}\odot\delta_{\bm{n}_{t}} (44) δ 𝐢 ~ t \displaystyle\delta_{\tilde{\mathbf{i}}_{t}}\ = 𝐢 t ′ ⊙ 𝒛 t ⊙ δ 𝒄 t + 𝐢 t ′ ⊙ δ 𝒏 t \displaystyle=\ \mathbf{i}^{\prime}_{t}\odot\bm{z}_{t}\odot\delta_{\bm{c}_{t}}\ +\ \mathbf{i}^{\prime}_{t}\odot\delta_{\bm{n}_{t}} (45) δ 𝐳 ~ t \displaystyle\delta_{\tilde{\mathbf{z}}_{t}}\ = 𝐢 t ⊙ φ ′ ​ ( 𝐳 ~ t ) ⊙ δ 𝒄 t \displaystyle=\ \mathbf{i}_{t}\odot\varphi^{\prime}(\tilde{\mathbf{z}}_{t})\ \odot\delta_{\bm{c}_{t}} (46) δ 𝐨 ~ t \displaystyle\delta_{\tilde{\mathbf{o}}_{t}}\ = 𝐨 t ′ ⊙ 𝒄 t ⊙ 𝒏 t − 1 ⊙ δ 𝒉 t \displaystyle=\ \mathbf{o}^{\prime}_{t}\odot\bm{c}_{t}\odot\bm{n}^{-1}_{t}\odot\delta_{\bm{h}_{t}} (47) δ 𝐱 t \displaystyle\delta_{\mathbf{x}_{t}}\ = ∑ 𝐠 ∈ { 𝐟 , 𝐢 , 𝐳 , 𝐨 } 𝑾 𝐠 ⊤ ​ δ 𝐠 ~ t \displaystyle=\ \sum_{\mathbf{g}\in\{\mathbf{f},\mathbf{i},\mathbf{z},\mathbf{o}\}}\bm{W}^{\top}_{\mathbf{g}}\delta_{\tilde{\mathbf{g}}_{t}} (48) δ 𝒉 t − 1 𝑹 \displaystyle\delta^{\bm{R}}_{\bm{h}_{t-1}}\ = ∑ 𝐠 ∈ { 𝐟 , 𝐢 , 𝐳 , 𝐨 } 𝑹 𝐠 ⊤ ​ δ 𝐠 ~ t \displaystyle=\ \sum_{\mathbf{g}\in\{\mathbf{f},\mathbf{i},\mathbf{z},\mathbf{o}\}}\bm{R}^{\top}_{\mathbf{g}}\delta_{\tilde{\mathbf{g}}_{t}} (49) δ 𝑹 𝐠 ⊤ \displaystyle\delta_{\bm{R}_{\mathbf{g}}}^{\top}\ = ∑ t 𝒉 t − 1 ​ δ 𝐠 ~ t ⊤ , \displaystyle=\ \sum_{t}\bm{h}_{t-1}\delta_{\tilde{\mathbf{g}}_{t}}^{\top}\ , 𝐠 ∈ { 𝐢 , 𝐟 , 𝐳 , 𝐨 } \displaystyle\mathbf{g}\in\{\mathbf{i},\mathbf{f},\mathbf{z},\mathbf{o}\} (50) δ 𝑾 𝐠 ⊤ \displaystyle\delta_{\bm{W}_{\mathbf{g}}}^{\top}\ = ∑ t 𝒙 t ​ δ 𝐠 ~ t ⊤ , \displaystyle=\ \sum_{t}\bm{x}_{t}\delta_{\tilde{\mathbf{g}}_{t}}^{\top}\ , 𝐠 ∈ { 𝐢 , 𝐟 , 𝐳 , 𝐨 } \displaystyle\mathbf{g}\in\{\mathbf{i},\mathbf{f},\mathbf{z},\mathbf{o}\} (51)

with the derivatives of the respective gate activation function 𝐢 t ′ = exp ′ ⁡ ( 𝐢 ~ t ) = exp ⁡ ( 𝐢 ~ t ) = 𝐢 t \mathbf{i}^{\prime}_{t}=\exp^{\prime}(\tilde{\mathbf{i}}_{t})=\exp(\tilde{\mathbf{i}}_{t})=\mathbf{i}_{t} , 𝐨 t ′ = σ ′ ​ ( 𝐨 ~ t ) \mathbf{o}^{\prime}_{t}=\sigma^{\prime}(\tilde{\mathbf{o}}_{t}) , and 𝐟 t ′ = σ ′ ​ ( 𝐟 ~ t ) \mathbf{f}^{\prime}_{t}=\sigma^{\prime}(\tilde{\mathbf{f}}_{t}) or 𝐟 t ′ = 𝐟 t \mathbf{f}^{\prime}_{t}=\mathbf{f}_{t} depending on the forget gate activation. φ ′ ​ ( z ) \varphi^{\prime}(z) is the derivative of the cell input activation function φ ⁡ ( z ) \varphi(z) .

The matrices 𝑹 𝒛 \bm{R}_{\bm{z}} , 𝑹 𝐢 \bm{R}_{\mathbf{i}} , 𝑹 𝐟 \bm{R}_{\mathbf{f}} , 𝑹 𝐨 \bm{R}_{\mathbf{o}} are block-diagonal which is analogous to multiple heads in the mLSTM. This way, the parameters reduce to d 2 / ( N h ) d^{2}/(N_{h}) , where N h N_{h} is the number of heads, limiting the cell interactions to individual heads. This parameter efficient formulation of cell interactions together with the exponential gating is called the new memory mixing. Finally, to stabilize the backward pass, we clip the magnitude of δ 𝒉 t 𝑹 \delta^{\bm{R}}_{\bm{h}_{t}} to 10 10 , as a means to prohibit exploding gradients for long context lengths.

##### Proof of Equivalence for sLSTM Stabilized Version.