Differential Transformer V2: Faster Decoding and Improved Stability

Differential Transformer V2

Tianzhu Ye, Li Dong, Yutao Sun, and Furu Wei at Microsoft introduced Differential Transformer V2, a novel attention mechanism designed to improve LLM training and decoding efficiency. DIFF V2 maintains decoding speeds comparable to standard Transformers while reducing language modeling loss, achieving a gap of 0.02 to 0.03 at 1T training tokens.

Why This Matters

Current transformer models struggle with scaling due to computational costs and numerical instability, particularly during pretraining with large learning rates. Ideal transformer models would achieve higher throughput and maintain stability, but existing architectures often require complex custom kernels or suffer from gradient spikes. These issues limit scalability and increase the cost of training large language models.

Key Insights

• FlashAttention Kernels: DIFF V2 avoids the need for custom attention kernels, unlike DIFF V1, by aligning head dimensions for query, key, and value.
• Context RMS Constraint: The original Softmax attention mechanism constrains the context RMS, potentially leading to instability; DIFF V2 addresses this by allowing the lower bound to approach zero.
• Parameter Efficiency: DIFF V2 saves approximately 25% of the attention module parameters compared to a standard Transformer with equivalent performance, enabling parameter reallocation.

Working Example

def DiffAttnV2(
q, k, v, lam
):
"""
q: (N, 2h, d)
k: (N, h_kv, d)
v: (N, h_kv, d)
lam: (N, h, 1)
"""
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

Practical Applications

• Large Language Models: Gemma 3n leverages techniques like YOCO alongside DIFF V2 to reduce prefilling complexity.
• Training Instability: DIFF V2’s design reduces gradient spikes during pretraining, allowing for the use of larger learning rates.

References:

• https://huggingface.co/blog/microsoft/diff-attn-v2