Differential Transformer V2

We compare DIFF V2 with DIFF V1 below:

(For simplicity, we omit the batch dimension and assume that both the input and output of the following flash_attn_func are three-dimensional tensors (tokens, heads, head dimension). Heads belonging to the same GQA group are arranged contiguously in the output)

Note DIFF V2 subtracts two heads that are in the same GQA group, which means they share the same key and value. This is crucial to performance.

def DiffAttnV1( layer_index, q1, q2, k1, k2, v, lam_q1, lam_k1, lam_q2, lam_k2):
    attn1 = flash_attn_func(q1, k1, v)
    attn2 = flash_attn_func(q2, k2, v)
    lam_init = 0.8 - 0.6 * exp(-0.3 * layer_index)
    lam1 = exp(sum(lam_q1 * lam_k1))
    lam2 = exp(sum(lam_q2 * lam_k2))
    lam = lam1 - lam2 + lam_init
    attn = attn1 - lam * attn2
    attn = rmsnorm(attn)
    attn = attn * (1 - lam_init)
    return attn

def DiffAttnV2( q, k, v, lam):
    attn = flash_attn_func(q, k, v)
    attn1, attn2 = (attn[:, 0::2], attn[:, 1::2])
    lam_val = sigmoid(lam)
    attn = attn1 - lam_val * attn2
    return attn

DIFF V2 doubles number of query heads while maintaining number of key value heads, and the extra dimension is reduced back to h*d after the differential operation so the projection remains the same as baseline Transformer. Since LLM decoding is typically memory-bound, this design allows DIFF V2 to achieve decoding speeds on par with standard Transformer. Besides, since head dimension is aligned between query, key and value, there is no need for custom attention kernels for DIFF V2.

In DIFF V1, we find per-head RMSNorm leads to massive gradients and numerical instability. In DIFF V2, after removing the per-head RMSNorm, the gradient norm scale becomes comparable to that of Transformer. We introduce a projected lambda for each token and each head to control the context RMS, which helps eliminate attention sinks and improve training stability.

Experiments on production-scale LLMs show notably lower language modeling loss compared to Transformer and reduced loss/gradient spikes during training. Theoretically, DIFF V2 can save approximately 25% of the attention-module parameters by explicitly constructing the differential operation.