Differential Transformers Explained

The Basics

Before diving into the new Differential Transformer, let's go over how a traditional Transformer works. At its core, Transformers use an attention mechanism to allow a model to focus on specific parts of an input sequence. This attention is computed using a softmax function:

$$Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$$

Where:

• $Q$ is the query matrix
• $K$ is the key matrix
• $V$ is the value matrix
• $d_k$ is the dimensionality of the key matrix

This mechanism assigns weights to different input tokens based on their relevance. Despite its success, current Transformers tend to be very distracted. The standard softmax function tends to over-allocate attention to irrelevant parts of the context. In long-context sequences, the model can focus too broadly, leading to inefficient learning. This broad focus also negatively impacts in-context learning.

The Differential Transformer addresses these challenges by introducing a new mechanism. Instead of relying on a single attention map, it calculates two distinct attention maps:

$$A_{diff} = softmax(A_1) - softmax(A_2)$$

Yes, that's right, it's that simple. This approach effectively removes redundant or noisy attention, promoting sparser and more focused attention. In turn, this prevents over-allocating attention to irrelevant tokens and allows the model to better manage long sequences and complex in-context learning scenarios.

Key Benefits:

• Sparse Attention Patterns: By reducing redundant attention, the model can better focus on critical parts of the input sequence.
• Improved Long-Context Modeling: Differential Attention allows the model to handle longer contexts more effectively, improving tasks like document summarization and question answering.
• In-Context Learning: The differential attention mechanism dynamically adapts based on the input context, enhancing the model's ability to learn from examples within the input.
• Hallucination Mitigation: In generation tasks, the DIFF Transformer reduces hallucinations by focusing more accurately on relevant context, leading to more coherent outputs.

The DIFF Transformer has broad applications, particularly in:

• Handling long texts while focusing on the core information in Text Summarization tasks
• Improved performance in QA systems which require nuanced understanding of context
• Robust Generation, mitigating hallucinations in current models (GPT, Claude, Llama, etc.)

Implementation

Start by modifying the attention mechanism within a Transformer architecture. Instead of directly computing the attention using the standard softmax approach, compute two separate attention maps and subtract them to generate a differential attention map:

Here's a high-level snapshot written in Python:

def diff_attention(Q, K, V):     
    A1 = softmax(Q @ K.T / sqrt(d_k))  # first attention map    
    A2 = softmax(Q @ K.T / sqrt(d_k))  # second attention map    
    diff = A1 - A2  # differential attention    
    return diff @ V`

This approach allows you to integrate differential attention into any Transformer-based architecture.

That's a Wrap

DIFF Transformers are a significant leap forward. By refining the attention mechanism, they address key weaknesses we all encounter in the traditional Transformer architecture, leading to more efficient, focused, and context-aware models. Implementing these ideas can enhance the performance of large-scale language models in your applications, from NLP tasks to GenAI.

Thanks for reading!

References

• ARXIV - Differential Transformer