2.5 Automatinis diferencijavimas

Prisiminkite iš 2.4 skyrelis, kad išvestinių skaičiavimas yra lemiamas žingsnis visuose optimizavimo algoritmuose, kuriuos naudosime giliesiems tinklams apmokyti. Nors skaičiavimai paprasti, atlikti juos rankomis gali būti nuobodu ir lengva suklysti, o šios problemos tik didėja, kai mūsų modeliai tampa sudėtingesni.

Laimei, visi šiuolaikiniai giliojo mokymosi karkasai šį darbą perima iš mūsų, siūlydami automatinį diferencijavimą (dažnai sutrumpintai autograd). Kai duomenis perduodame per kiekvieną paskesnę funkciją, karkasas sukuria skaičiavimo grafą, kuris seka, kaip kiekviena reikšmė priklauso nuo kitų. Išvestinėms apskaičiuoti automatinis diferencijavimas eina šiuo grafu atgal, taikydamas grandinės taisyklę. Skaičiavimo algoritmas, taikantis grandinės taisyklę tokiu būdu, vadinamas atgaliniu sklidimu.

Nors autograd bibliotekos pastarąjį dešimtmetį tapo karštu rūpesčiu, jos turi ilgą istoriją. Tiesą sakant, ankstyviausios nuorodos į autograd siekia daugiau nei pusę amžiaus (Wengert, 1964). Pagrindinės šiuolaikinio atgalinio sklidimo idėjos siekia 1980 m. daktaro disertaciją (Speelpenning, 1980) ir buvo toliau plėtojamos devintojo dešimtmečio pabaigoje (Griewank, 1989). Nors atgalinis sklidimas tapo numatytuoju gradientų skaičiavimo metodu, jis nėra vienintelis pasirinkimas. Pavyzdžiui, programavimo kalba Julia naudoja sklidimą pirmyn (Revels ir kt., 2016). Prieš nagrinėdami metodus, pirmiausia įvaldykime autograd paketą.

from mxnet import autograd, np, npx
npx.set_np()

import torch

import tensorflow as tf

from jax import numpy as jnp

2.5.1 Paprasta funkcija

Tarkime, kad mus domina (funkcijos $y = 2\mathbf{x}^{\top}\mathbf{x}$ diferencijavimas stulpelio vektoriaus $\mathbf{x}$ atžvilgiu.) Pradžioje priskiriame x pradinę reikšmę.

x = np.arange(4.0)
x

x = torch.arange(4.0)
x

x = tf.range(4, dtype=tf.float32)
x

x = jnp.arange(4.0)
x

# We allocate memory for a tensor's gradient by invoking `attach_grad`
x.attach_grad()
# After we calculate a gradient taken with respect to `x`, we will be able to
# access it via the `grad` attribute, whose values are initialized with 0s
x.grad

# Can also create x = torch.arange(4.0, requires_grad=True)
x.requires_grad_(True)
x.grad  # The gradient is None by default

x = tf.Variable(x)

(Dabar apskaičiuojame savo funkciją nuo x ir rezultatą priskiriame y.)

# Our code is inside an `autograd.record` scope to build the computational
# graph
with autograd.record():
    y = 2 * np.dot(x, x)
y

y = 2 * torch.dot(x, x)
y

# Record all computations onto a tape
with tf.GradientTape() as t:
    y = 2 * tf.tensordot(x, x, axes=1)
y

y = lambda x: 2 * jnp.dot(x, x)
y(x)

y.backward()
x.grad

y.backward()
x.grad

x_grad = t.gradient(y, x)
x_grad

from jax import grad
# The `grad` transform returns a Python function that
# computes the gradient of the original function
x_grad = grad(y)(x)
x_grad

(Jau žinome, kad funkcijos $y = 2\mathbf{x}^{\top}\mathbf{x}$ gradientas $\mathbf{x}$ atžvilgiu turėtų būti $4\mathbf{x}$.) Dabar galime patikrinti, ar automatinis gradiento skaičiavimas ir laukiamas rezultatas sutampa.

x.grad == 4 * x

x.grad == 4 * x

x_grad == 4 * x

x_grad == 4 * x

with autograd.record():
    y = x.sum()
y.backward()
x.grad  # Overwritten by the newly calculated gradient

x.grad.zero_()  # Reset the gradient
y = x.sum()
y.backward()
x.grad

with tf.GradientTape() as t:
    y = tf.reduce_sum(x)
t.gradient(y, x)  # Overwritten by the newly calculated gradient

y = lambda x: x.sum()
grad(y)(x)

2.5.2 Atgalinis sklidimas neskaliariniams kintamiesiems

Kai y yra vektorius, natūraliausias y išvestinės vektoriaus x atžvilgiu vaizdavimas yra matrica, vadinama Jakobio matrica, kurioje yra kiekvienos y komponentės dalinės išvestinės kiekvienos x komponentės atžvilgiu. Panašiai aukštesnės eilės y ir x atveju diferencijavimo rezultatas gali būti dar aukštesnės eilės tenzorius.

Nors Jakobio matricos pasirodo kai kuriuose pažangiuose mašininio mokymosi metoduose, dažniau norime susumuoti kiekvienos y komponentės gradientus pagal visą vektorių x, gaudami vektorių, kurio forma tokia pati kaip x. Pavyzdžiui, dažnai turime vektorių, vaizduojantį nuostolių funkcijos reikšmę, apskaičiuotą atskirai kiekvienam pavyzdžiui apmokymo pavyzdžių minirinkinyje. Čia tiesiog norime (susumuoti gradientus, apskaičiuotus atskirai kiekvienam pavyzdžiui).

with autograd.record():
    y = x * x  
y.backward()
x.grad  # Equals the gradient of y = sum(x * x)

x.grad.zero_()
y = x * x
y.backward(gradient=torch.ones(len(y)))  # Faster: y.sum().backward()
x.grad

with tf.GradientTape() as t:
    y = x * x
t.gradient(y, x)  # Same as y = tf.reduce_sum(x * x)

y = lambda x: x * x
# grad is only defined for scalar output functions
grad(lambda x: y(x).sum())(x)

2.5.3 Skaičiavimo atsiejimas

Kartais norime kai kuriuos skaičiavimus perkelti už įrašyto skaičiavimo grafo ribų. Pavyzdžiui, tarkime, naudojame įvestį sukurti kai kuriems pagalbiniams tarpiniams nariams, kurių gradiento skaičiuoti nenorime. Šiuo atveju turime atsieti atitinkamą skaičiavimo grafą nuo galutinio rezultato. Šis žaislinis pavyzdys tai paaiškina: tarkime, turime z = x * y ir y = x * x, bet norime sutelkti dėmesį į tiesioginę x įtaką z, o ne per y perduodamą įtaką. Šiuo atveju galime sukurti naują kintamąjį u, turintį tą pačią reikšmę kaip y, bet kurio kilmė (kaip jis buvo sukurtas) ištrinta. Taigi u grafe neturi protėvių, ir gradientai neteka per u į x. Pavyzdžiui, apskaičiavus z = x * u gradientą, gausime rezultatą u (o ne 3 * x * x, kaip galėjote tikėtis, nes z = x * x * x).

with autograd.record():
    y = x * x
    u = y.detach()
    z = u * x
z.backward()
x.grad == u

x.grad.zero_()
y = x * x
u = y.detach()
z = u * x

z.sum().backward()
x.grad == u

# Set persistent=True to preserve the compute graph. 
# This lets us run t.gradient more than once
with tf.GradientTape(persistent=True) as t:
    y = x * x
    u = tf.stop_gradient(y)
    z = u * x

x_grad = t.gradient(z, x)
x_grad == u

import jax

y = lambda x: x * x
# jax.lax primitives are Python wrappers around XLA operations
u = jax.lax.stop_gradient(y(x))
z = lambda x: u * x

grad(lambda x: z(x).sum())(x) == y(x)

Atkreipkite dėmesį, kad nors ši procedūra atsieja y protėvius nuo grafo, vedančio į z, skaičiavimo grafas, vedantis į y, išlieka, todėl galime apskaičiuoti y gradientą x atžvilgiu.

y.backward()
x.grad == 2 * x

x.grad.zero_()
y.sum().backward()
x.grad == 2 * x

t.gradient(y, x) == 2 * x

grad(lambda x: y(x).sum())(x) == 2 * x

2.5.4 Gradientai ir Python valdymo srautas

Iki šiol apžvelgėme atvejus, kai kelias nuo įvesties iki išvesties buvo aiškiai apibrėžtas funkcija, pvz., z = x * x * x. Programavimas suteikia daug daugiau laisvės, kaip skaičiuojame rezultatus. Pavyzdžiui, galime padaryti, kad jie priklausytų nuo pagalbinių kintamųjų arba sąlygotų pasirinkimus tarpiniais rezultatais. Vienas automatinio diferencijavimo privalumas yra tas, kad net jei skaičiavimo grafo kūrimui (funkcijos reikėjo praeiti Python valdymo srauto labirintą) (pvz., sąlygas, ciklus ir savavališkus funkcijų iškvietimus), (vis tiek galime apskaičiuoti gauto kintamojo gradientą.) Tam iliustruoti panagrinėkime šį kodo fragmentą, kuriame while ciklo iteracijų skaičius ir if sakinio įvertinimas abu priklauso nuo įvesties a reikšmės.

def f(a):
    b = a * 2
    while np.linalg.norm(b) < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c

def f(a):
    b = a * 2
    while b.norm() < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c

def f(a):
    b = a * 2
    while tf.norm(b) < 1000:
        b = b * 2
    if tf.reduce_sum(b) > 0:
        c = b
    else:
        c = 100 * b
    return c

def f(a):
    b = a * 2
    while jnp.linalg.norm(b) < 1000:
        b = b * 2
    if b.sum() > 0:
        c = b
    else:
        c = 100 * b
    return c

Toliau iškviečiame šią funkciją, kaip įvestį perduodami atsitiktinę reikšmę. Kadangi įvestis yra atsitiktinis dydis, nežinome, kokią formą įgis skaičiavimo grafas. Tačiau kai tik vykdome f(a) su konkrečia įvestimi, realizuojame konkretų skaičiavimo grafą ir vėliau galime paleisti backward.

a = np.random.normal()
a.attach_grad()
with autograd.record():
    d = f(a)
d.backward()

a = torch.randn(size=(), requires_grad=True)
d = f(a)
d.backward()

a = tf.Variable(tf.random.normal(shape=()))
with tf.GradientTape() as t:
    d = f(a)
d_grad = t.gradient(d, a)
d_grad

from jax import random
a = random.normal(random.PRNGKey(1), ())
d = f(a)
d_grad = grad(f)(a)

Nors mūsų funkcija f demonstravimo tikslais yra kiek dirbtinė, jos priklausomybė nuo įvesties gana paprasta: tai tiesinė a funkcija su dalimis apibrėžtu mastelio koeficientu. Todėl f(a) / a yra pastovių įrašų vektorius, ir, be to, f(a) / a turi sutapti su f(a) gradientu a atžvilgiu.

a.grad == d / a

a.grad == d / a

d_grad == d / a

d_grad == d / a

Dinaminis valdymo srautas giliajame mokyme labai dažnas. Pavyzdžiui, apdorojant tekstą, skaičiavimo grafas priklauso nuo įvesties ilgio. Šiais atvejais automatinis diferencijavimas tampa gyvybiškai svarbus statistiniam modeliavimui, nes gradiento neįmanoma apskaičiuoti a priori.

2.5.5 Diskusija

Dabar paragavote automatinio diferencijavimo galios. Bibliotekų, skaičiuojančių išvestines ir automatiškai, ir efektyviai, kūrimas labai padidino giliojo mokymosi praktikų produktyvumą, išlaisvindamas juos, kad galėtų sutelkti dėmesį į mažiau mechaniškus darbus. Be to, autograd leidžia projektuoti didžiulius modelius, kuriems gradientų skaičiavimas pieštuku ir popieriumi užimtų neįmanomai daug laiko. Įdomu tai, kad nors autograd naudojame modeliams optimizuoti (statistine prasme), pačių autograd bibliotekų optimizavimas (skaičiavimo prasme) yra turtinga tema, labai svarbi karkasų kūrėjams. Čia kompiliatorių ir grafų manipuliavimo priemonės naudojamos rezultatams apskaičiuoti patogiausiu ir atmintį tausojančiausiu būdu.

Kol kas pabandykite prisiminti šiuos pagrindus: (i) prie tų kintamųjų, kurių atžvilgiu norime išvestinių, pridėkite gradientus; (ii) įrašykite tikslinės reikšmės skaičiavimą; (iii) vykdykite atgalinio sklidimo funkciją; ir (iv) pasiekite gautą gradientą.

2.5.6 Pratimai

1. Kodėl antrąją išvestinę apskaičiuoti daug brangiau nei pirmąją?
2. Paleidę atgalinio sklidimo funkciją, nedelsdami paleiskite ją dar kartą ir pažiūrėkite, kas nutinka. Ištirkite.
3. Valdymo srauto pavyzdyje, kuriame skaičiuojame d išvestinę a atžvilgiu, kas nutiktų, jei kintamąjį a pakeistume atsitiktiniu vektoriumi ar matrica? Tuomet skaičiavimo f(a) rezultatas nebėra skaliaras. Kas nutinka rezultatui? Kaip tai išanalizuoti?
4. Tegul $f(x) = \sin(x)$. Nubraižykite $f$ ir jos išvestinės $f'$ grafikus. Nenaudokite fakto, kad $f'(x) = \cos(x)$, o rezultatą gaukite naudodami automatinį diferencijavimą.
5. Tegul $f(x) = ((\log x^2) \cdot \sin x) + x^{-1}$. Užrašykite priklausomybės grafą, sekantį rezultatus nuo $x$ iki $f(x)$.
6. Naudodami grandinės taisyklę apskaičiuokite minėtos funkcijos išvestinę $\frac{df}{dx}$, kiekvieną narį išdėstydami anksčiau sukonstruotame priklausomybės grafe.
7. Turint grafą ir tarpinius išvestinių rezultatus, skaičiuojant gradientą yra keletas variantų. Įvertinkite rezultatą vieną kartą pradėdami nuo $x$ iki $f$ ir vieną kartą nuo $f$ sekdami atgal iki $x$. Kelias nuo $x$ iki $f$ paprastai vadinamas tiesioginiu diferencijavimu (angl. forward differentiation), o kelias nuo $f$ iki $x$ – atgaliniu diferencijavimu.
8. Kada galėtumėte norėti naudoti tiesioginį, o kada atgalinį diferencijavimą? Užuomina: apsvarstykite reikalingų tarpinių duomenų kiekį, galimybę lygiagretinti žingsnius ir susijusių matricų bei vektorių dydį.