Transformer模型已經(jīng)成為大語言模型（LLMs）的標準架構(gòu)，但研究表明這些模型在準確檢索關(guān)鍵信息方面仍面臨挑戰(zhàn)。今天介紹一篇名叫Differential Transformer的論文，論文的作者觀察到一個關(guān)鍵問題：傳統(tǒng)Transformer模型傾向于過分關(guān)注不相關(guān)的上下文信息，這種"注意力噪聲"會影響模型的性能。

在這篇論文中，作者注意到transformer模型傾向于關(guān)注不相關(guān)的上下文。為了放大相關(guān)上下文的注意力分數(shù)，他們提出了一個新的注意力模型，稱為差分注意力模型。在這個模型中，他們將查詢和鍵值向量分成兩組，并計算兩個子注意力分數(shù)。

差分注意力機制

差分注意力機制（Differential Attention）的核心思想是通過計算兩個獨立的注意力圖譜之差來消除注意力噪聲。這種設(shè)計借鑒了電氣工程中差分放大器的原理，通過對比兩個信號的差異來消除共模噪聲。

讓我們看看論文中的第一個方程：

方程(1)

方程(1)顯示，我們首先像標準注意力計算一樣計算Q、K和V張量。關(guān)鍵點是我們將Q和K張量分成Q1、Q2和K1、K2子張量。

論文中輸入X、Q1、Q2、K1、K2和V張量的形狀

根據(jù)論文，Q和K張量的形狀應(yīng)該是Nx2d，因為Q1、Q2、K1和K2將是Nxd。輸入X的形狀是Nxd_model，這是論文中的嵌入維度。這就是為什么W_Q、W_K和W_V的可學習參數(shù)的形狀必須是d_modelx2d。

論文中用于lambda計算的方程(2)

方程(2)展示了如何計算可學習參數(shù)lambda。在這個方程中有一個初始lambda參數(shù)。lambda是一個標量參數(shù)，但lambda_q1、lambda_k1、lambda_q2和lambda_k2是向量。這一點很關(guān)鍵。向量lambda_q和lambda_k的運算是點積。

用于lambda初始化的方程(3)

實驗結(jié)果與性能提升

論文的實驗表明，相比傳統(tǒng)Transformer：

DIFF Transformer只需要約65%的模型參數(shù)量即可達到相同的性能，在訓練token數(shù)量方面也只需要約65%就能達到相同效果

在Needle-In-A-Haystack測試中：4K上下文長度：DIFF Transformer在多目標檢索任務(wù)中保持85%準確率；64K上下文長度：在深度為25%的位置檢測時，比傳統(tǒng)Transformer提升了76%的準確率

Python實現(xiàn)

下面我們根據(jù)論文的公式來做一個簡單的實現(xiàn)，首先方程(3)展示了我們?nèi)绾斡嬎鉲ambda_initial變量。現(xiàn)在讓我們把方程轉(zhuǎn)換成Python代碼：

def lambda_init_fn(depth):  
     return 0.8 - 0.6 * math.exp(-0.3 * depth)

然后再寫一個簡單的Python函數(shù)，使用方程(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類中，我們實現(xiàn)了一個多頭差分注意力機制。有dim_model（嵌入維度）、head_nums和depth參數(shù)。為Q1、Q2、K1和K2聲明了四個lambda可學習參數(shù)，并使用均值為0、標準差為0.1的隨機正態(tài)分布初始化它們。

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)) @ V

forward方法很直觀。我分別實現(xiàn)了方程(1)和方程(2)。forward方法直接實現(xiàn)了論文中的偽代碼。

多頭差分注意力架構(gòu)和偽代碼

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類是根據(jù)論文中的偽代碼編寫的。這里使用了RMSNorm而不是GroupNorm。

論文中使用多頭差分注意力機制的語言模型的方程

最后使用實現(xiàn)的MultiHeadDifferentialAttention機制構(gòu)建一個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)

性能優(yōu)化

論文還提供了兩種FlashAttention實現(xiàn)方式：

1、支持不同維度的實現(xiàn)：

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

固定維度的實現(xiàn)：

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 - λ A2

Differential Transformer論文提出的兩種FlashAttention實現(xiàn)方案各有特色。第一種實現(xiàn)（FlashDiffAttn_1）采用直接計算策略，允許Q、K、V具有不同的維度，這種靈活性使其更適合需要動態(tài)調(diào)整維度的場景，但可能在某些硬件上的優(yōu)化效果不如第二種方案。第二種實現(xiàn)（FlashDiffAttn_2）通過將計算分解為多個相同維度的子運算，雖然計算步驟增多，但每個步驟都能充分利用硬件優(yōu)化，特別是在支持張量核心的現(xiàn)代GPU上表現(xiàn)更好。

這兩種實現(xiàn)的選擇主要取決于具體應(yīng)用場景：如果模型架構(gòu)需要頻繁調(diào)整維度或者需要更靈活的注意力機制，建議使用第一種實現(xiàn)；如果追求極致的計算效率且維度相對固定，第二種實現(xiàn)可能是更好的選擇。從工程實踐角度看，第二種實現(xiàn)與現(xiàn)有的FlashAttention優(yōu)化庫的兼容性更好，更容易在現(xiàn)有基礎(chǔ)設(shè)施上部署和優(yōu)化。

局限性和未來研究方向

Differential Transformer雖然在多個方面展現(xiàn)出優(yōu)秀的性能，但仍然存在一些值得關(guān)注的局限性。首要的挑戰(zhàn)來自其計算效率方面。由于模型需要同時計算兩個獨立的注意力圖譜，這不可避免地增加了計算開銷。在實際測試中，相比傳統(tǒng)Transformer，DIFF Transformer在3B規(guī)模模型上的計算吞吐量降低了約9%，這種性能損失雖然可以通過更少的參數(shù)量來部分抵消，但在大規(guī)模部署場景中仍然需要認真考慮。

內(nèi)存使用是另一個重要的局限性。模型需要存儲兩組獨立的查詢和鍵值向量，這導(dǎo)致了更高的內(nèi)存占用。盡管這種設(shè)計對于提升模型性能是必要的，但在資源受限的環(huán)境下可能會造成部署困難。特別是在處理超長序列時，內(nèi)存壓力會進一步加大。

訓練穩(wěn)定性也是一個需要特別關(guān)注的問題。模型中λ參數(shù)的初始化策略對訓練過程的穩(wěn)定性有顯著影響。研究發(fā)現(xiàn)，不同的λinit取值會導(dǎo)致訓練收斂速度和最終性能的差異。雖然論文提出了一個基于層深度的初始化策略，但這種方案并非在所有場景下都能取得最優(yōu)效果，有時需要根據(jù)具體任務(wù)進行調(diào)整。

基于這些局限性，論文提出未來的研究可以沿著幾個重要方向展開。首先在計算效率優(yōu)化方面，可以探索更高效的注意力核心實現(xiàn)。這包括研究如何更好地利用現(xiàn)代硬件特性，例如開發(fā)專門的CUDA核心來加速差分注意力的計算。同時考慮到模型產(chǎn)生的稀疏注意力模式，可以設(shè)計特定的稀疏計算優(yōu)化策略，這不僅能提升計算效率，還能減少內(nèi)存占用。

λ參數(shù)的動態(tài)調(diào)整機制是另一個值得深入研究的方向。當前的參數(shù)計算方案雖然有效，但仍有優(yōu)化空間。可以考慮設(shè)計更靈活的自適應(yīng)機制，使λ參數(shù)能夠根據(jù)輸入內(nèi)容和任務(wù)特點動態(tài)調(diào)整，從而在不同場景下都能獲得最佳性能。這可能需要引入額外的上下文感知機制，或者設(shè)計新的參數(shù)更新策略。

在內(nèi)存優(yōu)化方面，量化技術(shù)提供了一個有前景的研究方向。考慮到DIFF Transformer在處理激活值異常方面的優(yōu)勢，可以探索專門的量化策略。比如，研究如何在保持模型性能的同時，對注意力權(quán)重和中間狀態(tài)進行更激進的量化，從而減少內(nèi)存占用。這對于模型在邊緣設(shè)備上的部署具有重要意義。

長文本建模能力的進一步提升也是一個重要研究方向。雖然當前模型在64K長度的實驗中表現(xiàn)出色，但隨著應(yīng)用需求的增長，可能需要處理更長的序列。這要求研究如何在更長序列上保持模型的效率和性能，可能需要開發(fā)新的注意力機制變體或優(yōu)化策略。

總結(jié)

DIFF Transformer通過創(chuàng)新的差分注意力機制成功提升了模型性能，特別是在長文本理解、關(guān)鍵信息檢索和模型魯棒性等方面。雖然存在一些計算效率和內(nèi)存使用的權(quán)衡，但考慮到顯著的性能提升和更少的參數(shù)需求，這是一個非常有價值的改進。這項工作為大語言模型的架構(gòu)設(shè)計提供了新的思路，也為后續(xù)研究指明了幾個重要的優(yōu)化方向。