很多人剛開始學(xué)AI以為最難啃的是微積分。可真到寫代碼的時(shí)候會(huì)發(fā)現(xiàn)最先攔住你的往往是線性代數(shù)——向量、矩陣、張量這些看起來像數(shù)學(xué)課的“抽象概念”其實(shí)是每一個(gè)模型骨架里的“語法規(guī)則”。這篇進(jìn)階篇數(shù)學(xué)篇第6講就專門把這些計(jì)算掰開揉碎向量、矩陣、張量在AI里到底怎么運(yùn)算、怎么用以及你寫模型時(shí)踩過的那些形狀錯(cuò)誤、廣播警告根源在哪里。適合已經(jīng)會(huì)用框架搭模型、但數(shù)學(xué)細(xì)節(jié)一直似懂非懂的同學(xué)也適合準(zhǔn)備面試前臨時(shí)抱佛腳的人。1. 為什么說向量、矩陣、張量是 AI 的“計(jì)算語法”1.1 從單個(gè)數(shù)字到高維容器標(biāo)量、向量、矩陣、張量一個(gè)普通的數(shù)字在數(shù)學(xué)里叫標(biāo)量在代碼里就是一個(gè)float或int。它沒有任何方向也沒有“批量”的概念。但AI模型每次推理都要處理成千上萬個(gè)樣本每個(gè)樣本又包含幾百上千個(gè)特征如果還一個(gè)數(shù)字一個(gè)數(shù)字地管理模型根本沒法寫。于是有了向量一串?dāng)?shù)字排成一行或一列通常表示一個(gè)樣本的特征。比如某個(gè)用戶有年齡、收入、活躍度三個(gè)數(shù)值就可以用一個(gè)三維向量[28, 15000, 0.8]表示。矩陣則是二維的像一個(gè)表格。它的第一個(gè)維度常常是“樣本數(shù)”第二個(gè)維度是“特征數(shù)”。一個(gè)包含100個(gè)用戶、每人3個(gè)特征的輸入就是一個(gè)形狀為(100, 3)的矩陣。張量的定義更寬泛它是任意維度的數(shù)組。三維張量可以表示視頻或圖像序列四維張量常用來表示一批圖片五維、六維也偶有出現(xiàn)。對(duì)AI工程師來說張量就是數(shù)據(jù)的通用容器它的每一維都有一個(gè)名字比如batch_size、channel、height、width。搞定了維度命名張量就不可怕了。數(shù)據(jù)類型維度常見形狀例子AI中的典型對(duì)象標(biāo)量0維()或[]損失值、學(xué)習(xí)率向量1維(3,)單個(gè)樣本的特征向量、詞向量矩陣2維(100, 3)一批樣本的特征表、單個(gè)線性層的權(quán)重張量3維及以上(16, 3, 224, 224)一批圖片、注意力分?jǐn)?shù)矩陣、序列批次1.2 “計(jì)算語法”到底指什么把數(shù)據(jù)裝進(jìn)張量只是第一步真正的關(guān)鍵是運(yùn)算規(guī)則。同樣是一個(gè)三維數(shù)組你可以把它當(dāng)成“一批2D圖片”也可以當(dāng)成“一條序列的隱藏狀態(tài)”這決定了維度怎么對(duì)位、乘法怎么執(zhí)行。我常把向量、矩陣、張量的運(yùn)算比作英語的語法單詞是數(shù)據(jù)語法是數(shù)據(jù)之間的組合規(guī)則。你背了再多單詞不知道主謂賓順序也說不出完整句子。深度學(xué)習(xí)框架里torch.matmul、torch.reshape、torch.transpose這些函數(shù)就是語法規(guī)則而forward函數(shù)本質(zhì)上就是把這些規(guī)則按一定順序串起來描述數(shù)據(jù)如何從輸入變成輸出。而且GPU 芯片內(nèi)部對(duì)矩陣乘法和張量搬運(yùn)做了大量定制優(yōu)化模型訓(xùn)練的速度瓶頸往往不在計(jì)算本身而在于你寫的張量運(yùn)算是否貼合底層硬件習(xí)慣。后面我會(huì)專門講內(nèi)存布局和廣播問題這些看著不起眼卻直接影響訓(xùn)練速度和顯存占用。2. 向量運(yùn)算點(diǎn)積、范數(shù)和相似度計(jì)算到底在算些什么2.1 點(diǎn)積AI里出場(chǎng)率最高的親密關(guān)系度量向量點(diǎn)積的算法很簡(jiǎn)單對(duì)應(yīng)元素相乘再把所有乘積加起來。寫成公式就是a · b a1*b1 a2*b2 ... an*bn。在NumPy里一行代碼搞定import numpy as np a np.array([1.0, 2.0, 3.0]) b np.array([2.0, 0.0, 1.0]) dot np.dot(a, b) # 1*2 2*0 3*1 5點(diǎn)積的幾何意義是“一個(gè)向量在另一個(gè)向量方向上的投影長(zhǎng)度”乘以“另一個(gè)向量的長(zhǎng)度”。所以它天然適合度量?jī)蓚€(gè)向量的相似度方向越一致點(diǎn)積越大方向垂直時(shí)點(diǎn)積等于0方向相反時(shí)點(diǎn)積為負(fù)。這在AI里太常用了。自注意力機(jī)制的分?jǐn)?shù)計(jì)算就是Q和K的點(diǎn)積除以維度的平方根推薦系統(tǒng)算用戶向量和物品向量的匹配度用的也是點(diǎn)積或余弦相似度文本檢索里計(jì)算查詢向量和文檔向量的相關(guān)性本質(zhì)還是點(diǎn)積?？梢哉f大部分模型的“相關(guān)性”概念都是用點(diǎn)積定義的。2.2 范數(shù)向量長(zhǎng)度是正則化和歸一化的底層邏輯范數(shù)用來衡量向量的大小。最常見的L2范數(shù)就是向量的歐幾里得長(zhǎng)度公式是sqrt(a1^2 a2^2 ... an^2)。L1范數(shù)則是絕對(duì)值之和。L2范數(shù)在訓(xùn)練里最出名的應(yīng)用是權(quán)重衰減也就是給損失函數(shù)加上權(quán)重矩陣的L2范數(shù)平方逼著參數(shù)不要長(zhǎng)得太大。L1范數(shù)則常用來做稀疏化讓一部分特征權(quán)重變成0相當(dāng)于自動(dòng)做特征選擇。歸一化也是靠范數(shù)完成的。比如把向量a除以它的L2范數(shù)得到單位向量。兩個(gè)單位向量再做點(diǎn)積結(jié)果就是余弦相似度數(shù)值穩(wěn)定在 -1 到 1 之間。很多向量召回模型在輸出層之后都會(huì)加一個(gè)L2歸一化目的就是讓點(diǎn)積結(jié)果只跟方向有關(guān)跟向量長(zhǎng)度無關(guān)。我自己做相似度模型的時(shí)候就踩過這個(gè)坑不加L2歸一化訓(xùn)練時(shí)模型會(huì)把所有向量都拉長(zhǎng)來提高點(diǎn)積分?jǐn)?shù)推理時(shí)結(jié)果完全不可控。加上歸一化之后訓(xùn)練穩(wěn)定很多。3. 矩陣運(yùn)算線性層、批量推理與參數(shù)壓縮背后的乘法邏輯3.1 矩陣乘法的形狀對(duì)位先記住這一條規(guī)則矩陣乘法C A B要求A的列數(shù)必須等于B的行數(shù)。結(jié)果C的形狀是“A的行數(shù) × B的列數(shù)”。比如A形狀(m, n)B形狀(n, p)結(jié)果就是(m, p)。這個(gè)規(guī)則是初學(xué)者最容易卡殼的地方?？纯催@張表A形狀B形狀是否合法C形狀(3, 4)(4, 2)合法(3, 2)(3, 4)(3, 4)不合法無(3, 4)(4,)合法(3,)為什么會(huì)這樣因?yàn)榫仃嚦朔ɡ顰的每一行要和B的每一列做點(diǎn)積所以要求行的長(zhǎng)度也就是A的列數(shù)和列的高度也就是B的行數(shù)一致。從另一種視角看矩陣乘法等于把B看作列向量的集合A B等價(jià)于對(duì)A的每一行分別與B的所有列向量做點(diǎn)積。這一下子就把矩陣乘法拆回第2章的向量點(diǎn)積了心理負(fù)擔(dān)小很多。3.2 線性層的本質(zhì)把“一組神經(jīng)元”同時(shí)作用到“一批樣本”上全連接層的公式是Y X W.T b。其中X是輸入形狀一般是(batch_size, input_dim)W是權(quán)重形狀是(output_dim, input_dim)W.T轉(zhuǎn)置成(input_dim, output_dim)b是偏置形狀(output_dim,)。為什么要轉(zhuǎn)置因?yàn)閃的每一行本來就是一個(gè)神經(jīng)元的權(quán)重向量。我們想讓X的每個(gè)樣本和W的每個(gè)神經(jīng)元做點(diǎn)積得到一個(gè)(batch_size, output_dim)的矩陣。如果直接X W維度會(huì)變成(batch_size, input_dim)完全不是我們要的。轉(zhuǎn)置一下XW.T就變成“每一行是樣本每一列是神經(jīng)元”的相似度矩陣輸出里第i行第j列就是第i個(gè)樣本對(duì)第j個(gè)神經(jīng)元的原始激活值。這個(gè)操作把兩層循環(huán)遍歷樣本、遍歷神經(jīng)元壓縮成一次矩陣乘法。我剛開始學(xué)的時(shí)候很疑惑為什么框架里能直接寫nn.Linear(in_features, out_features)底層不就是一個(gè)矩陣乘加偏置嗎后來才意識(shí)到矩陣乘法的價(jià)值就是把“大量重復(fù)的點(diǎn)積”集中起來并行執(zhí)行。CPU 和 GPU 底層的 BLAS 庫(kù)對(duì)矩陣乘法做了極致的優(yōu)化分塊、內(nèi)存緩存、SIMD 指令甚至多核并行。手動(dòng)寫for sample in batch: for neuron in weight:做點(diǎn)積速度連的百分之一都不到。3.3 矩陣運(yùn)算中的批量與維度擴(kuò)展模型訓(xùn)練時(shí)通常一次喂一個(gè) batch而不是一個(gè)樣本。假設(shè) batch size 是 64輸入特征數(shù)是 128輸出神經(jīng)元數(shù)是 256那么X是(64, 128)W是(256, 128)輸出Z是(64, 256)。一次把 64×256 個(gè)點(diǎn)積全部算完這才是真正的“批量計(jì)算”。有時(shí)候你還會(huì)需要把多個(gè)權(quán)重矩陣堆起來那就需要三維張量。舉個(gè)例子多頭注意力里每個(gè)頭都有自己的W_q、W_k、W_v頭數(shù)為8每個(gè)頭的輸入維度是d_k。你當(dāng)然可以 for 循環(huán)8次但更高效的做法是把8個(gè)權(quán)重矩陣疊成一個(gè)形狀(8, input_dim, d_k)的三維張量然后一次torch.matmul完成多頭的線性變換。這就是張量運(yùn)算延展矩陣乘法的經(jīng)典場(chǎng)景。4. 張量運(yùn)算廣播機(jī)制、維度變換與內(nèi)存布局4.1 張量是帶批次的矩陣圖片、文本和序列的高維真相很多初學(xué)者看到四維張量就頭暈其實(shí)只要記住“前面的維度是批量相關(guān)后面的維度是單樣本內(nèi)部結(jié)構(gòu)”就夠用了。圖像分類中一批 RGB 圖片的形狀是(batch_size, channels, height, width)比如(16, 3, 224, 224)。16是圖片張數(shù)3是RGB通道后面是每張圖片的高寬。如果你把單張圖片(3, 224, 224)當(dāng)成“一張?zhí)厥獾木仃嚒蹦?16, 3, 224, 224)就是16個(gè)這樣的矩陣疊在一起所以叫張量。自然語言處理中一句話經(jīng)過詞嵌入后得到(seq_len, hidden_size)一批句子則是(batch_size, seq_len, hidden_size)。注意力分?jǐn)?shù)矩陣是(batch_size, num_heads, seq_len_q, seq_len_k)簡(jiǎn)直是一個(gè)套娃。但只要你給每個(gè)維度起了名字就不會(huì)亂。4.2 廣播機(jī)制沒有復(fù)制數(shù)據(jù)卻像把數(shù)據(jù)擴(kuò)展了一樣兩個(gè)張量做逐元素運(yùn)算時(shí)形狀不一定要完全一樣。NumPy 和 PyTorch 遵循一套廣播規(guī)則從右往左對(duì)齊維度每個(gè)維度要么相等要么其中一個(gè)等于1要么一個(gè)維度缺失?？催@個(gè)例子import torch a torch.randn(3, 1) # 形狀 (3, 1) b torch.randn(1, 4) # 形狀 (1, 4) c a b # 結(jié)果形狀 (3, 4)這里是(3, 1)和(1, 4)相加廣播后相當(dāng)于a的第1維復(fù)制4份b的第0維復(fù)制3份最終得到(3, 4)。廣播看起來很方便但也埋了很多雷。有一次我調(diào)試一個(gè)模型發(fā)現(xiàn)某個(gè)中間結(jié)果在 batch size 大于1時(shí)出現(xiàn)奇怪的低精度排查半天才發(fā)現(xiàn)是形狀(batch, 1)的向量和一個(gè)(batch, feature)的矩陣相加廣播把向量“橫向拉伸”到了所有特征維上導(dǎo)致每個(gè)樣本都加了同樣的偏置。這不是錯(cuò)誤但完全不是我想要的。所以遇到形狀警告或者奇怪的數(shù)值第一反應(yīng)就是檢查是不是廣播在偷偷“幫倒忙”。4.3 維度變換是改“視圖”而不是改“數(shù)據(jù)”view、reshape、permute、transpose是張量維度操作的四件套但它們的底層邏輯完全不同。permute/transpose只交換維度順序不改變底層數(shù)據(jù)的排列順序。view/reshape則要求兩個(gè)視圖的元素順序一致如果之前做過transpose底層內(nèi)存順序已經(jīng)不是原來的順序直接view會(huì)報(bào)錯(cuò)或者產(chǎn)生錯(cuò)誤結(jié)果。解決辦法是先contiguous()讓張量在內(nèi)存里變得連續(xù)再view。我見過很多同學(xué)在attention實(shí)現(xiàn)里把(batch, heads, seq_len, d_head)轉(zhuǎn)成(batch, seq_len, heads * d_head)時(shí)直接view報(bào)錯(cuò)然后一臉懵。原因就是之前做過transpose。正確的做法是permute之后.contiguous().view(...)。理解內(nèi)存布局還有性能意義。GPU 對(duì)連續(xù)內(nèi)存的訪問效率遠(yuǎn)高于跳躍式訪問。如果你把NCHW的圖片張量transpose成NHWC雖然形狀看著沒問題但每次內(nèi)存讀取都不連續(xù)訓(xùn)練速度可能下降百分之二三十。所以盡量讓數(shù)據(jù)從加載到計(jì)算都保持同一套布局別在模型中間來回轉(zhuǎn)置。5. 手寫一個(gè)線性層用 NumPy 復(fù)現(xiàn)前向傳播和反向梯度5.1 前向傳播一行矩陣乘法就夠我們用一個(gè)小例子把前面所有知識(shí)串起來。假設(shè)輸入X有3個(gè)樣本每個(gè)樣本4個(gè)特征線性層有2個(gè)神經(jīng)元import numpy as np X np.array([ [1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0], ]) # 形狀 (3, 4) W np.array([ [0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8], ]) # 形狀 (2, 4)一行是一個(gè)神經(jīng)元 b np.array([0.1, 0.2]) # 形狀 (2,) Z X W.T b # (3, 2) print(Z)X W.T的形狀是(3, 2)加上b時(shí)廣播到每一行。這個(gè)Z就是兩個(gè)神經(jīng)元對(duì)三個(gè)樣本的原始輸出。如果后面接一個(gè)sigmoid它就是2個(gè)神經(jīng)元的前向激活值。很多人會(huì)疑惑為什么不把W定義成(input_dim, output_dim)直接用X W當(dāng)然也可以但 PyTorch 的nn.Linear官方實(shí)現(xiàn)里權(quán)重形狀是(out_features, in_features)所以需要轉(zhuǎn)置。我這里故意用同樣的定義方式就是為了讓你在看框架文檔時(shí)能對(duì)上號(hào)。5.2 反向傳播梯度也是矩陣乘法訓(xùn)練模型時(shí)反向傳播算梯度本質(zhì)上還是在做矩陣乘法。對(duì)于線性層Z X W.T b如果我們得到了損失對(duì)Z的梯度dZ那么權(quán)重的梯度就是dW dZ.T X。用剛才的數(shù)據(jù)試一下dZ np.array([ [1.0, -1.0], [0.5, 0.5], [-0.5, 0.2], ]) # 模擬上游梯度 dW dZ.T X # 形狀 (2, 4) print(dW)這里dZ.T形狀是(2, 3)X形狀是(3, 4)結(jié)果恰好是(2, 4)和W形狀一致。如果寫成X.T dZ也行關(guān)鍵是矩陣乘法要讓我們把“每個(gè)神經(jīng)元對(duì)每個(gè)輸入特征的梯度”都算出來。自己用 NumPy 寫一遍這個(gè)流程比看十遍原理都管用。做完之后你會(huì)特別直觀地理解為什么 autograd 能自動(dòng)求梯度它只是在計(jì)算圖的每個(gè)節(jié)點(diǎn)上應(yīng)用了類似的矩陣乘法規(guī)則而已。5.3 加上激活函數(shù)非線性藏在逐元素操作里線性層串再多矩陣乘法所有層加起來仍然是一個(gè)線性變換因?yàn)榫仃嚦朔M足結(jié)合律。真正讓深度學(xué)習(xí)變“深”的是激活函數(shù)ReLU、Sigmoid、Tanh 等都是逐元素的張量運(yùn)算。用代碼體現(xiàn)一下def relu(tensor): return np.maximum(tensor, 0) A1 relu(X W1.T b1) A2 relu(A1 W2.T b2)這里A1是(3, hidden_dim)的矩陣它同時(shí)是下一層的輸入。所以張量的形狀就是在這樣的逐層變換中流轉(zhuǎn)。每個(gè) batch 的數(shù)據(jù)從(batch, input_dim)變成(batch, hidden_dim)再變成(batch, output_dim)整個(gè)模型就是一系列矩陣乘法和逐元素激活函數(shù)的組合。6. 我踩過的張量形狀坑維度錯(cuò)位、廣播誤用和性能瓶頸6.1 形狀錯(cuò)誤“忘記轉(zhuǎn)置”是頭號(hào)坑我自己剛寫 Transformer 時(shí)卡得最久的就是注意力分?jǐn)?shù)計(jì)算。Q和K的形狀都是(batch, heads, seq_len, head_dim)按理說要對(duì)最后兩個(gè)維度做點(diǎn)積Q K.transpose(-2, -1)結(jié)果是(batch, heads, seq_len, seq_len)同時(shí)轉(zhuǎn)置的位置還關(guān)系到K要不要先contiguous。有次我把K轉(zhuǎn)置后直接view報(bào)錯(cuò)說“view size is not compatible with input size”。后來才意識(shí)到transpose會(huì)讓內(nèi)存布局變成非連續(xù)必須先.contiguous()。這個(gè)坑幾乎人人都會(huì)踩一次解決辦法就是記住transpose和permute之后如果要改變形狀或者接入某些要求連續(xù)內(nèi)存的操作先加contiguous()。6.2 廣播誤用相似度矩陣被悄悄“降維”另一個(gè)經(jīng)典問題是計(jì)算兩個(gè)矩陣兩兩之間的相似度。假設(shè)A是(m, d)B是(n, d)你想要(m, n)的相似度矩陣。最自然的寫法是兩層循環(huán)for i in range(m): for j in range(n): sim[i, j] np.dot(A[i], B[j])但這樣既冗長(zhǎng)又慢。高效寫法是sim A B.T # 形狀 (m, n)直接把所有點(diǎn)積一次算完。這個(gè)例子說明很多“看起來要循環(huán)”的場(chǎng)景其實(shí)都能用矩陣乘法表達(dá)。學(xué)會(huì)用替代顯式循環(huán)不僅代碼整潔速度優(yōu)勢(shì)尤其明顯。在一組 512 個(gè) query 和 512 個(gè) key 上對(duì)比測(cè)試循環(huán)要幾毫秒矩陣乘法可以到微秒級(jí)而且批量越大差距越明顯。6.3 精度和內(nèi)存隱藏的性能殺手我遇過一種情況模型在 PyTorch 上是float32但中間為了調(diào)試把數(shù)據(jù)轉(zhuǎn)成 NumPy某些操作默認(rèn)變成了float64再轉(zhuǎn)回 Tensor 時(shí)顯存直接翻倍。在batch size很大的場(chǎng)景下這一下就可能讓 OOM 出現(xiàn)。解決辦法是時(shí)刻關(guān)注dtype。NumPy 創(chuàng)建數(shù)組時(shí)默認(rèn)float64如果你只是做數(shù)值驗(yàn)證沒關(guān)系但如果你在模型訓(xùn)練循環(huán)里做自定義數(shù)據(jù)處理記得顯式指定dtypenp.float32或者用torch.float32。另外GPU 上減少不必要的張量復(fù)制也很重要。大多數(shù)框架的view和reshape不會(huì)復(fù)制數(shù)據(jù)但contiguous()在某些情況下會(huì)觸發(fā)新拷貝頻繁調(diào)用會(huì)拖慢整體速度。我在調(diào)優(yōu)時(shí)常用的一個(gè)原則是從一個(gè)張量到另一個(gè)張量盡量只改變每個(gè)維度的含義而不是反復(fù)交換維度順序。如果非要換就把所有轉(zhuǎn)置集中到數(shù)據(jù)預(yù)處理階段不要在模型主循環(huán)里做。最后再補(bǔ)充一個(gè)很多人忽略的點(diǎn)不要迷信einsum。它確實(shí)能讓復(fù)雜的高維矩陣乘法寫得很簡(jiǎn)潔比如注意力里的torch.einsum(bqhd,bkhd-bhqk, Q, K)。但einsum的缺點(diǎn)是報(bào)錯(cuò)信息非常隱晦內(nèi)部還可能走通用路徑性能未必優(yōu)于顯式的加transpose。我的習(xí)慣是先用顯式矩陣乘法調(diào)通功能確認(rèn)需求后再考慮用einsum簡(jiǎn)化代碼。這樣踩的坑最少效率也最高。機(jī)器學(xué)習(xí)這條路走到越深越會(huì)發(fā)現(xiàn)數(shù)學(xué)不是“紙上公式”而是嵌在每一行代碼里的思維框架。向量、矩陣、張量的運(yùn)算規(guī)則就是AI的底層語法搞懂了它們你看模型源碼會(huì)順暢很多排查 bug 也不用靠瞎猜。希望這篇“計(jì)算語法”筆記能幫你把線性代數(shù)和實(shí)際模型的最后一段縫隙也補(bǔ)上。