RetNet

Summary

• Introduced constant forgetting factor $\gamma$.
• Derived parallel and recurrent representation.
• Derived chunk-wise block-parallel form for efficient training.

Motivation

• Linear attention
  • Replaces $\exp$ with $\varphi (\cdot )$
  • Worse performance than transformers.
• Recurrent models
  • Sacrifices training parallelism.
• State space models.
  • Training parallelism & performance.

Methodology

Both recurrence and parallelism.

$$\begin{array}{c}{\mathbf{S}}_t=\mathbf{A}{\mathbf{S}}_{t-1}+{\mathbf{v}}_t{\mathbf{k}}_t^{\top }\\ {\mathbf{o}}_t={\mathbf{S}}_t{\mathbf{q}}_t\end{array}$$

$\mathbf{A}\in {\mathbb{R}}^{d\times d}$ is diagonalized

$$\mathbf{A}=\Lambda \left(\gamma e^{i\theta }\right){\Lambda }^{-1},\gamma ,\theta \in {\mathbb{R}}^d$$

$\Lambda$ can be absorbed into ${\mathbf{W}}_q$ and ${\mathbf{W}}_k$. Therefore, the formula can be simplified as

$$\text{Retention}(\mathbf{X})=\left(\mathbf{Q}{\mathbf{K}}^{\top }\odot \mathbf{D}\right)\mathbf{V}\text{where}{\mathbf{D}}_{nm}=\{\begin{array}{rlr}& {\gamma }^{n-m},& n\ge m\\ & 0,& n<m\end{array}$$

or in recurrent form

$${\mathbf{S}}_t=\gamma {\mathbf{S}}_{t-1}+{\mathbf{v}}_t{\mathbf{k}}_t^{\top },{\mathbf{o}}_t={\mathbf{S}}_t{\mathbf{q}}_t$$