DeltaNet

Parallelizing Linear Transformers with the Delta Rule over Sequence Length (arxiv link).

Motivation

• Linear attention is pure addition.
• Forgetting factor $\gamma$ cannot fully forget previous context.
• Model should learn to remove less important memory.

The Delta Rule

$${\mathbf{S}}_t={\mathbf{S}}_{t-1}-{\beta }_t\left({\mathbf{S}}_{t-1}{\mathbf{k}}_t-{\mathbf{v}}_t\right){\mathbf{k}}_t^{\top }$$

This can be regarded as optimizing an online regression loss using a single step of SGD:

$${\mathcal{L}}_t(\mathbf{S})=\frac{1}{2}{\Vert \mathbf{S}{\mathbf{k}}_t-{\mathbf{v}}_t\Vert }^2,{\mathbf{S}}_t={\mathbf{S}}_{t-1}-{\beta }_t{\nabla }_{{\mathbf{S}}_{t-1}}{\mathcal{L}}_t({\mathbf{S}}_{t-1})={\mathbf{S}}_{t-1}-{\beta }_t\left({\mathbf{S}}_{t-1}{\mathbf{k}}_t-{\mathbf{v}}_t\right)$$

An alternative interpretation for DeltaNet is from the perspective of key-value retrieval:

1. Retrieves the old value using the current key ${\mathbf{v}}_t^{\text{old}}={\mathbf{S}}_{t-1}{\mathbf{k}}_t$.
2. Obtains a new value ${\mathbf{v}}_t^{\text{new}}$ by interpolating between the old value and the current new value ${\mathbf{v}}_t$, which replaces ${\mathbf{v}}_t^{\text{old}}$ in the memory.

$${\mathbf{v}}_t^{\text{new}}={\beta }_t{\mathbf{v}}_t+(1-{\beta }_t){\mathbf{v}}_t^{\text{old}}{\mathbf{S}}_t={\mathbf{S}}_{t-1}\underset{\text{remove}}{\underbrace{-{\mathbf{v}}_t^{\text{old}}{\mathbf{k}}_t^{\top }}}\underset{\text{write}}{\underbrace{+{\mathbf{v}}_t^{\text{new}}{\mathbf{k}}_t^{\top }}}$$

DeltaNet Transformer

---
config:
  look: handDrawn
---
flowchart BT
in(("in"))
qk_proj("Linear")
qk_conv("Conv<sup>*</sup>")
qk_silu("SiLU")
qk_l2norm("L2 Norm")
v_proj("Linear")
v_conv("Conv<sup>*</sup>")
v_silu("SiLU")
beta_proj("Linear")
beta_sigmoid("Sigmoid")
delta_rule("Delta Rule")
rmsnorm("RMSNorm")
o_proj("Linear")
out(("out"))


in --> qk_proj --> qk_conv --> qk_silu --> qk_l2norm -- "q, k" --> delta_rule
in --> v_proj --> v_conv --> v_silu -- "v" --> delta_rule
in --> beta_proj --> beta_sigmoid -- "beta" --> delta_rule
delta_rule --> rmsnorm --> o_proj
o_proj --> out

• The convolution layers are depthwise-separable convolution that generalizes the shift SSM.
• Interleave DeltaNet layers and SWA layers.
• Two layers with global attention: 2^nd and the $(\frac{N}{2}+1)$-th layer.