DeltaNet

• https://sustcsonglin.github.io/blog/2024/deltanet-1/

Vanilla Softmax

$\begin{array}{rlrlrl}\mathrm{Parallel}\mathrm{training}:& & & \mathbf{O}=\mathrm{softmax}(\mathbf{Q}{\mathbf{K}}^{\top }\odot \mathbf{M})\mathbf{V}& & \in {\mathbb{R}}^{L\times d}\\ \mathrm{Iterative}\mathrm{inference}:& & & {\mathbf{o}}_{\mathbf{t}}=\sum _{j=1}^t\frac{\exp ({\mathbf{q}}_t^{\top }{\mathbf{k}}_j)}{{\sum }_{l=1}^t\exp ({\mathbf{q}}_t^{\top }{\mathbf{k}}_l)}{\mathbf{v}}_j& & \in {\mathbb{R}}^d\end{array}$

Linear Attention

$\begin{array}{rlrlrl}\mathrm{Parallel}\mathrm{training}：& & & \mathbf{O}=(\mathbf{Q}{\mathbf{K}}^{\top }\odot \mathbf{M})\mathbf{V}& & \in {\mathbb{R}}^{L\times d}\\ \mathrm{Iterative}\mathrm{inference}：& & & {\mathbf{o}}_{\mathbf{t}}=\sum _{j=1}^t({\mathbf{q}}_t^{\top }{\mathbf{k}}_j){\mathbf{v}}_j& & \in {\mathbb{R}}^d\end{array}$

• While removing softmax alone doesn’t immediately reduce computational complexity, it enables a crucial mathematical property: linearity.
• Linearity gives us associativity ⇒ chunkwise parallel form for prefill

Linear Attention as a linear RNN

• For inference, we can rearrange things $\begin{array}{rlrlrlrlrl}& & & & {\mathbf{o}}_{\mathbf{t}}=\sum _{j=1}^t{\mathbf{v}}_j({\mathbf{k}}_j^{\top }{\mathbf{q}}_t)& & & & & (transposition){\mathbf{k}}_j^{\top }{\mathbf{q}}_t={\mathbf{q}}_t^{\top }{\mathbf{k}}_j\in \mathbb{R}\\ & & & & =(\sum _{j=1}^t{\mathbf{v}}_j{\mathbf{k}}_j^{\top }){\mathbf{q}}_t& & & & & \text{By associativity}\end{array}$

• Let’s define a state matrix ${\mathbf{S}}_t={\sum }_{j=1}^t{\mathbf{v}}_j{\mathbf{k}}_j^{\top }$

  • Then we have a clear recurrent relationship
  • ${\mathbf{S}}_t={\mathbf{S}}_{t-1}+{\mathbf{v}}_t{\mathbf{k}}_t^{\top }\in {\mathbb{R}}^{d\times d},{\mathbf{o}}_t={\mathbf{S}}_t{\mathbf{q}}_t\in {\mathbb{R}}^d$

• Linear attention is essentially a linear RNN with a matrix-valued state $S$ that accumulates key-value outer-products, and keeps track of a compressed view of size $O(d^2)$ of the history.

• This is called a state size expansion from $O(d)$ to $O(d^2)$

• By casting linear attention as an RNN,

  • we reduce inference cost from $O(Ld)$ to $O(d^2)$
  • we reduce the space complexity to $O(Ld)$ to O($d^2$)

Limitations of Linear Attention - defining retrieval error

• The fixed-size state matrix in linear attention means it cannot perfectly preserve all historical information, making exact retrieval particularly challenging.

Retrieval error

• Linear attention implements a key-value associative memory, which is the sum of outer products between keys and values $\mathbf{S}=\sum {\mathbf{v}}_i{\mathbf{k}}_i^{\top }$

• Assuming all keys are normalized to unit length, when we try to retrieve a value associated with a specific key $k_j$ , i.e. the dot-product between the state matrix $S$ and the key $k_j$ should ideally give us back $v_j$ (because $k_j^{\top }{k}_j=1$)

  • $\begin{array}{rl}\mathbf{S}{\mathbf{k}}_j& =\sum {\mathbf{v}}_i({\mathbf{k}}_i^{\top }{\mathbf{k}}_j)\\ & ={\mathbf{v}}_j+\underset{\text{retrieval error}}{\underbrace{\sum _{i\ne j}({\mathbf{k}}_i^{\top }{\mathbf{k}}_j){\mathbf{v}}_i}}\end{array}$

• To minimize the retrieval error term, we need all they key vectors to be orthogonal ($k_i^{\top }{k}_j=0$)

  • However, in a space of dimension $d$, we can not define more than $d$ vectors that are all orthogonal to each other.
  • $d$ is the head-dimension, and it has been shown that increasing head-dimension improves performance, as it gives more space for storing distinct-key value pairs
  • There is a tradeoff between increasing head-dim, and allowing chunked linear attention to be hardware friendly (we want tiles that are small enough to keep in registers)

Gating or forgetting as a mechanism to improve retrieval

• In this key-value associative memory system, we can only add new key-value associations without the ability to erase existing information. As sequences grow longer, this leads to accumulating “retrieval errors” that degrade performance.
• We can narrow the performance gap with standard attention in language modeling tasks by incorporating a forgetting mechanism
  • ${\mathbf{S}}_t={\mathbf{G}}_t\odot {\mathbf{S}}_{t-1}+{\mathbf{v}}_t{\mathbf{k}}_t^{\top }$, where ${\mathbf{G}}_t\in {\mathbb{R}}^{d\times d}$
  • There are multiple different structured parameterization (for parameter efficiency), often with outer product structure.
    • ${\mathbf{G}}_t={\beta }_{\mathbf{t}}{{\alpha }_{\mathbf{t}}}^{\top }$ (decaying fast weight)
    • ${\mathbf{G}}_t=1{{\alpha }_{\mathbf{t}}}^{\top }$ (GLA)
    • ${\mathbf{G}}_t=\exp (-({\mathbf{\Delta }}_{\mathbf{t}}{1}^{\top })\odot \exp (A))$ (Mamba)
    • ${\mathbf{G}}_t={\gamma }_{t}11^{\top }$ (Mamba 2)

DeltaNet: Linear Attention with Delta Rule

What is the Delta Rule

• Very simple error-correction learning principle
  • principle: adjust the model’s parameters based on the difference (delta) between what we want (target) and what we actually get (prediction).
• Imagine teaching a child to aim at a target. If they shoot too far to the left, you’d tell them to adjust right; too far right, adjust left. T
• The size of the adjustment depends on
  • the delta size
  • the magnitude of the input itself (in the linear regression case)

Pseudocode

import numpy as np
 
def delta_rule(x, y, epochs=100, lr=0.1):
    """
    Simple delta rule implementation
    x: input features (N samples by D features)
    y: target values (N samples)
    """
    # Initialize weights
    w = np.zeros(x.shape[1])
    
    # Train
    for _ in range(epochs):
        for i in range(len(x)):
            # Forward pass
            pred = np.dot(x[i], w)
            
            # Compute error
            error = y[i] - pred
            
            # Update weights
            w += lr * error * x[i]
            
    return w

DeltaNet

DeltaNet applies this error-correction principle to linear attention. Instead of simply accumulating key-value outer product, it updates its state based on prediction errors:

\begin{align*} \mathbf{S}_{t} &= \mathbf{S}_{t-1} - \beta_t(\mathbf{S}_{t-1} \mathbf{k}_t - \mathbf{v}_t)\mathbf{k}_t^\top \\ &= \mathbf{S}_{t-1} - \beta_t \mathbf{S}_{t-1} \mathbf{k}_t \mathbf{k}_t^\top + \beta_t \mathbf{v}_t \mathbf{k}_t^\top \end{align*}

• The parallel to the Delta Rule becomes clear when we break down the components:

  • ${\beta }_t\in \mathbb{R}$ acts as the learning rate
  • ${\mathbf{k}}_t\in {\mathbb{R}}^d$ is the input data
  • ${\mathbf{v}}_t\in {\mathbb{R}}^d$ is the target
  • ${\mathbf{S}}_{t-1}{\mathbf{k}}_t\in {\mathbb{R}}^d$ is our current prediction (trying to retrieve ${\mathbf{v}}_t$ from the state matrix)

• Think of ${\mathbf{S}}_{t-1}{\mathbf{k}}_t$ as retrieving the “old value” associated with the current key ${\mathbf{k}}_t$ from memory. When we encounter a newly associated value ${\mathbf{v}}_t$ for the same key, rather than blindly overwriting, we make a careful update: \begin{align*} \mathbf{v}_t^{\text{new}} &= (1-\beta_t) \mathbf{v}_t^{\text{old}} + \beta_t \mathbf{v}_t, \\ \mathbf{S}_t &= \mathbf{S}_{t-1} - \underbrace{\mathbf{v}_t^{\text{old}} \mathbf{k}_t^\top}_{\text{erase}} + \underbrace{\mathbf{v}_t^{\text{new}} \mathbf{k}_t^\top}_{\text{write}} \end{align*}

• where ${\mathbf{v}}_t^{\text{new}}$ is a learned combination of the old and current values, controlled by a dynamic ${\beta }_t\in (0,1)$ : when ${\beta }_t=0$, the memory content remains intact, and when ${\beta }_t=1$, we completely replace the old associated value with the new one.

DeltaNet as the gradient update for MSE

• DeltaNet’s update rule can be derived by sequentially minimizing the mean squared error (MSE) between the desired output and the predicted output at each time step using gradient descent: ${\mathcal{L}}_t(\mathbf{S})=\frac{1}{2}\Vert \mathbf{S}{\mathbf{k}}_t-{\mathbf{v}}_t{\Vert }^2$
• Applying gradient descent to minimize this MSE loss gives:

$\begin{array}{rl}{\mathbf{S}}_t& ={\mathbf{S}}_{t-1}-{\eta }_t\nabla {\mathcal{L}}_t({\mathbf{S}}_{t-1})\\ & ={\mathbf{S}}_{t-1}-{\eta }_t\left({\mathbf{S}}_{t-1}{\mathbf{k}}_t-{\mathbf{v}}_t\right){\mathbf{k}}_t^{\top }\end{array}$

• When the learning rate is set to ${\beta }_t$, we recover DeltaNet
• In contrast, vanilla linear attention employs a linear loss function: ${\mathcal{L}}_t^{\prime }(\mathbf{S})=-⟨\mathbf{S}{\mathbf{k}}_t,{\mathbf{v}}_t⟩$