4.3 Jacobi-mátrix, láncszabály és VJP

4.3.1 Cél, előismeret és induló kérdés

A 4.2. Parciális derivált és gradiens alfejezetben egy $f:\mathbb R^d\to\mathbb R$ skalár függvény gradiensét értelmeztük. A neurális réteg azonban tipikusan vektorból vektort képez: egy affin réteg $\mathbb R^d$-ből $\mathbb R^m$-be, egy aktiváció réteg szintén $\mathbb R^m$-ből $\mathbb R^m$-be. A tanítás során a veszteség skalár, tehát a visszafelé haladó számításnak „csak” egy-egy irány menti szorzatot kell elvégeznie, nem kell felépítenie a teljes mátrix deriváltat.

Induló becslés: ha $f:\mathbb R^{10000}\to\mathbb R^{1000}$, akkor a teljes Jacobi-mátrix $1000\cdot10000=10^7$ számot tartalmazna. Ha viszont csak egy $g\in\mathbb R^{1000}$ kimeneti vektort vetítünk vissza, a kívánt bemeneti gradiens maga is csupán $10000$ szám. A kérdés: milyen mátrix–vektor szorzat és milyen transzponálási konvenció adja helyesen ezt az olcsóbb számítást?

A cél egy pontos, alakhelyes nyelv: Jacobi-mátrix, kompozíció, JVP, VJP, batch-redukció, valamint az affin réteg súly-, bias- és inputgradiensének levezetése. Előfeltétel a 4.2. Parciális derivált és gradiens parciális derivált- és gradiensfogalma; a 3.3. Vetítés, legkisebb négyzetek és kondicionáltság alfejezetből a legkisebb négyzetes normálegyenletre is röviden visszautalunk.

4.3.2 A Jacobi-mátrix: definíció és alak

Definíció (Jacobi-mátrix). Legyen $U\subseteq\mathbb R^d$ nyílt, $f:U\to\mathbb R^m$, és legyen $f$ differenciálható az $x\in U$ pontban. Ekkor az $f$ Jacobi-mátrixa

$$J_f(x)\in\mathbb R^{m\times d}, \qquad [J_f(x)]_{i,j}=\frac{\partial f_i}{\partial x_j}(x).$$

A mátrix sora az output, oszlopa az input komponenshez tartozik. A mátrix $i$-edik sora az $f_i$ skalár függvény gradiensének transzponáltja:

$$J_f(x)=\begin{bmatrix} (\nabla f_1(x))^\top\\ \vdots\\ (\nabla f_m(x))^\top \end{bmatrix}.$$

A differenciálhatóság definíciója éppen azt jelenti, hogy a maradékos lineáris közelítés minden irányból jó:

$$f(x+u)=f(x)+J_f(x)u+r(u), \qquad \lim_{u\to0}\frac{\|r(u)\|}{\|u\|}=0.$$

Alakellenőrzés: $J_f(x)$ mérete $m\times d$, $u$ mérete $d\times1$, tehát $J_f(x)u$ mérete $m\times1$, azonos $f(x+u)-f(x)$ alakjával. Ha $m=1$, akkor $J_f(x)$ egy $1\times d$ sor, és $J_f(x)^\top=\nabla f(x)$, ami visszaadja a 4.2. Parciális derivált és gradiens gradiensét.

4.3.3 Kompozíció és a láncszabály

A rétegzett háló matematikailag kompozíció: a következő réteg bemenete az előző réteg kimenete. A láncszabály mondja meg, hogyan szorzódnak a Jacobi-mátrixok.

Tétel (többváltozós láncszabály). Legyen $U\subseteq\mathbb R^d$ és $V\subseteq\mathbb R^m$ nyílt, $f:U\to V$ differenciálható $x\in U$-ban, $g:V\to\mathbb R^k$ differenciálható $f(x)\in V$-ben. Ekkor $h=g\circ f$ differenciálható $x$-ben, és

$$J_h(x)=J_g(f(x))\,J_f(x).$$

A szorzat alakjai: $J_g(f(x))\in\mathbb R^{k\times m}$, $J_f(x)\in\mathbb R^{m\times d}$, ezért a szorzat $\mathbb R^{k\times d}$-beli.

Bizonyításvázlat. A bizonyítás a differenciálhatóság elsőrendű maradékos alakját használja. Legyen $u\in\mathbb R^d$, és írjuk

$$f(x+u)=f(x)+J_f(x)u+r_f(u), \qquad \frac{\|r_f(u)\|}{\|u\|}\to0.$$

Legyen $\delta=f(x+u)-f(x)=J_f(x)u+r_f(u)$. Mivel $g$ differenciálható $f(x)$-ben,

$$g(f(x)+\delta)=g(f(x))+J_g(f(x))\delta+r_g(\delta), \qquad \frac{\|r_g(\delta)\|}{\|\delta\|}\to0.$$

Behelyettesítve $\delta$ értékét:

$$h(x+u)=h(x)+J_g(f(x))J_f(x)u+J_g(f(x))r_f(u)+r_g(\delta).$$

Az első maradéktag nyilván $o(\|u\|)$, mert $r_f$ az. A másodikhoz azt használjuk, hogy $\delta=O(\|u\|)$: $J_f$ folytonos lineáris leképezés, ezért alkalmas $C>0$ mátrixnormával $\|J_f(x)u\|\le C\|u\|$, és $\|r_f(u)\|/\|u\|$ korlátos elég kis $u$-ra. Így $\delta\to0$, amikor $u\to0$, tehát $\|r_g(\delta)\|/\|\delta\|\to0$, miközben $\|\delta\|/\|u\|$ korlátos. Ezért a teljes maradék $o(\|u\|)$, ami éppen a kívánt lineáris közelítést adja. A teljes $\varepsilon$-$\delta$ kidolgozás a szokásos elemi kalkulus bizonyítás; a fenti vázlat a saját levezetés, a feltételek a kimondott differenciálhatósági pontok.

A tétel feltételei elégségesek, de nem szükségesek. Példa: $f(x)=0$ konstans, $g(t)=|t|$. Ekkor $g$ nem differenciálható $f(x)=0$-ban, ezért a láncszabály képlete nem alkalmazható; a $g\circ f$ mégis azonosan nulla, tehát differenciálható. Ezért a tétel használatakor mindig ellenőrizni kell, hogy a belső függvény kimeneti pontjában a külső függvény tényleg differenciálható-e.

4.3.4 VJP, JVP, batch és reduction

Definíció (JVP). Adott $f:\mathbb R^d\to\mathbb R^m$ és adott $v\in\mathbb R^d$ irányvektor esetén a Jacobian–vector product

$$\mathrm{JVP}_f(x;v)=J_f(x)v\in\mathbb R^m.$$

Ez az irány menti lineáris megváltozás: $f(x+\varepsilon v)\approx f(x)+\varepsilon\,J_f(x)v$.

Definíció (VJP). Adott $g\in\mathbb R^m$ kimeneti vektor esetén a vector–Jacobi product

$$\mathrm{VJP}_f(x;g)=J_f(x)^\top g\in\mathbb R^d.$$

A megegyezésünk szerint az absztrakt vektorok oszlopvektorok, ezért a VJP-ben a Jacobi-mátrix transzponáltja áll. A mátrixok mérete: $J_f^\top\in\mathbb R^{d\times m}$, $g\in\mathbb R^{m\times1}$, tehát az eredmény $d\times1$.

A két művelet nem szinonima. A JVP egy inputirányt visz előre, a VJP egy outputgradienst visz hátra. Ha $m\ll d$, a JVP lehet olcsóbb irány; a klasszikus backpropagation egy skalár veszteség gradiense miatt viszont VJP-kből épül fel. A kapcsolatot a láncszabály adja: legyen $L:\mathbb R^m\to\mathbb R$ differenciálható, és $h(x)=L(f(x))$. Skalár kimenetű $h$-ra a Jacobi-mátrix $1\times d$ sor, transzponáltja a gradiens:

$$\nabla h(x)=J_f(x)^\top \nabla L(f(x)).$$

Itt $g=\nabla L(f(x))$, tehát $\nabla h(x)=\mathrm{VJP}_f(x;g)$.

Batch és reduction. A tanításban a veszteség általában egy batch-átlag:

$$L=\frac1N\sum_{i=1}^N \ell(z_i,y_i),$$

ahol $z_i=f(x_i)\in\mathbb R^m$ egy megfigyelés kimenete. Legyen $Z\in\mathbb R^{N\times m}$ a batchelt kimenet, és $G=\nabla_Z L\in\mathbb R^{N\times m}$; a $G$ $i$-edik sora $(\nabla_{z_i}L)^\top$. Az átlagolás a láncszabályon kívüli, de a $G$-ben lévő $1/N$ faktor:

$$\frac{\partial L}{\partial z_{i,p}}=\frac1N\frac{\partial \ell(z_i,y_i)}{\partial z_{i,p}}.$$

Ha valaki sum redukcióval számol, akkor $G_{\mathrm{sum,i:}}=(\nabla_{z_i}\ell)^\top$, és $G_{\mathrm{mean}}=(1/N)G_{\mathrm{sum}}$. A fontos szabály: a redukció nem a Jacobi-mátrix transzponálásába épül, hanem a kimeneti gradiens skálájába. Az autograd keretrendszer ezt a $G$ értéket szolgáltatja a hátramenetben, és a réteg feladata már csak az adott $G$-hez tartozó VJP-összeg kiszámítása.

4.3.5 Az affin réteg súly-, bias- és inputgradiense

Most teljesítjük a fejezet fő levezetési kötelezettségét. Legyen $X\in\mathbb R^{N\times d}$ batchelt bemenet, $W\in\mathbb R^{m\times d}$ súly, $b\in\mathbb R^m$ bias, és

$$Z=XW^\top+b,$$

ahol a bias sorirányban broadcast; komponensenként

$$Z_{i,p}=\sum_{q=1}^d X_{i,q}W_{p,q}+b_p.$$

Ez a torch.nn.Linear súlyalakja is: weight mérete (out_features, in_features), a forward pedig x @ weight.T + bias.

Legyen $L$ skalár, és tegyük fel, hogy $L$ mint $Z$ függvénye differenciálható. Jelölje

$$G=\nabla_Z L\in\mathbb R^{N\times m}, \qquad G_{i,r}=\frac{\partial L}{\partial Z_{i,r}}.$$

Az $X$, $W$ és $b$ gradienseit külön-külön vezetjük le.

Inputgradiens. Az $X_{i,q}$ csak az $i$-edik sor $Z_{i,:}$ kimenetén keresztül hat. A láncszabály:

$$\frac{\partial L}{\partial X_{i,q}} = \sum_{p=1}^m\frac{\partial L}{\partial Z_{i,p}}\frac{\partial Z_{i,p}}{\partial X_{i,q}} = \sum_{p=1}^m G_{i,p}W_{p,q} = (GW)_{i,q}.$$

Oszlopvektoros alakban egy megfigyelésre: $\nabla_{x_i}L=W^\top g_i$, ahol $g_i=\nabla_{z_i}L$ oszlopvektor. A teljes batch mátrixalakja:

$$\boxed{\nabla_X L=GW\in\mathbb R^{N\times d}}.$$

Súlygradiens. A $W_{p,q}$ súly minden $i$ sorhoz hozzájárul, ezért a gradiens az $i$ szerinti külső összeg. Mivel

$$\frac{\partial Z_{i,r}}{\partial W_{p,q}}=\delta_{r,p}X_{i,q},$$

a láncszabály:

$$\frac{\partial L}{\partial W_{p,q}} = \sum_{i=1}^N\sum_{r=1}^m \frac{\partial L}{\partial Z_{i,r}} \frac{\partial Z_{i,r}}{\partial W_{p,q}} = \sum_{i=1}^N G_{i,p}X_{i,q} = (G^\top X)_{p,q}.$$

Tehát

$$\boxed{\nabla_W L=G^\top X\in\mathbb R^{m\times d}}.$$

Biasgradiens. A $b_p$ az összes megfigyelés $Z_{i,p}$ eleméhez ugyanazzal az $1$ együtthatóval járul hozzá:

$$\frac{\partial L}{\partial b_p} = \sum_{i=1}^N G_{i,p}.$$

Mátrixos írásmóddal ez a $G$ oszlopösszege:

$$\boxed{(\nabla_b L)_p=\sum_{i=1}^N G_{i,p},\qquad \nabla_b L=G^\top \mathbf 1_N\in\mathbb R^m.}$$

PyTorch/Numpy írásmódban: db = G.sum(axis=0).

A három eredmény egységes konvenciója: az absztrakt oszlopvektoros VJP $W^\top g$, a batch-soros tárolás viszont $GW$. Aki a kettőt keveri, transzponálási hibát kap; erre a tipikus tévedésnél visszatérünk.

Az $L=\frac12\sum_{i,p}Z_{i,p}^2$ speciális esetben $G=Z$, tehát például $\nabla_WL=Z^\top X$; a 3.3. Vetítés, legkisebb négyzetek és kondicionáltság normálegyenletének $X^\top(Xw-y)$ alakja is ugyanezt a VJP-irányt mutatja $f(w)=Xw-y$ és $G=r$ választással.

4.3.6 Kézi kis példa és minimális kód

Legyen

$$f(x)=\begin{bmatrix} x_1\\ x_1+x_2^2\\ x_1x_2 \end{bmatrix}, \qquad x=(1,2)^\top.$$

A parciális deriváltak:

• $\nabla f_1=(1,0)$,
• $\nabla f_2=(1,2x_2)=(1,4)$,
• $\nabla f_3=(x_2,x_1)=(2,1)$,

ezért

$$J_f(1,2)=\begin{bmatrix} 1&0\\ 1&4\\ 2&1 \end{bmatrix}.$$

Válasszuk a kimeneti gradiens értékét $g=(1,-2,3)^\top$-nak. A VJP:

$$J_f^\top g = \begin{bmatrix} 1&1&2\\ 0&4&1 \end{bmatrix} \begin{bmatrix} 1\\-2\\3 \end{bmatrix} = \begin{bmatrix} 1\cdot1+1\cdot(-2)+2\cdot3\\ 0\cdot1+4\cdot(-2)+1\cdot3 \end{bmatrix} = \begin{bmatrix} 5\\-5 \end{bmatrix}.$$

A következő kód ugyanezt explicit Jacobi-mátrixszal és PyTorch-autograddal is ellenőrzi. A torch.dot(z,g) skalár veszteséget képez, tehát az autograd a láncszabály szerinti $\nabla_x(g^\top f(x))=J_f^\top g$ VJP-t adja vissza.

# lab/n015/vjp_lab.py
import torch

torch.manual_seed(20261020)  # a példa determinisztikus; véletlen nincs benne
DTYPE = torch.float64
DEVICE = torch.device("cpu")

x = torch.tensor([1.0, 2.0], dtype=DTYPE, device=DEVICE, requires_grad=True)
g = torch.tensor([1.0, -2.0, 3.0], dtype=DTYPE, device=DEVICE)

def f(u: torch.Tensor) -> torch.Tensor:
    return torch.stack([
        u[0],
        u[0] + u[1] ** 2,
        u[0] * u[1],
    ])

z = f(x)

# Az x = (1,2) pontban kézzel levezetett Jacobi-mátrix:
J = torch.tensor(
    [
        [1.0, 0.0],
        [1.0, 2.0 * x[1].item()],
        [x[1].item(), x[0].item()],
    ],
    dtype=DTYPE,
    device=DEVICE,
)

vjp_explicit = J.T @ g
# loss = g·f(x), ezért grad_x(loss) = J_f(x)^T g:
loss = torch.dot(z, g)
vjp_autograd, = torch.autograd.grad(loss, x)

print("J =")
print(J)
print("vjp_explicit =", vjp_explicit)
print("vjp_autograd =", vjp_autograd)
print("max_abs_diff =", (vjp_explicit - vjp_autograd).abs().max().item())

Futtatás a lab/n015 munkakönyvtárból:

python lab/n015/vjp_lab.py

Várt kimenet (a fenti kézi számításból, nem a jelen írás során futtatott naplóból):

J =
tensor([[1., 0.],
        [1., 4.],
        [2., 1.]], dtype=torch.float64)
vjp_explicit = tensor([5., -5.], dtype=torch.float64)
vjp_autograd = tensor([5., -5.], dtype=torch.float64)
max_abs_diff = 0.0

4.3.7 Labor: explicit Jacobian és autograd-VJP összehasonlítása

A labor kódja a fenti lab/n015/vjp_lab.py.

• Input: x=(1,2) és g=(1,-2,3), mindkettő float64; nincs külső adatfájl, a GlyphLab itt nem szükséges.
• Import: torch; a labor a CPU-autograd utat használja, device="cpu".
• Seed/dtype/device: torch.manual_seed(20261020), torch.float64, CPU. A seed itt csak a protokoll része, mert a kód nem használ véletlen számot.
• Adatút: forrásfájl lab/n015/vjp_lab.py, futtatási parancs fent.
• Ellenőrzés: a vjp_explicit a kézzel felírt $J_f^\top g$, a vjp_autograd a $g\cdot f(x)$ skalár veszteség gradiense. Alaktolerancia: mindkettő (2,); f64 esetén az eltérésnek nulla közeli értéknek kell lennie, elfogadott numerikus korlát atol=1e-12, rtol=1e-12. A kiírt max_abs_diff a várható példán 0.
• Értelmezés: az egyezés azt mutatja, hogy a skalár veszteség hátramenete nem a $3\times2$ Jacobi-mátrixot tárolja, hanem annak transzponáltjával szorozva a $g$ vektort. A $J_f^\top g$ irány tehát az, amit a backprop VJP-ként használ.

A megadott kimenet számított, nem mért laborjelentés: a kódot a tanulólaborban külön kell futtatni, és az eltérést a fenti toleranciával rögzíteni.

4.3.8 Feladatok

F1, fogalmi, nehézség 1. Legyen $W=\begin{bmatrix}1&2&3\\4&5&6\end{bmatrix}$, $b=0$, $f(x)=Wx+b$. Adja meg $J_f$ alakját, a JVP-t $v=(1,0,-1)^\top$, és a VJP-t $g=(1,-1)^\top$ esetén. Ne feledje megnevezni, melyik művelet melyik Jacobi-irányt használja.

F2, levezetési, nehézség 2. Legyen

$$f(x_1,x_2,x_3)=\begin{bmatrix}x_1+x_2x_3\\x_2^2+x_3\end{bmatrix}, \qquad x_0=(1,2,3)^\top.$$

Számítsa ki $J_f(x_0)$ mátrixot és a $g=(-1,2)^\top$ vektorhoz tartozó VJP-t. Minden lépést indokoljon, és ellenőrizze az alakokat.

F3, implementációs, nehézség 2. Írjon NumPy-függvényt

def affine_grads(X: np.ndarray, W: np.ndarray, b: np.ndarray, G: np.ndarray):
    ...

amely float64 bemenetekkel visszaadja a dX, dW, db hármast az $X\in\mathbb R^{N\times d}$, $W\in\mathbb R^{m\times d}$, $b\in\mathbb R^m$, $G\in\mathbb R^{N\times m}$ konvencióban. Tesztelje az $N=2,d=3,m=2$ esetet: $X=[[1,0,2],[0,1,1]]$, $W=[[1,2,0],[0,1,2]]$, $b=[-1,1]$, $G=[[2,-1],[0,-3]]$. Adja meg a három elvárt mátrixot/vektort.

F4, hibakeresési, nehézség 1. Az alábbi kód a batchelt affin réteg inputgradiensét szeretné visszaadni $G=\nabla_ZL$ ismeretében:

def buggy_dX(X, W, G):
    # X: N x d, W: m x d, G: N x m
    return W.T @ G

Keresse meg a hibát, magyarázza el az okát, következményét, és javítsa a kódot. Ellenőrizze a javítást az F3 numerikus példáján.

F5, kísérleti/transzfer, nehézség 3. Tervezzen CPU-s PyTorch-mérést annak eldöntésére, hogy az $f:\mathbb R^{1000}\to\mathbb R^{1000}$, $f(x)=W\sin(x)+b$ leképezésnél a teljes Jacobi-mátrix anyagiasítása vagy az autograd-VJP használ-e kisebb csúcsmemóriát és időt. Nincs szükség tényleges futtatásra; adjon protokollt, baseline-t, változó tényezőt, metrikát, ellenőrző egyezést és érvényességi korlátot. Jelölje meg, mi lenne a nyers eredménytábla két sora.

4.3.9 Megoldások és indoklás

F1 megoldás. $f$ affin, ezért a Jacobi-mátrix konstans: $J_f=W\in\mathbb R^{2\times3}$. Sor–oszlop konvenció: az 1. sor az $f_1=x_1+2x_2+3x_3$, a 2. sor az $f_2=4x_1+5x_2+6x_3$ parciális deriváltjait tartalmazza.

JVP: $J_fv=Wv=(1\cdot1+2\cdot0+3\cdot(-1),\,4\cdot1+5\cdot0+6\cdot(-1))=(-2,-2)^\top$, mérete $2\times1$.

VJP: $J_f^\top g=W^\top g=(1\cdot1+4\cdot(-1),\,2\cdot1+5\cdot(-1),\,3\cdot1+6\cdot(-1))=(-3,-3,-3)^\top$, mérete $3\times1$. Az első a Jacobi-mátrixszal, a második annak transzponáltjával szoroz.

F2 megoldás. Parciális deriváltak:

$$\frac{\partial f_1}{\partial x_1}=1,\quad \frac{\partial f_1}{\partial x_2}=x_3,\quad \frac{\partial f_1}{\partial x_3}=x_2,$$

$$\frac{\partial f_2}{\partial x_1}=0,\quad \frac{\partial f_2}{\partial x_2}=2x_2,\quad \frac{\partial f_2}{\partial x_3}=1.$$

Ezért a $2\times3$ Jacobi-mátrix:

$$J_f(x)=\begin{bmatrix} 1&x_3&x_2\\ 0&2x_2&1 \end{bmatrix}, \qquad J_f(1,2,3)=\begin{bmatrix} 1&3&2\\ 0&4&1 \end{bmatrix}.$$

A VJP:

$$J_f(1,2,3)^\top g = \begin{bmatrix} 1&0\\ 3&4\\ 2&1 \end{bmatrix} \begin{bmatrix} -1\\2 \end{bmatrix} = \begin{bmatrix} -1\\5\\0 \end{bmatrix}.$$

Alakellenőrzés: $J_f^\top$ mérete $3\times2$, $g$ mérete $2\times1$, az eredmény $3\times1$, valóban az inputtér gradiense.

F3 megoldás. A szükséges kód:

# lab/n015/f3_solution.py
import numpy as np

def affine_grads(X, W, b, G):
    assert X.dtype == np.float64 and W.dtype == np.float64
    assert G.dtype == np.float64 and b.dtype == np.float64
    dX = G @ W
    dW = G.T @ X
    db = G.sum(axis=0)
    return dX, dW, db

X = np.array([[1.0, 0.0, 2.0],
              [0.0, 1.0, 1.0]], dtype=np.float64)
W = np.array([[1.0, 2.0, 0.0],
              [0.0, 1.0, 2.0]], dtype=np.float64)
b = np.array([-1.0, 1.0], dtype=np.float64)
G = np.array([[2.0, -1.0],
              [0.0, -3.0]], dtype=np.float64)

dX, dW, db = affine_grads(X, W, b, G)
print(dX)
print(dW)
print(db)

Kézi mátrixszorzással:

$dX=GW$, ahol $G$ mérete $2\times2$, $W$ mérete $2\times3$, tehát $dX\in\mathbb R^{2\times3}$:

$$dX=\begin{bmatrix}2&3&-2\\0&-3&-6\end{bmatrix}.$$

$dW=G^\top X$, mérete $2\times3$:

$$dW=\begin{bmatrix}2&0&4\\-1&-3&-5\end{bmatrix}.$$

A biasgradiens a $G$ oszlopösszege:

$$db=(2+0,\,-1-3)=(2,-4).$$

A kód a b értékét a gradiensben nem használja, ami helyes: a biasgradiens független a bias aktuális értékétől. Szélsőeset-ellenőrzés: $G=0$ esetén mindhárom gradiens nulla, a függvény (N×d, m×d, m) alakkal tér vissza; $N=1$ esetén db=G[0,:].

F4 megoldás. A hiba oka: a kód a batch-soros G@W irány helyett W.T@G-t ír. Az absztrakt oszlopvektoros szabály valóban $W^\top g_i$ egy megfigyelésre, de a batchelt soros tárolásban a transzponálás után a teljes mátrixszorzat G@W; a $G$ sorai a $g_i^\top$ vektorok. A W.T@G általában alakhibát ad, mert (d×m) @ (N×m) nem szorozható össze, kivéve a véletlenül egybeeső dimenziós eseteket. Következmény: futásidejű broadcast/shape hiba, vagy ha a méretek egybeesnek, akkor numerikusan rossz gradiens, amely a súlyfrissítést elrontja.

Javítás:

def buggy_dX_corrected(X, W, G):
    return G @ W  # N x d

Ellenőrzés az F3 példáján: G@W a fenti dX értéket adja, és az np.testing.assert_allclose f64 alatt atol=1e-12 határral elfogadja.

F5 megoldásminta. A kérdés erőforrás-típusú transzfer: a teljes Jacobi-mátrix $1000\times1000$ f64 elem, azaz $8\cdot10^6$ bájt, nagyságrendileg 8 MiB csak a mátrixra; a VJP egyetlen $1000$ hosszú vektort ad vissza. A terv:

• Kérdés: azonos matematikai eredményt feltételezve, a Jacobi-anyagiasítás vagy a VJP kisebb csúcsmemóriájú-e CPU-n.

• Baseline: J = torch.autograd.functional.jacobian(f, x), majd v1 = J.T @ g; ez anyagiasítja a teljes mátrixot.

• Változat: y = f(x), loss = torch.dot(y,g), v2, = torch.autograd.grad(loss,x).

• Fix keret: dtype=float64, CPU, torch.manual_seed(20261020), $W$ véletlen $1000\times1000$ mátrix, $x$ és $g$ rögzített vektor; warmup háromszor, mérés ötvenszer.

• Metrika: medián és interkvartilis terjedelem a falidőből, valamint a folyamat csúcsmemóriája (ru_maxrss, Linuxon KiB-ban); minden ismétlés előtt új folyamat a nyomkövető overhead elkerülésére. A csúcsmemória lekérdezése:

  import resource
  peak_kib = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss

• Kontroll: torch.allclose(v1,v2,rtol=1e-12,atol=1e-12) igaz legyen, és mindkét output alakja (1000,).

• Értelmezési korlát: az eredmény erre a toy leképezésre, e framework-verzióra és CPU-ra érvényes; a mérés overheadje, a kernelindítási küszöb és a $d=m$ szimmetria miatt nem általános GPU-, batch- vagy termelési állítás. A nyers táblában a két módszer wall_time_median, wall_time_iqr, ru_maxrss_kib értéke, valamint a kontroll allclose eredménye szerepel.

4.3.10 Tipikus tévedés, önellenőrzés és visszahívás

A leggyakoribb hiba a transzponálási konvenció elhallgatása. Legyen $f(x)=Wx$, ahol

$$W=\begin{bmatrix}1&2\\3&4\end{bmatrix},\qquad g=(1,0)^\top.$$

Ekkor $J_f=W$, tehát a helyes VJP

$$J_f^\top g=W^\top g=\begin{bmatrix}1&3\\2&4\end{bmatrix}\begin{bmatrix}1\\0\end{bmatrix}=(1,2)^\top.$$

Ha valaki tévesen $J_fg=Wg$-t számít, az $(1,3)^\top$-t kap. A példa azért fontos, mert $m=d=2$, tehát az alak hiba nélkül is elfedi a tévedést: mindkét eredmény két hosszú vektor, mégis a második a JVP, nem a VJP. Batch-soros tárolásnál a fenti affin szabály szerint dX=G@W, nem W.T@G; a kettő egyetlen megfigyelés oszlopgradiensére átszámítva $g_i^\top W=(\nabla_{x_i}L)^\top$.

Önellenőrzés ehhez a fejezethez: minden láncszabály-szorzatnál a szomszédos Jacobi-méreteknek illeszkedniük kell ($k\times m$ és $m\times d$); a VJP bemenete $m$-hosszú, kimenete $d$-hosszú; a mean-redukció $1/N$-je a $G$-ben van, nem a $W^\top$-ben. Ha a kód olyan gradiensvektort ad, amelynek hossza $m$ a kívánt $d$ helyett, akkor szinte biztosan JVP-t számoltunk VJP helyett.

Visszahívás: a 4.2. Parciális derivált és gradiens alfejezetben a gradiens oszlopvektor volt; itt ezért lett a láncszabályból $J_f^\top g$. A 3.3. Vetítés, legkisebb négyzetek és kondicionáltság legkisebb négyzetes gradiense, $X^\top(Xw-y)$, ugyanennek a VJP-formának a speciális esete: $f(w)=Xw-y$, $J_f=X$, és a kimeneti gradiens az $r=Xw-y$ reziduum.

4.3.11 Rövid összefoglaló

A Jacobi-mátrix az $f:\mathbb R^d\to\mathbb R^m$ deriváltjának $m\times d$ mátrixa, sora az output, oszlopa az input. Differenciálható kompozícióban a Jacobikat a $J_gJ_f$ sorrendben szorozzuk. A backprop nem a teljes Jacobit építi fel, hanem a VJP-t: $J_f^\top g$, ahol $g=\nabla_{f(x)}L$. Az affin rétegre $Z=XW^\top+b$ esetén a gradiensek: $\nabla_XL=GW$, $\nabla_WL=G^\top X$, és $\nabla_bL$ a $G$ oszlopösszege. A mean-redukció $1/N$ szorzója a $G$-ben van, nem a transzponálási szabályban. A kidolgozott f64 példa várt VJP-értéke $(5,-5)^\top$, és az autograd ellenőrzés ugyanezt a $J^\top g$ irányt teszteli.

4.3.12 Következő alkalmazás

A 4.4. I. ellenőrző pont: matematika és kód megfelelése ellenőrző pont a matematika és kód megfelelését vizsgálja. Ott a most levezetett affin gradiensképleteket, a JVP/VJP irányt és a numerikus toleranciát egy összetettebb, kis dimenziós ellenőrzési feladaton kell együtt használni.

4.3.13 Források

• [B01] Deisenroth, Faisal és Ong (2020). Mathematics for Machine Learning. Cambridge University Press. https://mml-book.github.io/.
• [B02] Goodfellow, Bengio és Courville (2016). Deep Learning. MIT Press. https://www.deeplearningbook.org/.
• [D05] PyTorch contributors (2.14). Autograd mechanics. https://docs.pytorch.org/docs/2.14/notes/autograd.html.