Differential Transformer：通過差分注意力機制提升大語言模型性能
Transformer模型已經成為大語言模型（LLMs）的標準架構，但研究表明這些模型在準確檢索關鍵信息方面仍面臨挑戰。今天介紹一篇名叫Differential Transformer的論文，論文的作者觀察到一個關鍵問題：傳統Transformer模型傾向于過分關注不相關的上下文信息，這種"注意力噪聲"會影響模型的性能。
在這篇論文中，作者注意到transformer模型傾向于關注不相關的上下文。為了放大相關上下文的注意力分數，他們提出了一個新的注意力模型，稱為差分注意力模型。在這個模型中，他們將查詢和鍵值向量分成兩組，并計算兩個子注意力分數。
差分注意力機制
差分注意力機制（Differential Attention）的核心思想是通過計算兩個獨立的注意力圖譜之差來消除注意力噪聲。這種設計借鑒了電氣工程中差分放大器的原理，通過對比兩個信號的差異來消除共模噪聲。
讓我們看看論文中的第一個方程：
方程(1)
方程(1)顯示，我們首先像標準注意力計算一樣計算Q、K和V張量。關鍵點是我們將Q和K張量分成Q1、Q2和K1、K2子張量。
論文中輸入X、Q1、Q2、K1、K2和V張量的形狀
根據論文，Q和K張量的形狀應該是Nx2d，因為Q1、Q2、K1和K2將是Nxd。輸入X的形狀是Nxd_model，這是論文中的嵌入維度。這就是為什么W_Q、W_K和W_V的可學習參數的形狀必須是d_modelx2d。
論文中用于lambda計算的方程(2)
方程(2)展示了如何計算可學習參數lambda。在這個方程中有一個初始lambda參數。lambda是一個標量參數，但lambda_q1、lambda_k1、lambda_q2和lambda_k2是向量。這一點很關鍵。向量lambda_q和lambda_k的運算是點積。
用于lambda初始化的方程(3)
實驗結果與性能提升
論文的實驗表明，相比傳統Transformer：
DIFF Transformer只需要約65%的模型參數量即可達到相同的性能，在訓練token數量方面也只需要約65%就能達到相同效果
在Needle-In-A-Haystack測試中：4K上下文長度：DIFF Transformer在多目標檢索任務中保持85%準確率；64K上下文長度：在深度為25%的位置檢測時，比傳統Transformer提升了76%的準確率
Python實現
下面我們根據論文的公式來做一個簡單的實現，首先方程(3)展示了我們如何計算lambda_initial變量。現在讓我們把方程轉換成Python代碼：
def lambda_init_fn(depth):
return 0.8 - 0.6 * math.exp(-0.3 * depth)然后再寫一個簡單的Python函數，使用方程(3)。
class DifferentialAttention(nn.Module):
def __init__(self, dim_model: int, head_nums: int, depth: int):
super().__init__()
self.head_dim = dim_model // head_nums
self.Q = nn.Linear(dim_model, 2 * self.head_dim, bias=False)
self.K = nn.Linear(dim_model, 2 * self.head_dim, bias=False)
self.V = nn.Linear(dim_model, 2 * self.head_dim, bias=False)
self.scale = self.head_dim ** -0.5
self.depth = depth
self.lambda_q1 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
self.lambda_q2 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
self.lambda_k1 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
self.lambda_k2 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
self.rotary_emb = RotaryEmbedding(self.head_dim * 2)在DifferentialAttention類中，我們實現了一個多頭差分注意力機制。有dim_model（嵌入維度）、head_nums和depth參數。為Q1、Q2、K1和K2聲明了四個lambda可學習參數，并使用均值為0、標準差為0.1的隨機正態分布初始化它們。
def forward(self, x):
lambda_init = lambda_init_fn(self.depth)
Q = self.Q(x)
K = self.K(x)
seq_len = x.shape[1]
cos, sin = self.rotary_emb(seq_len, device=x.device)
Q, K = apply_rotary_pos_emb(Q, K, cos, sin)
Q1, Q2 = Q.chunk(2, dim=-1)
K1, K2 = K.chunk(2, dim=-1)
V = self.V(x)
A1 = Q1 @ K1.transpose(-2, -1) * self.scale
A2 = Q2 @ K2.transpose(-2, -1) * self.scale
lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1).float()).type_as(Q1)
lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1).float()).type_as(Q2)
lambda_ = lambda_1 - lambda_2 + lambda_init
return (F.softmax(A1, dim=-1) - lambda_ * F.softmax(A2, dim=-1)) @ Vforward方法很直觀。我分別實現了方程(1)和方程(2)。forward方法直接實現了論文中的偽代碼。
多頭差分注意力架構和偽代碼
class MultiHeadDifferentialAttention(nn.Module):
def __init__(self, dim_model: int, head_nums: int, depth: int):
super().__init__()
self.heads = nn.ModuleList([DifferentialAttention(dim_model, head_nums, depth) for _ in range(head_nums)])
self.group_norm = RMSNorm(dim_model)
self.output = nn.Linear(2 * dim_model, dim_model, bias=False)
self.lambda_init = lambda_init_fn(depth)
def forward(self, x):
o = torch.cat([self.group_norm(h(x)) for h in self.heads], dim=-1)
o = o * (1 - self.lambda_init)
return self.output(o)MultiHeadDifferentialAttention類是根據論文中的偽代碼編寫的。這里使用了RMSNorm而不是GroupNorm。
論文中使用多頭差分注意力機制的語言模型的方程
最后使用實現的MultiHeadDifferentialAttention機制構建一個transformer解碼器。
class DifferentialTransformer(nn.Module):
def __init__(self, dim: int, depth: int, heads: int = 8, head_dim: int = 64, vocab_size: int = 10000):
super().__init__()
self.vocab_size = vocab_size
self.layers = nn.ModuleList([
MultiHeadDifferentialAttention(dim, heads, depth_idx)
for depth_idx in range(depth)
])
self.ln1 = RMSNorm(dim)
self.ln2 = RMSNorm(dim)
self.ffn = FeedForward(dim, (dim // 3) * 8)
self.output = nn.Linear(dim, self.vocab_size)
def forward(self, x):
for attn in self.layers:
y = attn(self.ln1(x)) + x
x = self.ffn(self.ln2(y)) + y
return self.output(x)性能優化
論文還提供了兩種FlashAttention實現方式：
1、支持不同維度的實現：
def FlashDiffAttn_1(X, W_q, W_k, W_v, λ):
Q1, Q2 = split(X @ W_q)
K1, K2 = split(X @ W_k)
V = X @ W_v
A1 = flash_attn(Q1, K1, V)
A2 = flash_attn(Q2, K2, V)
return A1 - λ A2固定維度的實現：
def FlashDiffAttn_2(X, W_q, W_k, W_v, λ):
Q1, Q2 = split(X @ W_q)
K1, K2 = split(X @ W_k)
V1, V2 = split(X @ W_v)
A11 = flash_attn(Q1, K1, V1)
A12 = flash_attn(Q1, K1, V2)
A1 = Concat(A11, A12)
A21 = flash_attn(Q2, K2, V1)
A22 = flash_attn(Q2, K2, V2)
A2 = Concat(A21, A22)
return A1 - λ A2Differential Transformer論文提出的兩種FlashAttention實現方案各有特色。第一種實現（FlashDiffAttn_1）采用直接計算策略，允許Q、K、V具有不同的維度，這種靈活性使其更適合需要動態調整維度的場景，但可能在某些硬件上的優化效果不如第二種方案。第二種實現（FlashDiffAttn_2）通過將計算分解為多個相同維度的子運算，雖然計算步驟增多，但每個步驟都能充分利用硬件優化，特別是在支持張量核心的現代GPU上表現更好。
這兩種實現的選擇主要取決于具體應用場景：如果模型架構需要頻繁調整維度或者需要更靈活的注意力機制，建議使用第一種實現；如果追求極致的計算效率且維度相對固定，第二種實現可能是更好的選擇。從工程實踐角度看，第二種實現與現有的FlashAttention優化庫的兼容性更好，更容易在現有基礎設施上部署和優化。
局限性和未來研究方向
Differential Transformer雖然在多個方面展現出優秀的性能，但仍然存在一些值得關注的局限性。首要的挑戰來自其計算效率方面。由于模型需要同時計算兩個獨立的注意力圖譜，這不可避免地增加了計算開銷。在實際測試中，相比傳統Transformer，DIFF Transformer在3B規模模型上的計算吞吐量降低了約9%，這種性能損失雖然可以通過更少的參數量來部分抵消，但在大規模部署場景中仍然需要認真考慮。
內存使用是另一個重要的局限性。模型需要存儲兩組獨立的查詢和鍵值向量，這導致了更高的內存占用。盡管這種設計對于提升模型性能是必要的，但在資源受限的環境下可能會造成部署困難。特別是在處理超長序列時，內存壓力會進一步加大。
訓練穩定性也是一個需要特別關注的問題。模型中λ參數的初始化策略對訓練過程的穩定性有顯著影響。研究發現，不同的λinit取值會導致訓練收斂速度和最終性能的差異。雖然論文提出了一個基于層深度的初始化策略，但這種方案并非在所有場景下都能取得最優效果，有時需要根據具體任務進行調整。
基于這些局限性，論文提出未來的研究可以沿著幾個重要方向展開。首先在計算效率優化方面，可以探索更高效的注意力核心實現。這包括研究如何更好地利用現代硬件特性，例如開發專門的CUDA核心來加速差分注意力的計算。同時考慮到模型產生的稀疏注意力模式，可以設計特定的稀疏計算優化策略，這不僅能提升計算效率，還能減少內存占用。
λ參數的動態調整機制是另一個值得深入研究的方向。當前的參數計算方案雖然有效，但仍有優化空間。可以考慮設計更靈活的自適應機制，使λ參數能夠根據輸入內容和任務特點動態調整，從而在不同場景下都能獲得最佳性能。這可能需要引入額外的上下文感知機制，或者設計新的參數更新策略。
在內存優化方面，量化技術提供了一個有前景的研究方向。考慮到DIFF Transformer在處理激活值異常方面的優勢，可以探索專門的量化策略。比如，研究如何在保持模型性能的同時，對注意力權重和中間狀態進行更激進的量化，從而減少內存占用。這對于模型在邊緣設備上的部署具有重要意義。
長文本建模能力的進一步提升也是一個重要研究方向。雖然當前模型在64K長度的實驗中表現出色，但隨著應用需求的增長，可能需要處理更長的序列。這要求研究如何在更長序列上保持模型的效率和性能，可能需要開發新的注意力機制變體或優化策略。
總結
DIFF Transformer通過創新的差分注意力機制成功提升了模型性能，特別是在長文本理解、關鍵信息檢索和模型魯棒性等方面。雖然存在一些計算效率和內存使用的權衡，但考慮到顯著的性能提升和更少的參數需求，這是一個非常有價值的改進。這項工作為大語言模型的架構設計提供了新的思路，也為后續研究指明了幾個重要的優化方向。
