The Key to Going Linear: Analysis-Driven Transformer Linearization

Core Methodology§

The paper aims to linearize pretrained transformers (frozen backbone) by replacing softmax attention with a linear attention mechanism. The key insight is that softmax implicitly computes a rank-1 orthogonal projection of the value vectors, conditioned on the query. This is formalized as:

$$ \text{softmax}(QK^\top)V = \sum_{i} a_i v_i, \quad a_i = \frac{\exp(q^\top k_i)}{\sum_j \exp(q^\top k_j)} $$

The authors show that this can be approximated by a linear state update if the state is updated with a delta rule (akin to a linear gating with key-dependent orthogonal projections). In contrast, purely gated accumulation (e.g., linear attention with cumulative sum) fails to capture the dynamic reweighting of softmax.

They identify two sources of approximation error: (1) non-stationary token representations due to layer norms and residual connections, and (2) the need for unbounded memory state. To mitigate these, they introduce three structural interventions:

• Sink tokens: Additional learnable tokens that absorb attention mass and stabilize the state.
• Short convolutions: A 1D causal convolution applied to key and value sequences before linearization, smoothing out local context.
• Fixed-budget cache routing: A mechanism to bound the state size by routing tokens to a fixed number of buckets, ensuring O(1) memory per layer.

Implementation Details§

Given a frozen transformer with softmax attention, they replace it with:

$$ \text{State}_t = \text{State}_{t-1} + \Delta(\text{State}_{t-1}, k_t, v_t) $$

where $\Delta$ is either a gated accumulation or a delta rule. The delta rule is:

$$ \Delta S = (v_t - S_{t-1}\phi(k_t)) \otimes \phi(k_t) $$

with $\phi$ a kernel feature map (here, simply identity). Sink tokens are prepended to the sequence as learned embeddings. Short convolutions are applied per channel with kernel size 3. Cache routing uses a fixed number of slots (e.g., 64) and routes each token to the nearest slot via learned keys.

Code Snippet§

import torch
import torch.nn as nn

class LinearAttentionLayer(nn.Module):
    def __init__(self, d_model, num_sinks=1, conv_kernel=3, num_slots=64):
        super().__init__()
        self.sink_tokens = nn.Parameter(torch.randn(1, num_sinks, d_model))
        self.short_conv = nn.Conv1d(d_model, d_model, kernel_size=conv_kernel, padding=conv_kernel-1, bias=False)
        self.state = None
        self.num_slots = num_slots
        self.slot_keys = nn.Parameter(torch.randn(num_slots, d_model))

    def forward(self, q, k, v, cache=None):
        # Add sink tokens
        batch_size = q.shape[0]
        sinks = self.sink_tokens.expand(batch_size, -1, -1)
        k = torch.cat([sinks, k], dim=1)
        v = torch.cat([sinks, v], dim=1)

        # Short convolution (apply on last dim)
        k = self.short_conv(k.transpose(1,2)).transpose(1,2)[:, :-self.short_conv.padding[0], :]
        v = self.short_conv(v.transpose(1,2)).transpose(1,2)[:, :-self.short_conv.padding[0], :]

        # Fixed-budget cache routing
        # Simplified: assign each token to nearest slot, then accumulate
        # For brevity, assume state is maintained externally
        output = []
        for t in range(q.shape[1]):
            # delta update
            state = cache if cache is not None else torch.zeros(batch_size, self.num_slots, v.shape[-1], device=q.device)
            # Compute routing weights
            scores = torch.matmul(k[:, t:t+1, :], self.slot_keys.T)  # (B, 1, num_slots)
            weights = torch.softmax(scores, dim=-1)
            # Update state: delta rule
            state = state + torch.matmul(weights.transpose(1,2), (v[:, t:t+1, :] - state))
            # Attend using state
            out = torch.matmul(q[:, t:t+1, :], state.transpose(1,2))  # simplified
            output.append(out)
        return torch.cat(output, dim=1), state

Results§

The approach is scaled to LLaMA and Qwen models up to 32B parameters. It outperforms prior post hoc linearization baselines on MMLU and achieves long-context retrieval accuracy comparable to complex adaptive-caching frameworks (e.g., Infini-Attention).