Automatic Differentiation

Giới thiệu

Trong nhiều bài toán machine learning, chúng ta thường có một hàm để đánh giá xem model của chúng ta đưa ra kết quả chính xác đến mức nào. Hàm đó thường là một hàm đạo hàm được và có output là một số thực.

Một model trong một bài toán machine learning sẽ bao gồm các parameters (tôi được biết chúng còn được gọi là "weights", nhưng không chắc lắm), với mục đích cố gắng giải thích bộ data có sẵn theo một công thức nào đó. Một hàm mất mát (loss function) sẽ cho chúng ta biết model của chúng ta tốt hay kém. Loss function có output thường có output là một số thực và các parameters nói trên là independent variables. Loss function có thể sử dụng bộ data để tính toán, nhưng bộ data này chỉ được xem là hằng số.

Khi train một model, chúng ta sẽ thay đổi các parameters sao cho model chúng ta fit bộ data đó càng khớp càng tốt, i.e. loss function càng nhỏ càng tốt. Nếu loss function là hàm đạo hàm được, với một giá trị cụ thể của bộ parameters, ta có thể tính được xem loss function tại giá trị đó đang có xu hướng tăng hay giảm và tốc độ thay đổi là bao nhiêu. Đối với một loss function đơn giản, ta có thể tính tay công thức gradient của nó và khi cần tính gradient tại điểm nào thì thay điểm đó vào công thức gradient là ra. Tuy nhiên nếu loss function phức tạp, việc tính toán trở nên cực kỳ khó khăn.

autograd là một module trong mxnet giúp chúng ta tự động hóa việc tính đạo hàm bằng cách "theo dõi" các phép tính được thực hiện và dựng nên một computational graph. Giả sử ta có một hàm , để tính , ta chỉ cần gọi

f.backward()
print(x.grad)

Lưu ý là (cũng như ) là một vector, i.e. tập hợp nhiều independent variable khác. Khi tính partial derivative của tương ứng với , xin hiểu rằng ta đang tính partial derivative của và gom các kết quả đó lại thành một vector.

Ví dụ mở đầu

Hàm attach_grad thông báo cho thư viện đâu là independent variable, để tính gradient dựa trên đó.

from mxnet import autograd, nd

x = nd.arange(4)
x

[ 0. 1. 2. 3.]

x.attach_grad()

Sau khi đã "khai báo" independent variable, ta có thể bắt đầu sử dụng nó để tính toán. Quá trình tính toán phải được nằm trong khối with autograd.record():

Ví dụ ta có bốn biến là independent variable và hàm

with autograd.record():
    y = 2 * nd.dot(x, x)
y.backward() # Tìm gradient của y
y

[ 28.]

thì gradient của hàm trên là:

x.grad

[ 0. 4. 8. 12.]

4*x

[ 0. 4. 8. 12.]

Nếu ta dùng x để tính trong một hàm khác khi autograd.record() được bật thì x.grad sẽ không giữ giá trị cũ nữa.

Ví dụ, xét hàm:

with autograd.record():
    z = nd.dot(3 * x * x * x, nd.array([1, 1, 1, 1]))
z.backward() # Tìm gradient của z
z

[ 108.]

thì gradient của hàm trên là:

x.grad

[ 0. 9. 36. 81.]

9*x*x

[ 0. 9. 36. 81.]

Backward đối với non-scalar dependent variable

Nếu gọi z.backward() mà z không phải là scalar thì mxnet nó sẽ cộng tất cả các phần tử của z lại (việc cộng này được coi là thực hiện khi autograd.record() được bật) rồi mới tính backward.

Ví dụ, ở đoạn code trên, tôi viết

with autograd.record():
    z = nd.dot(3 * x * x * x, nd.array([1, 1, 1, 1]))

cốt là để cộng các phần tử của 3*x*x*x lại với nhau.

3*x*x*x

[ 0. 3. 24. 81.]

nd.dot(3*x*x*x, nd.array([1,1,1,1]))

[ 108.]

Với tính chất của hàm backward nói trên, ta có thể đơn giản hóa code trên một tí mà vẫn giữ nguyên kết quả:

with autograd.record():
    z = 3 * x * x * x

with autograd.record():
    z = 3 * x * x * x
z.backward() # Tìm gradient của z
z

[ 0. 3. 24. 81.]

x.grad

[ 0. 9. 36. 81.]

(x.grad trong hai cách làm đều bằng nhau)

Detach

"Automatic differentation" có nghĩa là theo dõi sự thay đổi của các dependent variable (ví dụ y, z), ghi nhận các phép tính thực thi trên các independent variable (ví dụ x), để suy ra gradient, hay nói đúng hơn là suy ra partial derivative của một số independent variable tại một điểm nhất định. Nhắc lại rằng independent variable là các biến mà chúng ta đã gọi hàm attach_grad().

"Detach" là việc tách một giá trị trung gian ra khỏi computational graph và xem giá trị đó như là một hằng số hoàn toàn mới, xóa bỏ những "ký ức" về quá trình tính ra giá trị đó.

Ví dụ 1, tính gradient của , ta có kết quả là:

with autograd.record():
    y = 3*x*x
    z = y*x
z.backward()
x.grad.asnumpy()

array([ 0., 9., 36., 81.], dtype=float32)

Ví dụ 2, với ta được

with autograd.record():
    y = 3*x*x
    z = y*x
y.backward()
x.grad.asnumpy()

array([ 0., 6., 12., 18.], dtype=float32)

Ví dụ 3, với ta được

Nhận thấy rằng . Mặt khác , nên ta có thể tính 3*x*x trước, rồi xem nó như một hằng số bằng hàm detach(), và tính tiếp giá trị của .

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

[ 0. 3. 12. 27.]

Trong đoạn code trên, t được gán cho giá trị đúng bằng y, nhưng t không hề có quá khứ. Tức là t "không biết" nó đã được tính như thế nào. y thì vẫn "nhớ" là nó bằng 3*x*x. Vậy nên khi tính z = t*x, nó không truy vết được 3*x*x mà chỉ đơn thuần biết giá trị của t như thể t được tính sẵn trước khi vào autograd.record().

Lưu ý là dòng code này

t = y.detach()

không hề xóa lịch sử của y. Nó chỉ đơn thuần copy giá trị của y đưa cho t mà không cho t biết lịch sử của y.

Để xóa lịch sử của y, ta có thể không cần tới một biến t nào khác, có thể làm thẳng như sau:

y = y.detach()

và câu lệnh trên tương đương với lệnh sau:

y.attach_grad()

Có hai nhận xét cho câu lệnh y.attach_grad() nói trên. Một là, nếu y chưa phải là một independent variable, câu lệnh này khiến nó trở thành một independent variable với giá trị hoàn toàn mới. Hai là, nếu y vốn là một independent variable, câu lệnh này đơn thuần "refresh" y, xóa đi các phép tính đã hình thành nên y.

with autograd.record():
    y = 3*x*x
    y.attach_grad()
    z = y*x
z.backward()
x.grad

[ 0. 3. 12. 27.]

Một lưu ý nhỏ là khi bạn gọi var.grad hay var.backward(), phải đảm bảo rằng var có một chỗ đứng trong computational graph, nếu không sẽ chẳng có kết quả gì, hoặc thậm chí gây ra lỗi.

Ví dụ, đoạn code dưới đây

with autograd.record():
    y = 3*x*x
    t = y.detach()
    z = t*x

cho thấy t không hề có chỗ đứng trong computational graph. Nó không phải là một independent variable (t.attach_grad() chưa được gọi). Nó cũng không phải là một intermediary variable, i.e. biến trung gian được tính từ các independent variable. Nó chỉ đơn thuần được xem là một hằng số.

Ngoài ra, sau khi backward một biến, biến đó sẽ bị xóa khỏi computational graph.

Head gradient

Xét hàm và , theo chain rule thì

Ta biết được , tính được như dưới đây:

with autograd.record():
    z = 3*x*x*x
z.backward()
A = x.grad.copy()
A

[ 0. 9. 36. 81.]

Và ta cũng tính được và như dưới đây:

with autograd.record():
    t = 3*x*x
    y = t.detach()
    y.attach_grad()
    z = y*x

z.backward()
B, C = x.grad.copy(), y.grad.copy()
B, C

( [ 0. 3. 12. 27.] , [ 0. 1. 2. 3.] )

Và ta cũng tính được như dưới đây:

t.backward()
D = x.grad.copy()
D

[ 0. 6. 12. 18.]

Ta kiểm chứng

A, C*D + B

( [ 0. 9. 36. 81.] , [ 0. 9. 36. 81.] )

Khi chúng ta backward biến t, có thể truyền vào một đối số như sau:

t.backward(u)

Khi đó, các derivative của , bất kể đối với biến nào, đều được nhân với u. Tức là, giả sử D = x.grad sau khi gọi t.backward() và E = x.grad sau khi gọi t.backward(u), thì .

Mặc định của u là một mảng toàn các số 1.

Ta có thể kiểm chứng lại :

with autograd.record():
    t = 3*x*x
    y = t.detach()
    y.attach_grad()
    z = y*x
z.backward()
assert (C - y.grad).sum() == 0
t.backward(y.grad)
E = x.grad.copy()
bool((E - C*D).sum() == 0)

True