Differential Transformer V2

Summary

Differential Transformer V2 introduces an architectural refinement to differential attention that pairs and subtracts query heads within the same grouped-query attention group. The architecture doubles query heads while preserving key-value head counts, which avoids increasing key-value cache demands and removes the need for custom kernels. DIFF V2 replaces the per-head RMSNorm of its predecessor with a per-token, per-head projected lambda passed through a sigmoid function. This formulation expands the context root-mean-square range to between zero and root two, preventing massive gradient magnifications, mitigating attention sinks, and stabilizing large learning rate pretraining. Ablations confirm that sharing keys and values across subtracted heads and constraining lambda scaling are necessary to prevent training instability and elevated loss.

Takeaways

1. 1

   DIFF V2 preserves the standard Transformer key-value head count and output projection dimensions, avoiding key-value cache expansion and custom decoding kernels.

2. 2

   Lowering the context root-mean-square lower bound to zero via sigmoid-scaled differential attention eliminates attention sinks and stabilizes pretraining gradients.

3. 3

   Ablations show that failing to share keys and values across subtracted head pairs or omitting sigmoid activation on lambda leads to higher loss and training instability.