Differential Transformer V2

(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. See design ablations section and Github code.

Full code at: unilm/Diff-Transformer/Diff-Transformer-V2 at master · microsoft/unilm In the script, h represents number of query heads, h_kv represents number of key-value heads, and d means head dimension. The in DIFF V2 is projected from for each token each head.

DIFF V2 introduces additional query heads compared to the baseline Transformer, but does not increase the number of key-value (KV) heads . 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 contrast, DIFF V1 can be slower during decoding because the value cache must be loaded twice, and a custom attention kernel is needed. DIFF V2 can also increase the arithmetic intensity of the attention module during decoding.

During pretraining , when using cutting-edge FlashAttention kernels on H-series and B-series GPUs, the throughput reduction introduced by DIFF V2 is negligible. For long-sequence prefilling , we recommend combining DIFF V2 with techniques such as YOCO (also used in Gemma 3n ), which already reduces prefilling complexity to linear time with respect to sequence length.

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.

An alternative perspective is to compare DIFF V2 with a Transformer that has the same query dimension 2h*d . Under this comparison, both models exhibit same attention kernel speed, while DIFF V2 has less parameters and flops in output projection.

The weights are non-negative and sum to 1 ( ). Assume the value vectors are uncorrelated and have an RMS of 1, the Context RMS is strictly bounded in range however the attention distribution changes :

A typical phenomenon is that when DIFF V1 is pre-trained at a large learning rate, the gradient norm experiences a larger increase compared to Transformer in the later stages, along with higher variance. In DIFF V2, after removing the per-head RMSNorm, the gradient norm scale becomes comparable to that of Transformer, and the gradient norm spike is reduced (will be discussed further below).

In the standard Scaled Dot-Product Attention (SDPA), let be the queries, keys, and values. The context vector is defined as:

Where is the attention weight matrix. Let's focus on a single row of , denoted as , which is a weighted sum of value vectors :

We define the Context RMS (Root Mean Square) to represent the magnitude of this output: