「注意力實(shí)際上是對數(shù)的」？今天，一篇博客再次掀起了AI社區(qū)對注意力機(jī)制的討論。

作者認(rèn)為，Transformers 中實(shí)現(xiàn)的注意力機(jī)制，在計(jì)算復(fù)雜度上應(yīng)該被視為對數(shù)級別的。

這篇博客，還得到了 Karpathy 的高度肯定：

有時我會在想象中的神經(jīng)網(wǎng)絡(luò)完整計(jì)算圖中將其描述為「廣度是免費(fèi)的，深度是昂貴的」。

據(jù)我所知，這首先是 Transformer 背后的主要見解 / 靈感。我第一次真正受到它的震撼是在很久以前我讀到 Neural GPU 論文的時候（https://arxiv.org/abs/1511.08228）。

另外，在「從比特到智能」中為什么還要包含 python？刪除 python，我認(rèn)為你可以將其減少約 10 倍，就像 llmc 一樣。

我們知道，標(biāo)準(zhǔn)的注意力機(jī)制（如 Transformer 中的自注意力）計(jì)算步驟如下：

其復(fù)雜度主要來源于：

• 點(diǎn)積計(jì)算：QK^? 的矩陣乘法，復(fù)雜度為 O (n^2d)，其中 n 是序列長度，d 是特征維度。
• Softmax 歸一化：對每個位置的注意力權(quán)重進(jìn)行歸一化，復(fù)雜度為 O (n^2)。

一般來說，研究者認(rèn)為總復(fù)雜度隨著序列長度 n 呈平方增長，這也是標(biāo)準(zhǔn) Transformer 難以處理長序列的核心瓶頸。

而這篇博客，卻提出了另外一個全新的視角。

關(guān)于如何理解這一觀點(diǎn)，我們看看博客內(nèi)容便知。

• 博客鏈接：https://supaiku.com/attention-is-logarithmic

以下是博客內(nèi)容：

時間復(fù)雜度是衡量算法快慢最常用的標(biāo)準(zhǔn)。在 20 世紀(jì) 80 年代，那時候計(jì)算機(jī)大多只有一個核心，大家還不知道什么是單指令多數(shù)據(jù)（SIMD）技術(shù)，所以用時間復(fù)雜度來評估算法基本是合理的。

但現(xiàn)在是 2025 年，單核計(jì)算機(jī)已經(jīng)很少見了，就連智能手機(jī)都有 4 到 8 個核心。在這種情況下，只用時間復(fù)雜度來衡量算法的快慢就不夠全面了。

舉個例子來說，一個時間復(fù)雜度為 O (n3) 但能夠并行的算法，和一個必須按順序執(zhí)行的算法，單從時間復(fù)雜度上看不出來它們的區(qū)別。而且，有些算法天生就是并行的，比如線性代數(shù)，但人們還在用時間復(fù)雜度來描述它們，這其實(shí)是很荒謬的。

我們需要一種更好的方式來衡量算法的復(fù)雜度。「work-depth 模型」分析提供了一個很好的思路。它不僅關(guān)注輸入大小對應(yīng)的操作數(shù)量，還能從理論下限的角度思考算法的復(fù)雜度。

我們不僅要考慮算法執(zhí)行的原始操作數(shù)量（即「work」），更要關(guān)注計(jì)算圖相對于輸入大小的「depth」，也就是不可并行的順序操作的最小數(shù)量。因?yàn)檫@些順序操作是不可避免的，無論你的計(jì)算機(jī)有多少個核心，它們都會造成阻塞。

我主要研究機(jī)器學(xué)習(xí)系統(tǒng)的性能工程，所以接下來我會重點(diǎn)討論適用于張量的算法。「work-depth 模型」雖然不完美，但很有用。

在此，我先拋出一個問題：逐個元素相乘的時間復(fù)雜度是多少？從這個問題出發(fā)，我會進(jìn)一步闡述我的觀點(diǎn)：Transformers 中實(shí)現(xiàn)的注意力機(jī)制，在計(jì)算復(fù)雜度上應(yīng)該被視為對數(shù)級別的。

案例 1：逐個元素相乘

給定兩個長度相同的向量 a 和 b，逐個元素相乘是將 a 中的每個元素與 b 中對應(yīng)索引位置的元素相乘，并將結(jié)果存儲在新向量 c 中（或者直接在原位置修改）。

代碼如下：

從時間復(fù)雜度的角度看，這好像是線性的。如果用單線程來跑，那確實(shí)就是線性的。

然而，如果仔細(xì)觀察，你會發(fā)現(xiàn)在這個問題的計(jì)算圖中，range (n) 中的各個步驟之間沒有依賴關(guān)系。它們完全獨(dú)立。那么為什么不并行執(zhí)行它們呢？

這正是每個線性代數(shù) / 張量庫在底層所做的事情。

你很快會發(fā)現(xiàn)，逐個元素相乘實(shí)際上根本不是線性時間的！它實(shí)際上看起來像是常數(shù)時間，直到達(dá)到一個神秘的臨界點(diǎn)。

具體來說，我們可以分析逐個元素相乘時的「work」和「depth」：

算法里的每一步操作，比如加載數(shù)據(jù)、做乘法、存儲，這些操作本身都不復(fù)雜，理論上只需要常數(shù)時間就能完成。只要你的計(jì)算機(jī)有足夠的并行計(jì)算能力，直到某個臨界點(diǎn)，這些操作的時間復(fù)雜度都是常數(shù)時間。

案例 2：向量求和

向量求和比相乘更復(fù)雜一些。在這里，我們可以清楚地看到兩個步驟之間存在依賴關(guān)系（因?yàn)槔奂有枰{(diào)用 c 的狀態(tài)）。這無法完全并行執(zhí)行。

不過，向量求和看起來好像每一步都得依賴前一步，但仔細(xì)想想，不難發(fā)現(xiàn)它只是每兩個步驟（或者說每對元素）之間有點(diǎn)關(guān)聯(lián)。

實(shí)際上，這個操作仍然可以并行化，方法是不在一個步驟中并行執(zhí)行每個操作，而是在一個步驟中對每隊(duì)執(zhí)行操作。

舉個例子，假設(shè)你有一個長度為 n 的列表，向量加法是這樣的：

1. 先把列表里每一對相鄰的數(shù)字（比如第 1 個和第 2 個、第 3 個和第 4 個……）加起來。因?yàn)橐还灿?n 個數(shù)字，所以會有 n/2 對。把每對的結(jié)果存到其中一個位置（比如偶數(shù)位置或者奇數(shù)位置）。

2. 再把上一步得到的每一對結(jié)果（現(xiàn)在每對是之前兩對的和）再加起來。這次會有 n/4 對。

3. 每次都是把上一步的結(jié)果兩兩相加，直到最后只剩下一個數(shù)字。這個數(shù)字就是整個列表所有數(shù)字的總和。

這樣一來，每次操作的步驟數(shù)量都會減半。比如，第一次是 n/2 對，第二次是 n/4 對，以此類推，總共只需要 log?(n) 步就能把所有數(shù)字加起來。

案例 3：張量積

張量積是一個基本操作。它獲取兩個張量的所有索引，并對所有請求的索引（其中一些可能是共享的）逐個相乘。

比如，求兩個矩陣的張量積并且共享一個軸的時候，結(jié)果會是一個三維的張量。不過，這個操作其實(shí)并不復(fù)雜，因?yàn)樗恍枰霾⑿械募虞d、存儲、逐個相乘，所以它的「depth」是固定的，不會隨著數(shù)據(jù)量變大而增加。

但要注意，這種情況只有在張量（或者張量的一部分）能夠完整地裝進(jìn)緩存的時候才成立。如果張量太大，裝不下緩存，那就會出現(xiàn)瓶頸，因?yàn)榫彺娌粔蛴玫臅r候，計(jì)算機(jī)就不得不按順序處理數(shù)據(jù)，這時候「depth」就會增加。

張量積在機(jī)器學(xué)習(xí)里其實(shí)不太常被提到，但置換、求和、矩陣乘法、哈達(dá)瑪積、直積、各種批處理操作等等，所有這些操作都可以看成是某種形式的張量積，再加上某種形式的歸約（把多余的維度去掉或者合并）。

這樣一來，能讓復(fù)雜的張量操作變得更加系統(tǒng)、更有數(shù)學(xué)美感，尤其是在高性能計(jì)算和分布式系統(tǒng)里，用起來特別方便。

案例 4：矩陣乘法

矩陣乘法（MATMUL）就是這樣一種張量運(yùn)算，它通過張量積的收縮得到了優(yōu)雅的描述。

給定兩個張量分別為（i j）和（j k）的張量 A、B，張量乘法構(gòu)造出一個張量 C，其元素 C [i,j,k] = A [i,j] * B [j,k]，然后沿 j 維相加（收縮）成一個形狀為（i k）的矩陣 D。(為了提高效率，C 通常不會完全實(shí)體化，而是在張量積的碎片之間進(jìn)行收縮融合）。

只需忽略外軸，就可以對矩陣進(jìn)行批處理 / 廣播。

底層內(nèi)容的偽代碼：

注意，這只是將 TENSOR 順序組合成 CONTRACT，其深度復(fù)雜度分別為 O (1) 和 O (logn)：

案例 5：softmax

softmax 一點(diǎn)也不特別。先按元素應(yīng)用 e^x，然后收縮，最后按元素除法。

下面照例進(jìn)行深度復(fù)雜性分析：

案例 6：注意力

注意力就不用多說了。以下是深度分析：

可以看到，通過整數(shù)個 matmuls 收縮和一系列元素單義操作的順序組合，注意力的漸近深度復(fù)雜度僅為 O（logn + logd），其中 n 和 d 分別為序列長度和嵌入維數(shù)。

實(shí)際上，這通常意味著 O（log sequence_length），因?yàn)?sequence_length 通常遠(yuǎn)大于 embedding_dim。

局限性

然而，深度分析并不完美，當(dāng)考慮到內(nèi)存訪問模式和高速緩存的友好性時，問題立即顯現(xiàn)出來。

特別是，當(dāng)出現(xiàn)以下情況時，該模型就會失效：

• 樹的最大寬度 >> 計(jì)算單元（不管是什么內(nèi)核）。
• 內(nèi)存訪問模式不連續(xù) / 不可矢量化？
• 物化變量與內(nèi)存層次結(jié)構(gòu)不匹配。

在實(shí)踐中，這主要意味著物化張量的大小必須保持在 L2- 左右的緩存范圍內(nèi)，深度復(fù)雜度邊界才能成立。

那么為什么注意力不是對數(shù)的呢？

事實(shí)上，由于注意力至少需要將 QK^T 部分實(shí)體化（通常是非常大的整數(shù)，非常大的整數(shù)），這幾乎肯定會溢出二級緩存（這要么迫使你在內(nèi)存中計(jì)算的速度慢于 OOM，要么迫使你通過將 QK^T 矩陣分片為部分關(guān)聯(lián)塊并傳入 softmax 來將其轉(zhuǎn)化為順序問題）。

這就意味著，對于普通計(jì)算機(jī)而言，注意力的深度復(fù)雜度更像是 O (n log n)。雖然這絕不是一個不可還原的問題，但我在下一節(jié)中會提出一些推測性的解決方案。

對未來計(jì)算的猜測？

那么，這對目前的芯片和未來的芯片意味著什么？

我認(rèn)為這意味著很多，前提是一個關(guān)鍵事實(shí)，即訓(xùn)練范式在很大程度上仍然是非并發(fā)的（即看起來像循環(huán)上的前向→后向傳遞，或 dualpipe 之類的混合），為什么？

因?yàn)槿绻沁@種情況，那么神經(jīng)網(wǎng)絡(luò)的權(quán)重（在 nn 次循環(huán)中占運(yùn)動操作量的大部分）在很大程度上就是靜態(tài)的，而且計(jì)算單元的局部性會越來越強(qiáng)。

我們已經(jīng)看到這種情況的發(fā)生。權(quán)重曾經(jīng)被卸載到磁盤或保存到內(nèi)存中，只有在專門的內(nèi)核中才會啟動到 GPU。

后來，每個人都開始完全使用設(shè)備內(nèi)存（VRAM 或 HBM）進(jìn)行訓(xùn)練。

現(xiàn)在，芯片制造商已經(jīng)意識到，通過將權(quán)重轉(zhuǎn)移到更快的內(nèi)存（如 L2）上，他們可以獲得另一個 OOM（在深度復(fù)雜性分析失敗的地方有效地砍掉整個部分）。