Cross-Attention/zh

    From Marovi AI
    This page is a translated version of the page Cross-Attention and the translation is 100% complete.
    Other languages:
    Article
    Topic area Deep Learning
    Prerequisites Attention Mechanism, Self-Attention, Transformer


    概述

    交叉注意力注意力機制的一種變體,其中查詢來自一個序列,而鍵和值來自另一個序列。它是讓模型根據另一個流的內容來調控某個流的生成或表示的標準機制,也是編碼器-解碼器Transformer、檢索增強模型以及大多數將文本與圖像、音頻或視頻綁定的現代多模態系統的架構骨幹。

    自注意力不同——後者中查詢、鍵和值都是同一輸入的投影——交叉注意力建立了一種從源序列到目標序列的非對稱信息流。目標序列提出問題;源序列提供證據。這種解耦使得交叉注意力在兩個流的長度不同、模態不同或角色不同時尤為自然契合:例如,在機器翻譯解碼器的詞元關注編碼器的輸出,或者在視覺-語言模型中文本詞元關注圖像塊。

    交叉注意力作為原始Transformer架構的一部分被提出[1],此後已成為在擴散模型[2]、Perceiver 風格架構[3]以及當代多模態大型語言模型中被反覆使用的基本構件。

    直覺

    一個有用的思維模型是把注意力看作在一個關聯記憶中進行的軟性、可微分的查找。自注意力是一種其內容就是執行查找的詞元本身的記憶;交叉注意力則是一種內容來自別處的記憶。翻譯模型的解碼器在生成下一個德語單詞時,會查詢編碼器對源英語句子的表示,以決定此刻應該關注哪些詞。查詢知道自己想要什麼("一個表示主語的名詞短語");鍵宣告它們能提供什麼("我是一個關於貓的名詞短語");值則提供實際混合到解碼器隱藏狀態中的內容。

    由於在一個解碼步內鍵和值來自一個固定的外部源,交叉注意力也是將條件信號注入生成模型的最自然的位置。例如,文本到圖像的擴散模型將去噪U-Net的空間特徵作為查詢,將編碼後的文本提示作為鍵和值,使得每個空間位置都能在每個去噪步中有選擇地從提示中提取語義內容。

    公式化

    設目標(查詢)序列長度為$ n $、隱藏維度為$ d $,即$ X_{\text{tgt}} \in \mathbb{R}^{n \times d} $;設源(鍵/值)序列長度為$ m $,即$ X_{\text{src}} \in \mathbb{R}^{m \times d} $。三個可學習的線性投影分別產生查詢、鍵和值:

    $ {\displaystyle Q = X_{\text{tgt}} W_Q, \quad K = X_{\text{src}} W_K, \quad V = X_{\text{src}} W_V} $

    其中$ W_Q, W_K \in \mathbb{R}^{d \times d_k} $$ W_V \in \mathbb{R}^{d \times d_v} $。縮放點積交叉注意力隨後計算

    $ {\displaystyle \operatorname{CrossAttn}(X_{\text{tgt}}, X_{\text{src}}) = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V \in \mathbb{R}^{n \times d_v}.} $

    關鍵的結構性事實是兩個操作數的非對稱性:輸出與目標長度相同,但其內容是從源中抽取的值向量的凸組合。注意力矩陣是形狀為$ n \times m $的矩形矩陣,而不是像自注意力中那樣是方陣。

    在實踐中,交叉注意力幾乎總是採用其多頭形式。設有$ h $個維度為$ d_k = d / h $的頭,查詢、鍵和值被分割,每個頭分別應用注意力,然後將各頭的輸出拼接並進行線性投影:

    $ {\displaystyle \operatorname{MultiHead}(X_{\text{tgt}}, X_{\text{src}}) = \operatorname{Concat}(\operatorname{head}_1, \ldots, \operatorname{head}_h)\, W_O.} $

    不同的頭可以各自專門化:有些按位置進行對齊,有些進行語義對齊,還有些則起到近似均勻平滑的作用。

    在編碼器-解碼器 Transformer 中的應用

    在原始Transformer中,每個解碼器塊包含三個子層:對部分已生成目標進行的掩碼自注意力、對編碼器棧最終輸出進行的交叉注意力,以及按位置施加的前饋網絡。交叉注意力子層是信息從源流向目標的唯一通道;將其移除後,解碼器就退化為一個看不到源句子的普通語言模型

    由此可以得出幾個實際的要點。由於編碼器表示只計算一次並在每個解碼步中復用,鍵和值可以跨步緩存,這使得交叉注意力比隨部分目標增長的自注意力便宜得多。大多數生產環境中的解碼器會為自注意力維護一個單獨的 KV 緩存(按詞元逐步增長),並為交叉注意力維護一個靜態、預先計算好的 KV 張量

    交叉注意力也是源端填充掩碼的應用位置:源序列中被填充的位置會被掩蔽,使得softmax對其分配零概率。相比之下,因果掩碼在交叉注意力中並不必要——解碼器在任何解碼步都可以關注任何源位置。

    變體

    若干變體在基本交叉注意力層之上進行擴展或修改,以應對特定的約束。

    帶門控的交叉注意力在交叉注意力的輸出上引入一個可學習的、通常初始化為零的門,使得新增的交叉注意力層不會破壞預訓練模型的穩定性。這是 Flamingo 用來將視覺上下文嫁接到一個凍結的語言模型上的機制[4],更普遍地,它也是參數高效多模態適配的一種常見模式。

    Perceiver 風格的交叉注意力使用一小組可學習的潛在向量作為查詢,針對一段非常長的輸入序列進行查詢,將輸入壓縮為一個與其長度無關的固定大小表示。這打破了標準自注意力對輸入長度的二次依賴,也正是 Perceiver 系列能夠在不依賴特定模態分詞器的前提下處理原始像素、音頻採樣和點雲的原因。

    擴散模型中的交叉注意力通過將網絡的空間特徵作為查詢、將條件嵌入作為鍵和值,來基於文本或類別嵌入對去噪網絡進行條件化。同一機制在每一層、每一步去噪中反覆應用,正是這一點賦予潛在擴散模型對生成圖像的細粒度可控性。

    記憶與檢索式交叉注意力將源序列推廣為一個被檢索到的片段數據庫。諸如 RETRO 以及 kNN 增強Transformer之類的架構會檢索最近鄰段落並對其進行交叉注意力,從而把模型的參數容量與其在推理時可訪問的知識解耦開來。

    交叉注意力與自注意力的比較

    交叉注意力自注意力之間的差異是結構性的而非算法性的:計算的是相同的縮放點積,但鍵和值來自不同的源。由此帶來了若干實際後果。

    注意力矩陣是矩形的,通常不是方陣,因此其代價為$ O(nm) $而非$ O(n^2) $;對於關注長源的短目標,這比在兩者拼接序列上做自注意力要便宜得多。填充掩碼僅作用於源端;在使用因果掩碼時,它作用於目標自身的自注意力,而非交叉注意力。由於源表示在解碼過程中是固定的,其鍵和值可以預先計算一次並復用,這在推理時是一項可觀的收益。

    一個更細微的要點是:如果源表示已經包含了來自先前編碼器的位置信息,那麼交叉注意力在源端就不需要再添加位置編碼。在多模態場景下,源模態和目標模態的位置結構差異很大(例如,源是二維圖像塊,目標是一維文本),位置信息通常駐留在編碼器內部,而不是在交叉注意力的邊界處添加。

    局限性

    交叉注意力沿用了標準注意力在$ n \times m $矩形上的二次內存開銷。當源非常長時——例如長文檔、高分辨率圖像或時長達數小時的音頻——注意力矩陣會成為主導開銷,因此需要各種稀疏、低秩或內存高效的近似方法[5]

    交叉注意力對源流和目標流之間的分布偏移也是出了名的脆弱。一個被訓練為關注乾淨編碼器輸出的解碼器,在編碼器被替換或微調後可能會顯著退化,因為鍵的幾何結構可能以查詢未曾預料的方式發生變化。聯合訓練、門控或精心設計的適配器通常可以緩解這一問題。

    最後,交叉注意力本身並不能解決落地(grounding)幻覺問題。該機制僅規定信息如何流動,並不強制目標忠實地反映源。使用交叉注意力訓練的模型完全可能——並且確實會——忽略其條件,尤其是在自回歸設定下,目標自身的自注意力可能壓過交叉注意力信號。

    參考文獻