戰(zhàn)：torch.func.jvp原理、性能陷阱與工程取舍)
大多數(shù)用過 PyTorch 的人對loss.backward()這套反向傳播都很熟但一旦被問到雅可比向量積怎么算很多人第一反應(yīng)是打開torch.autograd.functional.jacobian把整個(gè)雅可比矩陣求出來再手動做一次矩陣乘法。這種做法在小模型上沒毛病一旦輸出維度上千顯存和耗時(shí)立刻教做人。雅可比向量積Jacobian-Vector Product簡稱 JVP存在的意義就是你根本不需要顯式構(gòu)造那個(gè) m×n 的雅可比矩陣也能拿到它和一個(gè)向量的乘積代價(jià)只有一次前向傳播的量級。這篇內(nèi)容圍繞 JVP 在 PyTorch 里的幾種實(shí)現(xiàn)路徑展開重點(diǎn)拆torch.func.jvp的用法、torch.autograd.functional.jvp的性能陷阱、以及踩過的坑。適合已經(jīng)會寫autograd.grad、想進(jìn)一步搞清楚前向模式自動微分、并且需要在二階優(yōu)化、雅可比正則、靈敏度分析這些場景里真正用起來的人。1. 把 JVP 從矩陣?yán)锊鸪鰜碚l乘在左邊誰乘在右邊1.1 一個(gè) 3×2 的例子先把概念釘死先不談自動微分回到線性代數(shù)本身。假設(shè)有一個(gè)函數(shù) f: R2 → R3在點(diǎn) x 處它的雅可比矩陣 J 是一個(gè) 3 行 2 列的矩陣。JVP 就是拿一個(gè) 2 維向量 v從右邊乘上去得到 J·v結(jié)果是一個(gè) 3 維向量。注意 v 的維度必須和輸入維度一致結(jié)果的維度必須和輸出維度一致。反過來VJPVector-Jacobian Product是拿一個(gè) 3 維向量 u 從左邊乘u?·J結(jié)果是一個(gè) 2 維向量維度跟著輸入走。這兩個(gè)操作的形狀規(guī)則完全不同這是后面所有坑的根源。拿一個(gè)具體例子import torch def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x torch.tensor([2.0, 3.0], dtypetorch.float64)手工求偏導(dǎo)J 應(yīng)該是J [[x1, x0 ], [0, 2*x1 ], [cos(x0), 0 ]]代入 x [2, 3]得到[[3, 2], [0, 6], [cos(2), 0]]其中 cos(2) ≈ -0.4161?，F(xiàn)在取 v [1, 0]J·v 就是 J 的第一列即[3, 0, -0.4161]。取 v [0, 1]得到第二列[2, 6, 0]。這兩個(gè)結(jié)果后面用來驗(yàn)證代碼是否正確。1.2 為什么 PyTorch 默認(rèn)只給你 VJPloss.backward()背后跑的是反向模式自動微分它天然產(chǎn)出的是 VJP不是 JVP。原因很實(shí)際深度學(xué)習(xí)里絕大多數(shù)函數(shù)是從一個(gè)巨大的參數(shù)向量映射到一個(gè)標(biāo)量損失也就是 f: R^N → R1。這種情況下輸出維度 m 1只需要一次反向傳播就能拿到完整的梯度向量代價(jià)極低。而如果換成 JVP 的思路你得做 N 次前向傳播才能拿全梯度完全不可行。所以框架的設(shè)計(jì)重心一直壓在 VJP 上前向模式自動微分在很長一段時(shí)間里屬于有但不太好用的狀態(tài)。這直接導(dǎo)致一個(gè)現(xiàn)象當(dāng)你遇到 f: R^n → R^m 且 n 很小、m 很大的場景時(shí)比如一個(gè)幾十維的控制參數(shù)映射到上千維的仿真輸出用反向模式硬算會非常虧而這恰恰是 JVP 的主場。用一個(gè)表把兩種模式的成本關(guān)系擺清楚這張表值得貼在顯示器邊上目標(biāo)操作需要調(diào)用幾次適合的形狀PyTorch 對應(yīng)接口完整雅可比逐列前向n 次前向n 小 m 大torch.func.jacfwd完整雅可比逐行反向m 次反向m 小 n 大torch.func.jacrev單次 JVPJ·v1 次前向量級輸入小torch.func.jvp單次 VJPu?·J1 次反向量級輸出小torch.func.vjp/autograd.grad1.3 前向模式到底在算什么dual number 的直覺前向模式自動微分的實(shí)現(xiàn)方式叫對偶數(shù)dual number。把每個(gè)標(biāo)量 a 擴(kuò)展成 a b·ε 的形式規(guī)定 ε ≠ 0 但 ε2 0。任何函數(shù)作用上去之后實(shí)部就是函數(shù)值ε 的系數(shù)就是方向?qū)?shù)。舉個(gè)手算的例子x 2 1·ε那么 x2 4 4ε ε2 4 4εε 的系數(shù) 4 正好是 2x 在 x2 處的導(dǎo)數(shù)乘上方向 1。這個(gè)機(jī)制的好處是一次前向傳播就能同時(shí)算出函數(shù)值和沿著某個(gè)方向 v 的方向?qū)?shù)而且天然支持任意深的網(wǎng)絡(luò)不需要保存中間激活。代價(jià)是每個(gè)標(biāo)量都要額外存一份 ε 系數(shù)顯存占用大約翻倍但和前向傳播同一量級。PyTorch 里這套機(jī)制的底層入口是torch.autograd.forward_adtorch.func.jvp是它的高層封裝。理解這一層很關(guān)鍵因?yàn)楹竺鏁吹絫orch.autograd.functional.jvp走的根本不是這條路。2. PyTorch 里算出 JVP 的四條路各自的真實(shí)代價(jià)2.1 torch.func.jvp目前最該用的那條PyTorch 2.0 之后torch.func命名空間下提供了一套函數(shù)式變換 APIjvp是其中之一。它的簽名是jvp(func, primals, tangents)primals是函數(shù)位置參數(shù)的元組tangents是切線方向的元組兩者結(jié)構(gòu)必須一致。返回值是一個(gè)二元組(output, jvp_out)output是函數(shù)原本的輸出jvp_out是 J·v 的結(jié)果。import torch from torch.func import jvp def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x torch.tensor([2.0, 3.0], dtypetorch.float64) v torch.tensor([1.0, 0.0], dtypetorch.float64) out, jvp_out jvp(f, (x,), (v,)) print(out) # tensor([ 6.0000, 9.0000, 0.9093], dtypetorch.float64) print(jvp_out) # tensor([ 3.0000, 0.0000, -0.4161], dtypetorch.float64)跑出來和手算的[3, 0, -0.4161]完全對上。注意兩個(gè)細(xì)節(jié)第一x本身不需要requires_gradTrue前向模式不依賴計(jì)算圖第二f內(nèi)部不允許有原地操作in-place比如x[0] 1這種寫法會直接報(bào)錯(cuò)因?yàn)?dual tensor 的 ε 分量不允許被就地覆蓋。2.2 torch.autograd.functional.jvp名字一樣實(shí)現(xiàn)完全不同很多人第一次搜到的是torch.autograd.functional.jvp。它確實(shí)能用簽名也幾乎一樣from torch.autograd.functional import jvp out, jvp_out jvp(f, (x,), (v,), create_graphFalse, strictFalse)但這里有個(gè)非常容易被忽略的事實(shí)這個(gè)函數(shù)內(nèi)部是用反向模式實(shí)現(xiàn)的。它的做法是先構(gòu)造出 VJP 的形式再對切向量 v 求一次導(dǎo)本質(zhì)上是 reverse-over-reverse也就是兩次反向傳播拼出來的 JVP。這意味著什么意味著當(dāng)你的函數(shù)是輸入維度小、輸出維度大時(shí)torch.autograd.functional.jvp反而會比真正的torch.func.jvp慢一大截因?yàn)樗刂敵鼍S度反向走一遍。我在一個(gè) 50 維輸入、2000 維輸出的仿真函數(shù)上做過對比torch.func.jvp耗時(shí)大約是functional.jvp的三分之一差距相當(dāng)明顯。提示torch.autograd.functional.jvp的create_graphTrue是它為數(shù)不多的優(yōu)勢場景——需要 JVP 結(jié)果本身再參與求導(dǎo)比如構(gòu)造二階以上的混合導(dǎo)數(shù)時(shí)它比函數(shù)式 API 更省事。日常只要一階 JVP優(yōu)先torch.func.jvp。2.3 手搓 double-backward 與有限差分第三條路是老派做法算兩次autograd.grad。這不是真正的 JVP而是構(gòu)造 VJP 再套一層def manual_jvp(f, x, v): x x.detach().requires_grad_(True) y f(x) # 先拿一個(gè)啞元的 VJP把 v 塞進(jìn) grad_outputs g torch.autograd.grad(y, x, grad_outputstorch.zeros_like(y), create_graphTrue) ...說實(shí)話這條路現(xiàn)在沒什么推薦價(jià)值了代碼繞、易錯(cuò)、還容易在create_graph上翻車除非你在維護(hù)一份必須兼容老版本 PyTorch 的代碼庫否則直接跳過。第四條路是有限差分(f(x h*v) - f(x)) / h。它的缺點(diǎn)寫在教科書里——截?cái)嗾`差和舍入誤差互相拉扯h 太小結(jié)果被噪聲吃掉h 太大非線性誤差明顯。不過它有個(gè)獨(dú)特價(jià)值作為數(shù)值校驗(yàn)的金標(biāo)準(zhǔn)。你寫完任何 JVP 實(shí)現(xiàn)拿有限差分對一遍能立刻發(fā)現(xiàn)是不是方向搞反了、是不是 v 的維度用錯(cuò)了。我習(xí)慣用h 1e-6配 float64 做對照誤差在 1e-6 量級就算通過。2.4 四條路的橫向?qū)Ρ嚷窂绞欠裾媲跋蚰Ｊ较鄬臅r(shí)需要 requires_grad支持二階推薦場景torch.func.jvp是1x基線否配合create_graph式組合首選尤其 n 小時(shí)的批量 JVPtorch.autograd.functional.jvp否反向兩次2~4x否strictFalse原生支持需要高階混合導(dǎo)手搓autograd.grad否3x 以上是麻煩老版本兼容有限差分近似2x否不支持?jǐn)?shù)值校驗(yàn)3. torch.func.jvp 的實(shí)戰(zhàn)用法從單方向到整塊雅可比3.1 參數(shù)結(jié)構(gòu)不匹配是最常見的報(bào)錯(cuò)來源jvp(func, primals, tangents)這三者的結(jié)構(gòu)約束很嚴(yán)格。primals和tangents必須是元組長度等于func的位置參數(shù)個(gè)數(shù)每個(gè)位置的張量形狀必須逐一對上。如果你的函數(shù)是def f(a, b)那就得寫jvp(f, (a, b), (va, vb))漏一個(gè)或者包錯(cuò)層級都會報(bào)形狀錯(cuò)誤。def g(a, b): return (a * b).sum() b.pow(2).sum() a torch.randn(3, dtypetorch.float64) b torch.randn(4, dtypetorch.float64) va torch.ones(3, dtypetorch.float64) vb torch.zeros(4, dtypetorch.float64) out, jvp_out jvp(g, (a, b), (va, vb)) # 輸出是標(biāo)量對 a 方向的偏導(dǎo)之和這里vb傳全零意思是只要 a 方向的導(dǎo)數(shù)。這是一個(gè)非常實(shí)用的技巧當(dāng)你只關(guān)心某一個(gè)輸入的分量時(shí)其余位置傳零張量即可前向模式會自動跳過它們對結(jié)果的貢獻(xiàn)嚴(yán)格說是它們的 ε 分量為零不影響計(jì)算但也不省時(shí)間只是省得你寫新函數(shù)。如果函數(shù)的輸出不是一個(gè)張量而是一個(gè)元組jvp的返回也會是同結(jié)構(gòu)的元組一一對應(yīng)。這一點(diǎn)比torch.autograd.grad友好得多后者在輸出為元組時(shí)經(jīng)常讓人抓狂。3.2 用 vmap 把逐方向的 JVP 拼成完整雅可比單次 JVP 只給一列。想拿完整雅可比把標(biāo)準(zhǔn)基向量逐個(gè)喂進(jìn)去就行但 Python 循環(huán)太慢正確做法是vmapfrom torch.func import jvp, vmap def full_jacobian_via_jvp(f, x): n x.numel() basis torch.eye(n, dtypex.dtype, devicex.device) # 每次調(diào)用返回 J e_i即雅可比的第 i 列 cols vmap(lambda v: jvp(f, (x,), (v,))[1])(basis) return cols.T # 轉(zhuǎn)置成標(biāo)準(zhǔn)的 (m, n) 形狀 J full_jacobian_via_jvp(f, x) print(J) # tensor([[ 3.0000, 2.0000], # [ 0.0000, 6.0000], # [-0.4161, 0.0000]], dtypetorch.float64)結(jié)果和手算完全一致。這里有個(gè)容易搞反的地方vmap出來后cols的形狀是(n, m)第 i 行是第 i 列所以必須轉(zhuǎn)置。我第一次寫的時(shí)候忘了轉(zhuǎn)后面所有矩陣乘法全錯(cuò)排查了半天才發(fā)現(xiàn)是形狀問題。實(shí)際上官方已經(jīng)把這個(gè)封裝好了torch.func.jacfwd干的就是這件事from torch.func import jacfwd J jacfwd(f)(x) # 直接得到 (m, n)jacfwd內(nèi)部就是 vmap jvp性能基本一致能直接用就用它少寫一堆膠水代碼。3.3 jacfwd 和 jacrev按形狀選方向別選反這一對函數(shù)經(jīng)常被混用但它們的成本差異可以到一個(gè)數(shù)量級。規(guī)則很簡單jacfwd前向模式需要 n 次前向量級的計(jì)算適合n m輸入維度小于輸出維度jacrev反向模式需要 m 次反向量級的計(jì)算適合m n輸出維度小于輸入維度import time x_small torch.randn(10, dtypetorch.float64) x_large torch.randn(2000, dtypetorch.float64) def f_wide(x): # 10 - 2000輸入小輸出大 return torch.sin(x).repeat(200).reshape(-1)[:2000] * 1.5 def f_tall(x): # 2000 - 10輸入大輸出小 return torch.sin(x).reshape(200, 10).sum(0)對f_widejacfwd只要 10 次前向jacrev要 2000 次反向差距懸殊。對f_tall正好相反。跑一下就能直觀感受到我這里實(shí)測f_wide上jacfwd比jacrev快兩個(gè)數(shù)量級還多f_tall上則反過來。注意不要無腦用jacrev也不要無腦用jacfwd。判斷依據(jù)只有一條——比較輸入元素總數(shù)和輸出元素總數(shù)誰小就用誰做方向。這是前向/反向模式選擇的唯一標(biāo)準(zhǔn)和函數(shù)長什么樣沒關(guān)系。還有一個(gè)隱蔽的點(diǎn)jacfwd在輸入維度超過幾千時(shí)因?yàn)橐淮涡?vmap 出幾千個(gè)方向的 dual tensor顯存會爆。這時(shí)候該用chunk_size分批或者干脆切回jacrev。3.4 手工構(gòu)造 dual tensor 的底層寫法有時(shí)候你需要繞開torch.func的封裝直接操作 dual level比如需要在自定義 autograd Function 內(nèi)部嵌一段前向模式邏輯。底層寫法是這樣from torch.autograd.forward_ad import make_dual, unpack_dual, dual_level x torch.tensor([2.0, 3.0], dtypetorch.float64) v torch.tensor([1.0, 0.0], dtypetorch.float64) with dual_level(): x_dual make_dual(x, v) y_dual f(x_dual) y, jvp_out unpack_dual(y_dual) print(jvp_out) # tensor([ 3.0000, 0.0000, -0.4161])dual_level()是必須的上下文管理器出了這個(gè) with 塊之后 dual tensor 就失效了再拿去做運(yùn)算是未定義行為。另外注意只有支持前向模式的算子才能在里面用像某些第三方擴(kuò)展的自定義算子如果沒注冊 forward AD 規(guī)則會拋RuntimeError。遇到這種要么退回torch.autograd.functional.jvp要么給算子補(bǔ)一份jvp規(guī)則。4. 踩過的坑讓 JVP 結(jié)果對不上的七個(gè)原因4.1 原地操作和 BatchNorm 的 training 狀態(tài)torch.func.jvp對原地操作零容忍。x 1、x.relu_()、x.masked_fill_()這類都會直接報(bào)錯(cuò)因?yàn)樵诮o dual tensor 覆寫 ε 分量時(shí)前向模式無法確定正確的語義。解決辦法是全局搜一遍下劃線結(jié)尾的算子改成非原地版本。更隱蔽的是 BatchNorm。model.train()狀態(tài)下 BN 會更新 running_mean/running_var這屬于副作用side effect函數(shù)式變換默認(rèn)禁止副作用。要么切model.eval()要么用torch.func提供的functional_call把 buffer 顯式傳進(jìn)去。我遇到過一次 JVP 結(jié)果每次調(diào)用都不一樣的詭異情況查了兩小時(shí)才發(fā)現(xiàn)是 BN 在偷偷更新統(tǒng)計(jì)量。4.2 create_graph 開著然后拿結(jié)果去 backward這是最高頻的翻車點(diǎn)。假設(shè)你這么寫_, jv jvp(f, (x,), (v,)) jv.sum().backward() # RuntimeError 或者結(jié)果全 None因?yàn)閖vp默認(rèn)不構(gòu)造可微圖jv是葉子張量backward 拿不到梯度。想對 JVP 的結(jié)果再求導(dǎo)得用torch.func的變換組合方式from torch.func import jvp, grad # 求 d/dx [J(x) v]即二階信息 def hvp_scalar(x, v): _, jv jvp(grad(lambda t: f(t).sum()), (x,), (v,)) return jv.sum()也就是說把jvp和高階變換組合起來而不是指望 autograd 的圖自動接上。這是函數(shù)式 API 和傳統(tǒng) autograd 在思維方式上最大的差別。4.3 隨機(jī)性、dropout 和數(shù)值精度模型里有 dropout 時(shí)jvp的兩次內(nèi)部調(diào)用一次算輸出、一次算切線如果各自采樣一份 mask結(jié)果就會莫名其妙。torch.func系列變換提供了randomness參數(shù)可以設(shè)成differentiate或error后者會在檢測到隨機(jī)算子時(shí)直接報(bào)錯(cuò)逼你顯式處理。我的習(xí)慣是調(diào)試階段一律設(shè)error把問題提前暴露出來。精度方面JVP 涉及大量累加float32 在方向?qū)?shù)上經(jīng)常只留 4 到 5 位有效數(shù)字。做數(shù)值校驗(yàn)時(shí)務(wù)必切到 float64否則有限差分和 JVP 的差異會被舍入噪聲蓋住你會誤以為代碼寫錯(cuò)了其實(shí)是精度不夠。4.4 一個(gè)完整的排查鏈路我把自己遇到的一次典型問題完整記錄下來這個(gè)排查思路可以復(fù)用?，F(xiàn)象一個(gè) 8 維輸入、512 維輸出的函數(shù)torch.func.jvp的結(jié)果和有限差分的差異穩(wěn)定在 3% 左右不是隨機(jī)噪聲而是系統(tǒng)性偏差。第一步把輸入換成 float64 重跑差異降到 0.1%說明有大半是精度問題但仍有殘差。第二步把函數(shù)拆到最簡只留一層線性變換此時(shí) JVP 和手算完全一致說明前向模式本身沒問題。第三步逐層加回發(fā)現(xiàn)加上某個(gè)歸一化層后偏差出現(xiàn)。第四步檢查那個(gè)層發(fā)現(xiàn)它在內(nèi)部對輸入做了clamp(min0)而測試點(diǎn) x 的分量有負(fù)數(shù)clamp 在 0 處不可導(dǎo)有限差分跨過了這個(gè)折點(diǎn)于是兩邊不一致。第五步把測試點(diǎn)挪到所有分量都為正的區(qū)域差異消失。結(jié)論是JVP 本身沒錯(cuò)是有限差分在不可導(dǎo)點(diǎn)附近失效。這個(gè)教訓(xùn)很值錢——做數(shù)值校驗(yàn)時(shí)先確認(rèn)測試點(diǎn)遠(yuǎn)離 ReLU、abs、clamp 這類折點(diǎn)否則校驗(yàn)工具本身就在說謊。5. JVP 真正常用的四個(gè)場景5.1 Hessian-vector product二階優(yōu)化的基石海森矩陣是 n×n參數(shù)量上億時(shí)根本存不下但你往往只需要 H·v 這個(gè)乘積。而 H·v 恰好等于梯度的 JVPimport torch from torch.func import jvp, grad def loss_fn(x): return (x ** 3).sum() x torch.tensor([1.0, 2.0, 3.0], dtypetorch.float64) v torch.tensor([1.0, 1.0, 1.0], dtypetorch.float64) _, hv jvp(grad(loss_fn), (x,), (v,)) print(hv) # tensor([ 6., 12., 18.])正好是 6x 逐元素乘 v手算驗(yàn)證loss Σx3梯度是 3x2海森是對角陣 diag(6x)H·v 6x·v [6, 12, 18]。完全對上。這條式子在整個(gè)二階優(yōu)化領(lǐng)域是基礎(chǔ)設(shè)施。共軛梯度法解 H·p -g 需要反復(fù)算 H·v信賴域方法、自然梯度、K-FAC 的近似版本全都建立在這個(gè)乘積上。用傳統(tǒng)寫法torch.autograd.grad(g, x, grad_outputsv)也能算但那個(gè)要求 g 必須由create_graphTrue得到顯存開銷大jvp(grad(f), ...)走前向模式顯存友好得多。5.2 雅可比正則與輸入梯度懲罰對抗訓(xùn)練和魯棒性研究里常加一項(xiàng)雅可比正則懲罰模型對輸入變化的敏感度形式是 ||?f/?x||2。直接算完整雅可比代價(jià)太高用 Hutchinson 估計(jì)def jacobian_frobenius_sq(model, x, num_samples4): 估計(jì) ||J||_F^2x 是單個(gè)樣本 def f_single(xi): return model(xi.unsqueeze(0)).squeeze(0) total 0.0 for _ in range(num_samples): v torch.randint(0, 2, x.shape, dtypex.dtype, devicex.device) * 2 - 1 _, jv jvp(f_single, (x,), (v,)) total total (jv ** 2).sum() return total / num_samples原理是 E[||J·v||2] ||J||_F2只要 v 的分量是零均值、單位方差的隨機(jī)向量Rademacher 分布 ±1 就夠。四次采樣的估計(jì)就足夠穩(wěn)定成本只有四次前向。這個(gè)技巧在很多需要控制 Lipschitz 常數(shù)的工作里都是標(biāo)配。5.3 靈敏度分析與參數(shù)掃描工程仿真里經(jīng)常有這種需求一個(gè)模型有 30 個(gè)可調(diào)參數(shù)輸出是 5000 個(gè)測點(diǎn)的響應(yīng)曲線想知道參數(shù)動一點(diǎn)點(diǎn)輸出怎么變。這就是標(biāo)準(zhǔn)的 JVP 場景輸入小輸出大反向模式完全使不上勁。def simulate(params): # params: (30,) 返回: (5000,) ... params torch.randn(30, dtypetorch.float64) direction torch.zeros(30, dtypetorch.float64) direction[3] 1.0 # 只看第 4 個(gè)參數(shù)的單獨(dú)影響 _, sensitivity jvp(simulate, (params,), (direction,))這里的sensitivity就是第 4 個(gè)參數(shù)單獨(dú)變動一個(gè)單位時(shí)5000 個(gè)測點(diǎn)各自的響應(yīng)。用 VJP 做同樣的事得對 5000 個(gè)輸出各反向一次成本差兩個(gè)數(shù)量級。神經(jīng) ODE 的伴隨方法里也大量依賴這類操作只不過那邊用的是 VJP因?yàn)槟繕?biāo)通常是標(biāo)量損失。方向選擇永遠(yuǎn)跟著形狀走這一條不會變。5.4 用 trace 估計(jì)代替完整雅可比有時(shí)候你只想要雅可比的跡trace比如在流模型里估算對數(shù)行列式的一階近似。跡可以用 Hutchinson 直接估def trace_hutchinson(f, x, num_samples8): n x.numel() total 0.0 for _ in range(num_samples): v torch.randint(0, 2, (n,), dtypex.dtype, devicex.device) * 2 - 1 _, jv jvp(f, (x,), (v,)) total total (v * jv).sum() return total / num_samples因?yàn)?E[v?Jv] tr(J)。對一個(gè) 100 維的對角函數(shù)八次采樣就能把跡估到 5% 以內(nèi)的相對誤差而完整雅可比要 100 次前向。這個(gè)差距在參數(shù)維度上千時(shí)會變得不可接受。6. 性能與工程取舍什么時(shí)候不該用 JVP6.1 批量化vmap 一次比循環(huán)一百次強(qiáng)torch.func.jvp真正的威力在于它能被vmap包裹把一百個(gè)方向變成一次核函數(shù)調(diào)用。手寫 for 循環(huán)的話每個(gè)方向都要重新走一遍算子調(diào)度Python 開銷直接吃掉所有收益。實(shí)測在一個(gè)中等規(guī)模的網(wǎng)絡(luò)上vmap 版本比 Python 循環(huán)快 20 倍以上。如果不想用vmapPyTorch 還提供了另一條路torch.autograd.grad的is_grads_batchedTrue參數(shù)把多個(gè) VJP 的 cotangent 堆成一個(gè) batch 一次算完。y f(x) u_batched torch.eye(y.numel(), dtypey.dtype).reshape(y.numel(), *y.shape) grads torch.autograd.grad(y, x, grad_outputsu_batched, is_grads_batchedTrue)[0] # grads 形狀 (y.numel(), *x.shape)每一片是一行雅可比這條路的優(yōu)點(diǎn)是底層實(shí)現(xiàn)更貼近 autograd 原語某些老算子上比 vmap 兼容性更好。缺點(diǎn)是它算的是 VJP 批量方向選擇依然受形狀限制。6.2 顯存前向模式不是免費(fèi)的前向模式的顯存開銷常常被低估。每個(gè)參與運(yùn)算的張量都要額外存一份 ε 分量理論上翻倍再疊加vmap的批維度實(shí)際可能是三到四倍。當(dāng)輸入維度 n 達(dá)到幾千、中間激活又很大的時(shí)候jacfwd直接 OOM 是常事。幾條降顯存的手段按有效性排序用chunk_size參數(shù)分批 vmap把 n 個(gè)方向拆成若干組一組一組算再拼起來在torch.no_grad()外層包一層只讓需要求導(dǎo)的部分進(jìn)計(jì)算圖換jacrev如果你的形狀允許的話反向模式的顯存模型完全不同有時(shí)候反而更省降低精度到 bf16前向模式對精度的容忍度比反向模式高一些因?yàn)闆]有梯度累積的鏈?zhǔn)椒糯蟮闰?yàn)證數(shù)值穩(wěn)定性。6.3 一條選擇決策線把這些經(jīng)驗(yàn)壓縮成一條可以照著走的判斷線如果你的目標(biāo)是一個(gè)方向上的方向?qū)?shù)且函數(shù)的輸入維度不大用torch.func.jvp一次完事。如果目標(biāo)是完整雅可比先看形狀輸入元素少于輸出元素走jacfwd反之走jacrev。如果目標(biāo)是海森矩陣乘向量用jvp(grad(f), ...)組合別去碰functional.hvp那個(gè)內(nèi)部是反向兩次慢。如果目標(biāo)是雅可比范數(shù)或跡這類標(biāo)量統(tǒng)計(jì)量別算完整雅可比用 Hutchinson 加隨機(jī)方向估計(jì)采樣四到八次就夠。我個(gè)人的體會是torch.func這套 API 剛上手時(shí)最大的障礙不是語法而是思維方式的切換——它不依賴全局的requires_grad狀態(tài)也不依賴.backward()一切都是顯式的函數(shù)變換。一旦適應(yīng)了這種寫法代碼反而比傳統(tǒng)的 autograd 更清晰因?yàn)閷κ裁辞髮?dǎo)、沿著什么方向全都寫在函數(shù)簽名里不用去猜計(jì)算圖長什么樣。真正需要警惕的只有兩點(diǎn)一是別在需要前向模式的場景里誤用torch.autograd.functional.jvp二是做數(shù)值校驗(yàn)時(shí)務(wù)必用 float64 并且避開折點(diǎn)這兩條踩過一次就夠記一輩子了。 PyTorch JVP實(shí)戰(zhàn)：torch.func.jvp原理、性能陷阱與工程取舍
大多數(shù)用過 PyTorch 的人對loss.backward()這套反向傳播都很熟但一旦被問到雅可比向量積怎么算很多人第一反應(yīng)是打開torch.autograd.functional.jacobian把整個(gè)雅可比矩陣求出來再手動做一次矩陣乘法。這種做法在小模型上沒毛病一旦輸出維度上千顯存和耗時(shí)立刻教做人。雅可比向量積Jacobian-Vector Product簡稱 JVP存在的意義就是你根本不需要顯式構(gòu)造那個(gè) m×n 的雅可比矩陣也能拿到它和一個(gè)向量的乘積代價(jià)只有一次前向傳播的量級。這篇內(nèi)容圍繞 JVP 在 PyTorch 里的幾種實(shí)現(xiàn)路徑展開重點(diǎn)拆torch.func.jvp的用法、torch.autograd.functional.jvp的性能陷阱、以及踩過的坑。適合已經(jīng)會寫autograd.grad、想進(jìn)一步搞清楚前向模式自動微分、并且需要在二階優(yōu)化、雅可比正則、靈敏度分析這些場景里真正用起來的人。1. 把 JVP 從矩陣?yán)锊鸪鰜碚l乘在左邊誰乘在右邊1.1 一個(gè) 3×2 的例子先把概念釘死先不談自動微分回到線性代數(shù)本身。假設(shè)有一個(gè)函數(shù) f: R2 → R3在點(diǎn) x 處它的雅可比矩陣 J 是一個(gè) 3 行 2 列的矩陣。JVP 就是拿一個(gè) 2 維向量 v從右邊乘上去得到 J·v結(jié)果是一個(gè) 3 維向量。注意 v 的維度必須和輸入維度一致結(jié)果的維度必須和輸出維度一致。反過來VJPVector-Jacobian Product是拿一個(gè) 3 維向量 u 從左邊乘u?·J結(jié)果是一個(gè) 2 維向量維度跟著輸入走。這兩個(gè)操作的形狀規(guī)則完全不同這是后面所有坑的根源。拿一個(gè)具體例子import torch def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x torch.tensor([2.0, 3.0], dtypetorch.float64)手工求偏導(dǎo)J 應(yīng)該是J [[x1, x0 ], [0, 2*x1 ], [cos(x0), 0 ]]代入 x [2, 3]得到[[3, 2], [0, 6], [cos(2), 0]]其中 cos(2) ≈ -0.4161?，F(xiàn)在取 v [1, 0]J·v 就是 J 的第一列即[3, 0, -0.4161]。取 v [0, 1]得到第二列[2, 6, 0]。這兩個(gè)結(jié)果后面用來驗(yàn)證代碼是否正確。1.2 為什么 PyTorch 默認(rèn)只給你 VJPloss.backward()背后跑的是反向模式自動微分它天然產(chǎn)出的是 VJP不是 JVP。原因很實(shí)際深度學(xué)習(xí)里絕大多數(shù)函數(shù)是從一個(gè)巨大的參數(shù)向量映射到一個(gè)標(biāo)量損失也就是 f: R^N → R1。這種情況下輸出維度 m 1只需要一次反向傳播就能拿到完整的梯度向量代價(jià)極低。而如果換成 JVP 的思路你得做 N 次前向傳播才能拿全梯度完全不可行。所以框架的設(shè)計(jì)重心一直壓在 VJP 上前向模式自動微分在很長一段時(shí)間里屬于有但不太好用的狀態(tài)。這直接導(dǎo)致一個(gè)現(xiàn)象當(dāng)你遇到 f: R^n → R^m 且 n 很小、m 很大的場景時(shí)比如一個(gè)幾十維的控制參數(shù)映射到上千維的仿真輸出用反向模式硬算會非常虧而這恰恰是 JVP 的主場。用一個(gè)表把兩種模式的成本關(guān)系擺清楚這張表值得貼在顯示器邊上目標(biāo)操作需要調(diào)用幾次適合的形狀PyTorch 對應(yīng)接口完整雅可比逐列前向n 次前向n 小 m 大torch.func.jacfwd完整雅可比逐行反向m 次反向m 小 n 大torch.func.jacrev單次 JVPJ·v1 次前向量級輸入小torch.func.jvp單次 VJPu?·J1 次反向量級輸出小torch.func.vjp/autograd.grad1.3 前向模式到底在算什么dual number 的直覺前向模式自動微分的實(shí)現(xiàn)方式叫對偶數(shù)dual number。把每個(gè)標(biāo)量 a 擴(kuò)展成 a b·ε 的形式規(guī)定 ε ≠ 0 但 ε2 0。任何函數(shù)作用上去之后實(shí)部就是函數(shù)值ε 的系數(shù)就是方向?qū)?shù)。舉個(gè)手算的例子x 2 1·ε那么 x2 4 4ε ε2 4 4εε 的系數(shù) 4 正好是 2x 在 x2 處的導(dǎo)數(shù)乘上方向 1。這個(gè)機(jī)制的好處是一次前向傳播就能同時(shí)算出函數(shù)值和沿著某個(gè)方向 v 的方向?qū)?shù)而且天然支持任意深的網(wǎng)絡(luò)不需要保存中間激活。代價(jià)是每個(gè)標(biāo)量都要額外存一份 ε 系數(shù)顯存占用大約翻倍但和前向傳播同一量級。PyTorch 里這套機(jī)制的底層入口是torch.autograd.forward_adtorch.func.jvp是它的高層封裝。理解這一層很關(guān)鍵因?yàn)楹竺鏁吹絫orch.autograd.functional.jvp走的根本不是這條路。2. PyTorch 里算出 JVP 的四條路各自的真實(shí)代價(jià)2.1 torch.func.jvp目前最該用的那條PyTorch 2.0 之后torch.func命名空間下提供了一套函數(shù)式變換 APIjvp是其中之一。它的簽名是jvp(func, primals, tangents)primals是函數(shù)位置參數(shù)的元組tangents是切線方向的元組兩者結(jié)構(gòu)必須一致。返回值是一個(gè)二元組(output, jvp_out)output是函數(shù)原本的輸出jvp_out是 J·v 的結(jié)果。import torch from torch.func import jvp def f(x): return torch.stack([x[0] * x[1], x[1] ** 2, torch.sin(x[0])]) x torch.tensor([2.0, 3.0], dtypetorch.float64) v torch.tensor([1.0, 0.0], dtypetorch.float64) out, jvp_out jvp(f, (x,), (v,)) print(out) # tensor([ 6.0000, 9.0000, 0.9093], dtypetorch.float64) print(jvp_out) # tensor([ 3.0000, 0.0000, -0.4161], dtypetorch.float64)跑出來和手算的[3, 0, -0.4161]完全對上。注意兩個(gè)細(xì)節(jié)第一x本身不需要requires_gradTrue前向模式不依賴計(jì)算圖第二f內(nèi)部不允許有原地操作in-place比如x[0] 1這種寫法會直接報(bào)錯(cuò)因?yàn)?dual tensor 的 ε 分量不允許被就地覆蓋。2.2 torch.autograd.functional.jvp名字一樣實(shí)現(xiàn)完全不同很多人第一次搜到的是torch.autograd.functional.jvp。它確實(shí)能用簽名也幾乎一樣from torch.autograd.functional import jvp out, jvp_out jvp(f, (x,), (v,), create_graphFalse, strictFalse)但這里有個(gè)非常容易被忽略的事實(shí)這個(gè)函數(shù)內(nèi)部是用反向模式實(shí)現(xiàn)的。它的做法是先構(gòu)造出 VJP 的形式再對切向量 v 求一次導(dǎo)本質(zhì)上是 reverse-over-reverse也就是兩次反向傳播拼出來的 JVP。這意味著什么意味著當(dāng)你的函數(shù)是輸入維度小、輸出維度大時(shí)torch.autograd.functional.jvp反而會比真正的torch.func.jvp慢一大截因?yàn)樗刂敵鼍S度反向走一遍。我在一個(gè) 50 維輸入、2000 維輸出的仿真函數(shù)上做過對比torch.func.jvp耗時(shí)大約是functional.jvp的三分之一差距相當(dāng)明顯。提示torch.autograd.functional.jvp的create_graphTrue是它為數(shù)不多的優(yōu)勢場景——需要 JVP 結(jié)果本身再參與求導(dǎo)比如構(gòu)造二階以上的混合導(dǎo)數(shù)時(shí)它比函數(shù)式 API 更省事。日常只要一階 JVP優(yōu)先torch.func.jvp。2.3 手搓 double-backward 與有限差分第三條路是老派做法算兩次autograd.grad。這不是真正的 JVP而是構(gòu)造 VJP 再套一層def manual_jvp(f, x, v): x x.detach().requires_grad_(True) y f(x) # 先拿一個(gè)啞元的 VJP把 v 塞進(jìn) grad_outputs g torch.autograd.grad(y, x, grad_outputstorch.zeros_like(y), create_graphTrue) ...說實(shí)話這條路現(xiàn)在沒什么推薦價(jià)值了代碼繞、易錯(cuò)、還容易在create_graph上翻車除非你在維護(hù)一份必須兼容老版本 PyTorch 的代碼庫否則直接跳過。第四條路是有限差分(f(x h*v) - f(x)) / h。它的缺點(diǎn)寫在教科書里——截?cái)嗾`差和舍入誤差互相拉扯h 太小結(jié)果被噪聲吃掉h 太大非線性誤差明顯。不過它有個(gè)獨(dú)特價(jià)值作為數(shù)值校驗(yàn)的金標(biāo)準(zhǔn)。你寫完任何 JVP 實(shí)現(xiàn)拿有限差分對一遍能立刻發(fā)現(xiàn)是不是方向搞反了、是不是 v 的維度用錯(cuò)了。我習(xí)慣用h 1e-6配 float64 做對照誤差在 1e-6 量級就算通過。2.4 四條路的橫向?qū)Ρ嚷窂绞欠裾媲跋蚰Ｊ较鄬臅r(shí)需要 requires_grad支持二階推薦場景torch.func.jvp是1x基線否配合create_graph式組合首選尤其 n 小時(shí)的批量 JVPtorch.autograd.functional.jvp否反向兩次2~4x否strictFalse原生支持需要高階混合導(dǎo)手搓autograd.grad否3x 以上是麻煩老版本兼容有限差分近似2x否不支持?jǐn)?shù)值校驗(yàn)3. torch.func.jvp 的實(shí)戰(zhàn)用法從單方向到整塊雅可比3.1 參數(shù)結(jié)構(gòu)不匹配是最常見的報(bào)錯(cuò)來源jvp(func, primals, tangents)這三者的結(jié)構(gòu)約束很嚴(yán)格。primals和tangents必須是元組長度等于func的位置參數(shù)個(gè)數(shù)每個(gè)位置的張量形狀必須逐一對上。如果你的函數(shù)是def f(a, b)那就得寫jvp(f, (a, b), (va, vb))漏一個(gè)或者包錯(cuò)層級都會報(bào)形狀錯(cuò)誤。def g(a, b): return (a * b).sum() b.pow(2).sum() a torch.randn(3, dtypetorch.float64) b torch.randn(4, dtypetorch.float64) va torch.ones(3, dtypetorch.float64) vb torch.zeros(4, dtypetorch.float64) out, jvp_out jvp(g, (a, b), (va, vb)) # 輸出是標(biāo)量對 a 方向的偏導(dǎo)之和這里vb傳全零意思是只要 a 方向的導(dǎo)數(shù)。這是一個(gè)非常實(shí)用的技巧當(dāng)你只關(guān)心某一個(gè)輸入的分量時(shí)其余位置傳零張量即可前向模式會自動跳過它們對結(jié)果的貢獻(xiàn)嚴(yán)格說是它們的 ε 分量為零不影響計(jì)算但也不省時(shí)間只是省得你寫新函數(shù)。如果函數(shù)的輸出不是一個(gè)張量而是一個(gè)元組jvp的返回也會是同結(jié)構(gòu)的元組一一對應(yīng)。這一點(diǎn)比torch.autograd.grad友好得多后者在輸出為元組時(shí)經(jīng)常讓人抓狂。3.2 用 vmap 把逐方向的 JVP 拼成完整雅可比單次 JVP 只給一列。想拿完整雅可比把標(biāo)準(zhǔn)基向量逐個(gè)喂進(jìn)去就行但 Python 循環(huán)太慢正確做法是vmapfrom torch.func import jvp, vmap def full_jacobian_via_jvp(f, x): n x.numel() basis torch.eye(n, dtypex.dtype, devicex.device) # 每次調(diào)用返回 J e_i即雅可比的第 i 列 cols vmap(lambda v: jvp(f, (x,), (v,))[1])(basis) return cols.T # 轉(zhuǎn)置成標(biāo)準(zhǔn)的 (m, n) 形狀 J full_jacobian_via_jvp(f, x) print(J) # tensor([[ 3.0000, 2.0000], # [ 0.0000, 6.0000], # [-0.4161, 0.0000]], dtypetorch.float64)結(jié)果和手算完全一致。這里有個(gè)容易搞反的地方vmap出來后cols的形狀是(n, m)第 i 行是第 i 列所以必須轉(zhuǎn)置。我第一次寫的時(shí)候忘了轉(zhuǎn)后面所有矩陣乘法全錯(cuò)排查了半天才發(fā)現(xiàn)是形狀問題。實(shí)際上官方已經(jīng)把這個(gè)封裝好了torch.func.jacfwd干的就是這件事from torch.func import jacfwd J jacfwd(f)(x) # 直接得到 (m, n)jacfwd內(nèi)部就是 vmap jvp性能基本一致能直接用就用它少寫一堆膠水代碼。3.3 jacfwd 和 jacrev按形狀選方向別選反這一對函數(shù)經(jīng)常被混用但它們的成本差異可以到一個(gè)數(shù)量級。規(guī)則很簡單jacfwd前向模式需要 n 次前向量級的計(jì)算適合n m輸入維度小于輸出維度jacrev反向模式需要 m 次反向量級的計(jì)算適合m n輸出維度小于輸入維度import time x_small torch.randn(10, dtypetorch.float64) x_large torch.randn(2000, dtypetorch.float64) def f_wide(x): # 10 - 2000輸入小輸出大 return torch.sin(x).repeat(200).reshape(-1)[:2000] * 1.5 def f_tall(x): # 2000 - 10輸入大輸出小 return torch.sin(x).reshape(200, 10).sum(0)對f_widejacfwd只要 10 次前向jacrev要 2000 次反向差距懸殊。對f_tall正好相反。跑一下就能直觀感受到我這里實(shí)測f_wide上jacfwd比jacrev快兩個(gè)數(shù)量級還多f_tall上則反過來。注意不要無腦用jacrev也不要無腦用jacfwd。判斷依據(jù)只有一條——比較輸入元素總數(shù)和輸出元素總數(shù)誰小就用誰做方向。這是前向/反向模式選擇的唯一標(biāo)準(zhǔn)和函數(shù)長什么樣沒關(guān)系。還有一個(gè)隱蔽的點(diǎn)jacfwd在輸入維度超過幾千時(shí)因?yàn)橐淮涡?vmap 出幾千個(gè)方向的 dual tensor顯存會爆。這時(shí)候該用chunk_size分批或者干脆切回jacrev。3.4 手工構(gòu)造 dual tensor 的底層寫法有時(shí)候你需要繞開torch.func的封裝直接操作 dual level比如需要在自定義 autograd Function 內(nèi)部嵌一段前向模式邏輯。底層寫法是這樣from torch.autograd.forward_ad import make_dual, unpack_dual, dual_level x torch.tensor([2.0, 3.0], dtypetorch.float64) v torch.tensor([1.0, 0.0], dtypetorch.float64) with dual_level(): x_dual make_dual(x, v) y_dual f(x_dual) y, jvp_out unpack_dual(y_dual) print(jvp_out) # tensor([ 3.0000, 0.0000, -0.4161])dual_level()是必須的上下文管理器出了這個(gè) with 塊之后 dual tensor 就失效了再拿去做運(yùn)算是未定義行為。另外注意只有支持前向模式的算子才能在里面用像某些第三方擴(kuò)展的自定義算子如果沒注冊 forward AD 規(guī)則會拋RuntimeError。遇到這種要么退回torch.autograd.functional.jvp要么給算子補(bǔ)一份jvp規(guī)則。4. 踩過的坑讓 JVP 結(jié)果對不上的七個(gè)原因4.1 原地操作和 BatchNorm 的 training 狀態(tài)torch.func.jvp對原地操作零容忍。x 1、x.relu_()、x.masked_fill_()這類都會直接報(bào)錯(cuò)因?yàn)樵诮o dual tensor 覆寫 ε 分量時(shí)前向模式無法確定正確的語義。解決辦法是全局搜一遍下劃線結(jié)尾的算子改成非原地版本。更隱蔽的是 BatchNorm。model.train()狀態(tài)下 BN 會更新 running_mean/running_var這屬于副作用side effect函數(shù)式變換默認(rèn)禁止副作用。要么切model.eval()要么用torch.func提供的functional_call把 buffer 顯式傳進(jìn)去。我遇到過一次 JVP 結(jié)果每次調(diào)用都不一樣的詭異情況查了兩小時(shí)才發(fā)現(xiàn)是 BN 在偷偷更新統(tǒng)計(jì)量。4.2 create_graph 開著然后拿結(jié)果去 backward這是最高頻的翻車點(diǎn)。假設(shè)你這么寫_, jv jvp(f, (x,), (v,)) jv.sum().backward() # RuntimeError 或者結(jié)果全 None因?yàn)閖vp默認(rèn)不構(gòu)造可微圖jv是葉子張量backward 拿不到梯度。想對 JVP 的結(jié)果再求導(dǎo)得用torch.func的變換組合方式from torch.func import jvp, grad # 求 d/dx [J(x) v]即二階信息 def hvp_scalar(x, v): _, jv jvp(grad(lambda t: f(t).sum()), (x,), (v,)) return jv.sum()也就是說把jvp和高階變換組合起來而不是指望 autograd 的圖自動接上。這是函數(shù)式 API 和傳統(tǒng) autograd 在思維方式上最大的差別。4.3 隨機(jī)性、dropout 和數(shù)值精度模型里有 dropout 時(shí)jvp的兩次內(nèi)部調(diào)用一次算輸出、一次算切線如果各自采樣一份 mask結(jié)果就會莫名其妙。torch.func系列變換提供了randomness參數(shù)可以設(shè)成differentiate或error后者會在檢測到隨機(jī)算子時(shí)直接報(bào)錯(cuò)逼你顯式處理。我的習(xí)慣是調(diào)試階段一律設(shè)error把問題提前暴露出來。精度方面JVP 涉及大量累加float32 在方向?qū)?shù)上經(jīng)常只留 4 到 5 位有效數(shù)字。做數(shù)值校驗(yàn)時(shí)務(wù)必切到 float64否則有限差分和 JVP 的差異會被舍入噪聲蓋住你會誤以為代碼寫錯(cuò)了其實(shí)是精度不夠。4.4 一個(gè)完整的排查鏈路我把自己遇到的一次典型問題完整記錄下來這個(gè)排查思路可以復(fù)用?，F(xiàn)象一個(gè) 8 維輸入、512 維輸出的函數(shù)torch.func.jvp的結(jié)果和有限差分的差異穩(wěn)定在 3% 左右不是隨機(jī)噪聲而是系統(tǒng)性偏差。第一步把輸入換成 float64 重跑差異降到 0.1%說明有大半是精度問題但仍有殘差。第二步把函數(shù)拆到最簡只留一層線性變換此時(shí) JVP 和手算完全一致說明前向模式本身沒問題。第三步逐層加回發(fā)現(xiàn)加上某個(gè)歸一化層后偏差出現(xiàn)。第四步檢查那個(gè)層發(fā)現(xiàn)它在內(nèi)部對輸入做了clamp(min0)而測試點(diǎn) x 的分量有負(fù)數(shù)clamp 在 0 處不可導(dǎo)有限差分跨過了這個(gè)折點(diǎn)于是兩邊不一致。第五步把測試點(diǎn)挪到所有分量都為正的區(qū)域差異消失。結(jié)論是JVP 本身沒錯(cuò)是有限差分在不可導(dǎo)點(diǎn)附近失效。這個(gè)教訓(xùn)很值錢——做數(shù)值校驗(yàn)時(shí)先確認(rèn)測試點(diǎn)遠(yuǎn)離 ReLU、abs、clamp 這類折點(diǎn)否則校驗(yàn)工具本身就在說謊。5. JVP 真正常用的四個(gè)場景5.1 Hessian-vector product二階優(yōu)化的基石海森矩陣是 n×n參數(shù)量上億時(shí)根本存不下但你往往只需要 H·v 這個(gè)乘積。而 H·v 恰好等于梯度的 JVPimport torch from torch.func import jvp, grad def loss_fn(x): return (x ** 3).sum() x torch.tensor([1.0, 2.0, 3.0], dtypetorch.float64) v torch.tensor([1.0, 1.0, 1.0], dtypetorch.float64) _, hv jvp(grad(loss_fn), (x,), (v,)) print(hv) # tensor([ 6., 12., 18.])正好是 6x 逐元素乘 v手算驗(yàn)證loss Σx3梯度是 3x2海森是對角陣 diag(6x)H·v 6x·v [6, 12, 18]。完全對上。這條式子在整個(gè)二階優(yōu)化領(lǐng)域是基礎(chǔ)設(shè)施。共軛梯度法解 H·p -g 需要反復(fù)算 H·v信賴域方法、自然梯度、K-FAC 的近似版本全都建立在這個(gè)乘積上。用傳統(tǒng)寫法torch.autograd.grad(g, x, grad_outputsv)也能算但那個(gè)要求 g 必須由create_graphTrue得到顯存開銷大jvp(grad(f), ...)走前向模式顯存友好得多。5.2 雅可比正則與輸入梯度懲罰對抗訓(xùn)練和魯棒性研究里常加一項(xiàng)雅可比正則懲罰模型對輸入變化的敏感度形式是 ||?f/?x||2。直接算完整雅可比代價(jià)太高用 Hutchinson 估計(jì)def jacobian_frobenius_sq(model, x, num_samples4): 估計(jì) ||J||_F^2x 是單個(gè)樣本 def f_single(xi): return model(xi.unsqueeze(0)).squeeze(0) total 0.0 for _ in range(num_samples): v torch.randint(0, 2, x.shape, dtypex.dtype, devicex.device) * 2 - 1 _, jv jvp(f_single, (x,), (v,)) total total (jv ** 2).sum() return total / num_samples原理是 E[||J·v||2] ||J||_F2只要 v 的分量是零均值、單位方差的隨機(jī)向量Rademacher 分布 ±1 就夠。四次采樣的估計(jì)就足夠穩(wěn)定成本只有四次前向。這個(gè)技巧在很多需要控制 Lipschitz 常數(shù)的工作里都是標(biāo)配。5.3 靈敏度分析與參數(shù)掃描工程仿真里經(jīng)常有這種需求一個(gè)模型有 30 個(gè)可調(diào)參數(shù)輸出是 5000 個(gè)測點(diǎn)的響應(yīng)曲線想知道參數(shù)動一點(diǎn)點(diǎn)輸出怎么變。這就是標(biāo)準(zhǔn)的 JVP 場景輸入小輸出大反向模式完全使不上勁。def simulate(params): # params: (30,) 返回: (5000,) ... params torch.randn(30, dtypetorch.float64) direction torch.zeros(30, dtypetorch.float64) direction[3] 1.0 # 只看第 4 個(gè)參數(shù)的單獨(dú)影響 _, sensitivity jvp(simulate, (params,), (direction,))這里的sensitivity就是第 4 個(gè)參數(shù)單獨(dú)變動一個(gè)單位時(shí)5000 個(gè)測點(diǎn)各自的響應(yīng)。用 VJP 做同樣的事得對 5000 個(gè)輸出各反向一次成本差兩個(gè)數(shù)量級。神經(jīng) ODE 的伴隨方法里也大量依賴這類操作只不過那邊用的是 VJP因?yàn)槟繕?biāo)通常是標(biāo)量損失。方向選擇永遠(yuǎn)跟著形狀走這一條不會變。5.4 用 trace 估計(jì)代替完整雅可比有時(shí)候你只想要雅可比的跡trace比如在流模型里估算對數(shù)行列式的一階近似。跡可以用 Hutchinson 直接估def trace_hutchinson(f, x, num_samples8): n x.numel() total 0.0 for _ in range(num_samples): v torch.randint(0, 2, (n,), dtypex.dtype, devicex.device) * 2 - 1 _, jv jvp(f, (x,), (v,)) total total (v * jv).sum() return total / num_samples因?yàn)?E[v?Jv] tr(J)。對一個(gè) 100 維的對角函數(shù)八次采樣就能把跡估到 5% 以內(nèi)的相對誤差而完整雅可比要 100 次前向。這個(gè)差距在參數(shù)維度上千時(shí)會變得不可接受。6. 性能與工程取舍什么時(shí)候不該用 JVP6.1 批量化vmap 一次比循環(huán)一百次強(qiáng)torch.func.jvp真正的威力在于它能被vmap包裹把一百個(gè)方向變成一次核函數(shù)調(diào)用。手寫 for 循環(huán)的話每個(gè)方向都要重新走一遍算子調(diào)度Python 開銷直接吃掉所有收益。實(shí)測在一個(gè)中等規(guī)模的網(wǎng)絡(luò)上vmap 版本比 Python 循環(huán)快 20 倍以上。如果不想用vmapPyTorch 還提供了另一條路torch.autograd.grad的is_grads_batchedTrue參數(shù)把多個(gè) VJP 的 cotangent 堆成一個(gè) batch 一次算完。y f(x) u_batched torch.eye(y.numel(), dtypey.dtype).reshape(y.numel(), *y.shape) grads torch.autograd.grad(y, x, grad_outputsu_batched, is_grads_batchedTrue)[0] # grads 形狀 (y.numel(), *x.shape)每一片是一行雅可比這條路的優(yōu)點(diǎn)是底層實(shí)現(xiàn)更貼近 autograd 原語某些老算子上比 vmap 兼容性更好。缺點(diǎn)是它算的是 VJP 批量方向選擇依然受形狀限制。6.2 顯存前向模式不是免費(fèi)的前向模式的顯存開銷常常被低估。每個(gè)參與運(yùn)算的張量都要額外存一份 ε 分量理論上翻倍再疊加vmap的批維度實(shí)際可能是三到四倍。當(dāng)輸入維度 n 達(dá)到幾千、中間激活又很大的時(shí)候jacfwd直接 OOM 是常事。幾條降顯存的手段按有效性排序用chunk_size參數(shù)分批 vmap把 n 個(gè)方向拆成若干組一組一組算再拼起來在torch.no_grad()外層包一層只讓需要求導(dǎo)的部分進(jìn)計(jì)算圖換jacrev如果你的形狀允許的話反向模式的顯存模型完全不同有時(shí)候反而更省降低精度到 bf16前向模式對精度的容忍度比反向模式高一些因?yàn)闆]有梯度累積的鏈?zhǔn)椒糯蟮闰?yàn)證數(shù)值穩(wěn)定性。6.3 一條選擇決策線把這些經(jīng)驗(yàn)壓縮成一條可以照著走的判斷線如果你的目標(biāo)是一個(gè)方向上的方向?qū)?shù)且函數(shù)的輸入維度不大用torch.func.jvp一次完事。如果目標(biāo)是完整雅可比先看形狀輸入元素少于輸出元素走jacfwd反之走jacrev。如果目標(biāo)是海森矩陣乘向量用jvp(grad(f), ...)組合別去碰functional.hvp那個(gè)內(nèi)部是反向兩次慢。如果目標(biāo)是雅可比范數(shù)或跡這類標(biāo)量統(tǒng)計(jì)量別算完整雅可比用 Hutchinson 加隨機(jī)方向估計(jì)采樣四到八次就夠。我個(gè)人的體會是torch.func這套 API 剛上手時(shí)最大的障礙不是語法而是思維方式的切換——它不依賴全局的requires_grad狀態(tài)也不依賴.backward()一切都是顯式的函數(shù)變換。一旦適應(yīng)了這種寫法代碼反而比傳統(tǒng)的 autograd 更清晰因?yàn)閷κ裁辞髮?dǎo)、沿著什么方向全都寫在函數(shù)簽名里不用去猜計(jì)算圖長什么樣。真正需要警惕的只有兩點(diǎn)一是別在需要前向模式的場景里誤用torch.autograd.functional.jvp二是做數(shù)值校驗(yàn)時(shí)務(wù)必用 float64 并且避開折點(diǎn)這兩條踩過一次就夠記一輩子了。