二次方高牆:注意力為何無法擴展
回到第二篇導引,我們看過每個 Transformer 的引擎:自注意力讓每個圖塊都能觀察其他所有圖塊,並透過縮放點積注意力決定該聽誰的。這種「全部對全部」的對話正是 ViT 強大的來源,但它同時也是個陷阱。如果每個 token 都必須跟其他每個 token 比較,那麼 token 數量加倍,計算量大約就會變成四倍。本節要把這句話講精確,因為弄懂這個唯一的瓶頸,就能解釋 Swin Transformer 後續的每一個設計選擇。
用白話講一下這個計數論證。若有 N 個 token,自注意力會對每一組 i、j 之間做兩兩比較。這種有序配對的數量是 N 乘以 N,也就是 N²。而且每次比較也不是免費的——每個 token 都是長度為 D(模型維度)的向量,所以單一個點積大約要花 D 次乘加運算。把配對數量乘上每對的成本,就得到下面的總量。
全域注意力對 token 數 N 是二次方;視窗化注意力(後面幾節)則是線性。
先讀左邊。全域注意力的成本是 O(N²·D):N 是 token 的總數(每個影像圖塊一個),D 是每個 token 向量的維度,而 N² 就來自我們剛剛數過的全配對比較。大 O 符號的意思就是「正比於」——我們忽略常數倍率,專注在影像變大時成本如何擴展。右邊的 O(N·M²·D) 則是 Swin 將達成的目標:M 是以 token 計算的視窗大小(例如 7,代表 7×7 的視窗)。請注意這個關鍵差異——一旦我們把 M 固定成一個小常數,M² 也就只是個常數,於是唯一還會隨影像變大而成長的就只剩下單獨那個 N。這就把一條二次方曲線變成了一條直線。
我們代入真實數字。一個典型的偵測骨幹可能會把影像處理成 56×56 的 token 網格。也就是 N = 56 × 56 = 3,136 個 token。因此兩兩比較的數量是 N² = 3,136² ≈ 980 萬——姑且說每一個注意力層、每一個注意力頭就有一千萬次點積。現在想像換成 112×112 的網格(更高解析度的影像):N 跳到 12,544,而 N² 暴增到約 1.57 億——每邊解析度只翻一倍,工作量卻多了十六倍。這就是「二次方高牆」:早在你達到偵測與分割所需的精細空間細節之前,成本就已經把硬體甩在後頭了。
視窗化注意力:在局部思考
Swin Transformer的第一個大想法簡單到幾乎令人不好意思:別讓每個 token 跟其他每個 token 對話。改成把 token 網格切成不重疊的方形視窗——Swin 每個視窗用 7×7 個圖塊——然後只在每個視窗「之內」執行一般的自注意力。左上角視窗裡的一個 token,只會關注它自己視窗內的另外 48 個 token,其他誰都不理。昂貴的全配對比較依然會發生,但只在單一視窗的 49 個 token 之間進行,絕不會一次涵蓋全部 3,136 個 token。
一張照片被切成規則的方形圖塊網格,這些圖塊再進一步被分組成更大、互不重疊的視窗區塊。
這裡有個能讓你記住的比喻。想像一座擠滿人的體育場,每個人都想分享消息。如果每一個人都得跟其他每一個人講話,喧鬧會多到根本沒辦法溝通——那就是全域注意力。視窗化注意力則說:只跟你自己那一排的人講話。每一排都嗡嗡地進行自己那段簡短的對話,所有排同時開講,而講話的總量只隨「排數」成長——也就是隨人群規模呈線性。換取便宜的代價是:目前你這一排完全不知道後面那一排在說什麼。
請留意這在理念上發生了什麼事。在第四篇導引我們學到,ViT 對資料的飢渴源自它丟掉了 CNN 內建的局部性偏好——CNN 假設鄰近的像素最重要,而純 ViT 必須靠如山的資料從頭學起。視窗化注意力刻意把那個局部性歸納偏好重新放回來:在設計上,一個 token 就只能關注它附近的鄰居。從這個意義上說,Swin 是對 CNN 智慧的部分回歸,而且是在不放棄注意力機制的前提下達成的。這份局部性不只是個提速的小手段,它同時也是一個有用的先驗,能幫 Swin 用更少的資料學習。
# Why windowed attention is LINEAR in the number of tokens.
# x is a feature map of shape (H, W, D); M is the window side (e.g. 7).
def window_partition(x, M):
# Cut the H x W grid into non-overlapping M x M windows.
windows = []
for i in range(0, H, M):
for j in range(0, W, M):
windows.append(x[i:i+M, j:j+M]) # each window: (M, M, D)
return windows # there are (H/M)*(W/M) of them
def windowed_attention(x, M):
out = []
for w in window_partition(x, M):
tokens = w.reshape(M*M, D) # M^2 tokens in this window
out.append(self_attention(tokens)) # cost ~ O((M^2)^2 * D) per window
return reassemble(out)
# Total cost = (#windows) * (cost per window)
# = (H*W / M^2) * O(M^4 * D)
# = O(H*W * M^2 * D)
# = O(N * M^2 * D) <-- linear in N because M is fixed.位移視窗:讓鄰居互通
如果視窗彼此永遠看不到對方,我們就失去了全域視野。Swin 的神來之筆——也是整個架構命名所本的那個唯一想法——就是位移視窗注意力。訣竅是:交替使用兩種層。第一層用我們剛描述的規則視窗網格。下一層則把整個視窗網格沿對角線位移半個視窗(對 7×7 的視窗來說,就是往下約 3 個圖塊、往右約 3 個圖塊)。因為網格移動了,新的視窗現在會橫跨舊視窗的邊界——於是上一層住在不同視窗的 token,突然被分到了同一組,終於能彼此對話了。
兩層視窗網格疊在一起:第二層網格位移了半個視窗,使得每個新視窗都與四個舊視窗重疊。
回到體育場。想像每一輪「只跟同排講話」結束後,我們請所有人換座位,讓排與排之間的分界線落在新的位置。原本分處兩個不同排兩端的人,現在坐到了一起,交換各自聽到的消息。每隔一輪就這麼做一次,於是一層一層地,從體育場某個角落起頭的傳言,會逐漸漣漪般擴散到整個人群——即使任何單一場對話都從未涉及超過一個小群體。位移視窗給了 Swin 一個不斷成長的「有效感受野」:在足夠多層之後,任何一個 token 都可能已經間接影響到任何另一個 token,全域脈絡就此恢復。
這裡有一個實務上的小麻煩。位移之後,影像邊緣的視窗會不完整——奇怪尺寸的局部視窗會在邊界處戳出來。在一堆大小不規則的視窗拼布上跑注意力,既慢又雜亂。Swin 用一個巧妙的循環位移(cyclic shift)來解決:它把影像像環面(torus)一樣捲起來,把從上邊與左邊掉出去的 token 滑回到下邊與右邊。現在每個視窗又都是完整、規則的 M×M 區塊,所以快速的批次運算依然行得通。代價是:捲繞會把原本不相鄰的 token(比如影像的上邊與下邊)放進同一個視窗——而這些其實不該彼此關注。
解法是一個注意力遮罩(attention mask)。在每個捲繞後的視窗內,Swin 標記出哪些 token 配對是真正的鄰居、哪些只是因捲繞而碰巧成為鄰居。對於那些假的配對,它會在 softmax 之前替它們的注意力分數加上一個很大的負數,於是 softmax 之後它們的權重幾乎為零——這些配對雖然被算了,但毫無貢獻。(回想第二篇導引:softmax 會把分數變成權重,而一個極度負的分數會變成接近零的權重。)做完注意力後,再把循環位移反轉回去,讓每個 token 回到它原本的位置。整體效果非常漂亮:到處都是完整、規則、快速的視窗,而每個 token 卻只真正關注它真實的空間鄰居。
圖塊合併:建構特徵金字塔
目前為止,每個 Swin 層都在同一個 token 網格上運作。但真實的視覺需要不只一種尺度:精細的網格擅長發現微小細節,粗糙的網格則擅長理解整個物件與場景佈局。Swin 用圖塊合併來打造這種多尺度視野。每隔一段、在各階段之間,它會把每個 2×2 的鄰近 token 區塊融合成單一一個 token。網格縮小了,每個存留下來的 token 現在概括了更大的區域,而且——這點很關鍵——模型在同一時間把每個 token 向量加寬,以維持足夠的表徵容量。
一次圖塊合併會把每個空間邊長減半,並把通道深度加倍。
我們仔細拆解這個箭頭。左邊,特徵圖有高 H、寬 W,每個 token 有 C 個通道(通道數就是每個 token 向量的長度)。圖塊合併會看每一個不重疊的 2×2 鄰域——那是 4 個 token、每個長度 C——並把它們串接成一個長度 4C 的長向量。把四個長度 C 的向量首尾相接,自然就得到 4C。接著一個線性層把那個 4C 向量投影降到 2C。右邊的結果就是一個高減半(H/2)、寬減半(W/2)、通道加倍(2C)的網格。所以 token 數量掉到四分之一(因為 H/2 × W/2 = HW/4),而每個 token 的豐富度則變成兩倍。
def patch_merging(x):
# x: a feature map of shape (H, W, C)
# Pick out the four positions of every 2x2 block.
x0 = x[0::2, 0::2, :] # top-left -> (H/2, W/2, C)
x1 = x[1::2, 0::2, :] # bottom-left -> (H/2, W/2, C)
x2 = x[0::2, 1::2, :] # top-right -> (H/2, W/2, C)
x3 = x[1::2, 1::2, :] # bottom-right -> (H/2, W/2, C)
x = concat([x0, x1, x2, x3], axis=-1) # glue the 4 neighbours -> (H/2, W/2, 4C)
x = linear(in=4*C, out=2*C)(x) # learnable squeeze 4C -> 2C
return x # (H/2, W/2, 2C): half the grid, double the depth一疊逐漸變小的特徵圖網格,每一層的通道數都比下一層多,構成一座金字塔。
我們帶一個具體的 Swin-T 走過它的四個階段,從一張 224×224 的影像切成 4×4 圖塊開始,得到一個有 C 通道的 56×56 token 網格。第一階段在 56×56×C 上做注意力。一次圖塊合併接著產生第二階段的 28×28×2C;再一次得到第三階段的 14×14×4C;最後一次得到第四階段的 7×7×8C。四張特徵圖、四種尺度——金字塔底部是精細細節,頂部是整個物件的語意。這刻意地仿照了 CNN:ResNet 裡的池化層做的正是同一套「解析度減半、通道加寬」的舞步,來建構它們經典的特徵金字塔。
Swin 作為通用骨幹
現在來看回報。在電腦視覺裡,「骨幹」是你拿來掛上各種任務頭部的特徵抽取主幹:分類頭、偵測頭、分割頭。多年來那個骨幹都是像 ResNet 這樣的 CNN,正是因為偵測與分割的框架都期待一座多尺度特徵圖金字塔,去餵給像特徵金字塔網路(FPN)這樣的結構。由於Swin Transformer輸出的正是那座金字塔——它的四個階段分別位於輸入解析度的 1/4、1/8、1/16 與 1/32——它能直接嵌入那些既有框架,作為 CNN 的即插即用替代品,完全不需要重新設計頭部。
看清這三個想法如何卡合成一台連貫的機器會很有幫助。視窗化注意力給了 Swin 局部性歸納偏好,以及對影像尺寸線性的成本。位移視窗讓資訊一層一層地滲過視窗邊界,恢復了那個讓注意力值得使用的全域脈絡。圖塊合併則把結果堆疊成一個多尺度的階層。局部性+全域流動+階層:這正是當初讓 CNN 成為優秀視覺骨幹的那三項特質——如今在一個仍能學習豐富、依內容而變的注意力的 Transformer 裡被重建了出來。
一座來自骨幹的多尺度特徵金字塔,分別接到偵測與分割的預測頭部。
這正是純 ViT 作為骨幹時的困難所在。只有單一低解析度的網格、沒有原生的金字塔,傳統 ViT 很難把偵測或分割頭部所需的整個尺度範圍交給它——你得嫁接額外的機制去假造一個出來。Swin 讓那種嫁接變得多餘,這也是它如此迅速流行起來的重要原因之一。
真實世界的判決與理論相符。Swin 問世時,在旗艦級的偵測與分割基準上(COCO 物件偵測、ADE20K 語意分割)創下了新的最先進數字,同時還保持高效,並迅速被學界與業界採用為通用的視覺骨幹。在本軌的脈絡裡,Swin 是那個終於讓視覺 Transformer 對密集、高解析度、真實世界影像變得實用的架構——而不只是用來分類整齊的 224×224 縮圖。本質上,它是一個成功的CNN 與 Transformer 思維的混合體:CNN 風格的階層與局部性,加上 Transformer 風格的注意力。
視覺 Transformer 的未來走向
退一步,看看本軌畫出的這條弧線。最初的視覺 Transformer證明了純粹的注意力、餵入影像圖塊,能與 CNN 匹敵——但前提是要在巨量的標註資料集上訓練。接著資料高效影像 Transformer(DeiT)展示了如何僅用普通的 ImageNet 就訓練出強大的 ViT,靠的是蒸餾與精心設計的訓練配方。而Swin Transformer則重新設計了架構本身,讓注意力能擴展到真實的高解析度影像,並充當通用骨幹。三步棋:證明這個想法、讓它資料高效、讓它能擴展。
不過,對 ViT 資料飢渴最根本的解答,是乾脆完全不再依賴標籤。自監督預訓練(self-supervised pretraining)從原始、無標註的影像中學習,方法是設計一個不需要任何人工標註的任務。最具代表性的例子是遮罩影像建模(masked image modeling),如 MAE(遮罩自編碼器):隨機遮住一張影像大部分的圖塊——常常高達 75%——並訓練 Transformer 從剩下的少數圖塊重建出消失的像素。為了正確補出一個被遮住的圖塊,模型被迫去理解物件、紋理與脈絡,於是它免費學到了豐富的特徵。這正是大型語言模型透過預測被遮住或下一個詞來預訓練的視覺版近親。
第二個互補的家族是自蒸餾(self-distillation),也就是 DINO 背後的想法。這裡有網路的兩個副本——一個「學生」和一個緩慢更新的「老師」——各自看到同一張影像的不同隨機裁切與增強,而學生被訓練去匹配老師的輸出。在完全沒有標籤的情況下,網路學到的特徵組織得如此良好,以致於它的注意力圖竟會自行在物件邊界上亮起來,而它的特徵也會按語意類別自然聚成群。遮罩影像建模與自蒸餾合在一起,才是通往資料高效的真正道路:先在一片無標註影像之海上預訓練一次,再在你真正在意的任何小型標註任務上微調。
再放大一層,最大的趨勢是統一。你在第三篇導引認識的那個一模一樣的 Transformer 區塊,如今同時支撐著文字、影像、音訊與影片的最先進模型。因為「影像即圖塊」不過是一個序列、「文字即 token」也不過是一個序列,單一架構就能一次吸收兩者——這正是現代多模態模型的動力來源,它們能為影像加上說明、回答關於照片的問題,或從一句話生成圖片。視覺與語言,這兩個曾經各自獨立的研究世界,正越來越用同一種架構語言交談:對 token 序列的注意力。