Optimizer

Gradient chỉ độ dốc, optimizer chọn bước

Các thuật toán dùng gradient để cập nhật trọng số, như SGD hay Adam, gọi chung là optimizer. Cùng một gradient, mỗi optimizer đi một bước khác nhau.

SGD cập nhật $w_{t+1}=w_t-\eta g_t$, với learning rate $\eta$ quyết định độ dài bước. Chữ stochastic (ngẫu nhiên) đến từ việc tính $g_t$ trên một minibatch lấy ngẫu nhiên thay vì trên toàn bộ dữ liệu. Gradient của một minibatch lệch khỏi gradient đầy đủ một ít, nhưng khi minibatch được lấy ngẫu nhiên đều thì kỳ vọng của nó đúng bằng gradient đầy đủ: mỗi bước rẻ hơn nhiều, đổi lại đường đi hơi nhiễu.

Có ba cỡ dữ liệu cho mỗi bước:

• Gradient descent dùng cả tập: gradient chính xác nhưng mỗi bước đắt.
• SGD thuần dùng một mẫu: rất rẻ nhưng rất nhiễu; một epoch $N$ bước tốn cỡ một bước trên cả tập mà cập nhật được $N$ lần.
• Minibatch đứng giữa và là cách dùng mặc định.

Nhiễu của gradient minibatch có độ lệch chuẩn tỉ lệ $1/\sqrt B$: tăng batch từ 16 lên 64 thì nhiễu giảm một nửa, còn batch quá nhỏ làm đường đi chập chờn. Tăng batch $k$ lần thì số bước mỗi epoch giảm $k$ lần; quy tắc tỉ lệ tuyến tính cho SGD nhân learning rate với $k$, thường kèm warmup, và vẫn phải kiểm lại trên validation.

Một mặt loss kiểm nhỏ là $J(w_1,w_2)=w_1^2+8w_2^2$. Gradient $(2w_1,16w_2)$ thay đổi gắt theo chiều $w_2$, nên một bước SGD tách thành hai phép nhân độc lập: $w_1\leftarrow(1-2\eta)w_1$ và $w_2\leftarrow(1-16\eta)w_2$. Chiều $w_2$ chỉ ổn định khi $|1-16\eta|<1$, tức $0<\eta<1/8$; ngay cả với $\eta$ sát $1/8$, chiều $w_1$ vẫn chỉ được nhân $1-2/8=0.75$ mỗi bước.

Tổng quát với hàm bậc hai, chiều cong nhất (trị riêng lớn nhất $\lambda_{\max}$ của Hessian, ở đây 16) đặt ngưỡng $\eta<2/\lambda_{\max}$, còn chiều thoải nhất (ở đây 2) quyết định tốc độ. Độ cong hai chiều càng chênh, thung lũng càng hẹp và SGD càng chậm. Con số $1/8$ là ngưỡng của riêng hàm này; trong mạng neural, độ cong đổi theo từng vùng của mặt loss.

Nếu biết độ cong của mọi chiều, có thể chia bước cho độ cong: phương pháp Newton nhân gradient với nghịch đảo Hessian (ma trận đạo hàm bậc hai). Trên hàm bậc hai lồi, một bước Newton tới thẳng cực tiểu từ bất kỳ đâu; vì vậy với ít tham số, như logistic regression, Newton cần rất ít bước dù mỗi bước đắt.

Cái giá nằm ở kích thước. Mạng có $d$ tham số có Hessian $d\times d$; với $d=10^4$ đã là $10^8$ phần tử, và nghịch đảo tốn cỡ $d^3=10^{12}$ phép tính mỗi bước, trong khi phép cập nhật bậc một chỉ tốn cỡ $d$ phép tính sau khi backward đã cho gradient. Vì vậy deep learning dùng optimizer bậc một; momentum và Adam là những cách rẻ để bù một phần độ cong mà không cần dựng Hessian.

Momentum nhớ hướng trước

Một quy ước thường dùng, cũng là dạng SGD momentum của PyTorch khi không bật dampening hay Nesterov, là

$v$ cộng dồn gradient với hệ số nhớ $\mu\in[0,1)$. Gradient cùng dấu qua nhiều bước thì được cộng lên, bước dài dần; gradient đổi dấu qua lại thì triệt nhau một phần. Vì vậy momentum tiến nhanh dọc đáy thung lũng thoải và hãm dao động ở chiều dốc, nhưng quán tính quá lớn cũng đưa đường đi vượt đáy. $v$ là trạng thái số được optimizer lưu giữa các bước, không phải vận tốc vật lý.

Hình chạy 25 bước trên $J=w_1^2+8w_2^2$ từ điểm $(-2,1)$ với $\eta=0.09$. SGD nhân $w_2$ với $1-16\cdot0.09=-0.44$ mỗi bước, nên $w_2$ đổi dấu liên tục: đường đi zigzag nhưng co nhanh. Chiều $w_1$ được nhân $0.82$ mỗi bước nên tiến đều về đáy. Thêm momentum $\mu=0.8$ với cùng $\eta$, quán tính tích lũy ở những bước đầu đưa đường đi vượt qua đáy rồi vòng lại: momentum bước dài hơn hẳn SGD, nên $\eta$ và $\mu$ phải chọn cùng nhau.

Một số sách viết $v_t=\mu v_{t-1}+(1-\mu)g_t$, tức trung bình có trọng số thay cho tổng. Ở quy ước này với $\mu = 0.9$ và gradient luôn bằng 1, $v$ là $0.1$ rồi $0.19$, nên bước đầu ngắn hơn hẳn. Khi gradient ổn định, hai quy ước chênh nhau đúng hệ số $1/(1-\mu)$ ở độ dài bước, nên learning rate chép từ tài liệu dùng quy ước này sang quy ước kia phải đổi thang. Mọi phép tính trong bài dùng quy ước đầu.

Nesterov là biến thể tính gradient sau khi đã nhích trước theo quán tính (nesterov=True trong PyTorch).

Adam điều chỉnh từng tọa độ

Adam giữ trung bình có trọng số của gradient, $m$ (giống momentum dạng có $1-\beta_1$), và của bình phương gradient, $v$ (ý tưởng của RMSprop):

Ở đây bình phương và chia là theo từng phần tử. Vì khởi tạo $m_0=v_0=0$, các trung bình đầu bị kéo về 0; Adam sửa bằng $\hat m_t=m_t/(1-\beta_1^t)$ và $\hat v_t=v_t/(1-\beta_2^t)$. Bước mới là

Với $\beta_1 = 0.9$ và gradient luôn bằng 1, $m_2 = 0.19$ còn $\hat m_2 = 0.19/(1 - 0.81) = 1$: bước sửa lệch (bias correction) trả về đúng gradient.

Phép chia cho $\sqrt{\hat v}$ làm độ dài bước gần như không phụ thuộc thang gradient. Nếu một tọa độ luôn có gradient $g$ cố định thì $\hat m=g$, $\hat v=g^2$, và bước bằng $\eta\,g/|g|=\pm\eta$, dù $g$ là 100 hay 0.01. Tọa độ có gradient lớn không còn bước quá dài, tọa độ có gradient nhỏ vẫn được đi bước đủ dài. Vì thang bước khác hẳn SGD, learning rate tốt của Adam (mặc định của PyTorch là $10^{-3}$) không chép sang SGD được, và ngược lại. $\varepsilon$ tránh chia cho 0 khi $\hat v$ gần 0.

Adam không biết feature nào quan trọng; nó chỉ chỉnh thang bước theo lịch sử gradient của từng tọa độ. Thích nghi theo tọa độ cũng không bảo đảm validation tốt hơn: nhiều thực nghiệm đã công bố, nhất là ở phân loại ảnh, ghi nhận SGD momentum với lịch rate được dò kỹ đạt validation ngang hoặc tốt hơn Adam (Wilson và cộng sự, 2017). Vì vậy chọn optimizer vẫn là một thí nghiệm.

Tính thử một bước rồi đối chiếu với PyTorch. Lấy $J(w)=(w-3)^2$, bắt đầu $w=2$, gradient $g=-2$. SGD với $\eta=0.1$ cho $w=2.2$. Adam bước đầu với $\beta_1=0.9,\beta_2=0.999$ có $m=-0.2,v=0.004$, sau bias correction $\hat m=-2,\hat v=4$; bước có độ lớn gần 0.1 nên $w$ gần 2.1. Không phải Adam tính sai đạo hàm: như đã nói ở trên, bước đầu của Adam dài cỡ $\eta$ dù gradient là $-2$ hay $-200$.

import numpy as np
import torch

torch.set_num_threads(1)
torch.manual_seed(23)
w = 2.
g = 2 * (w - 3)
m = .9 * 0 + .1 * g
v = .999 * 0 + .001 * g ** 2
m_hat = m / (1 - .9)
v_hat = v / (1 - .999)
w_adam = w - .1 * m_hat / (np.sqrt(v_hat) + 1e-8)
print('SGD:', round(w - .1 * g, 6))
print('Adam states:', round(m, 6), round(v, 6))
print('Adam:', round(w_adam, 6))

p = torch.nn.Parameter(torch.tensor(2., dtype=torch.float64))
adam = torch.optim.Adam([p], lr=.1)
adam.zero_grad(set_to_none=True)
((p - 3) ** 2).backward()
adam.step()
print('library match:', np.isclose(p.item(), w_adam))

Kết quảSGD: 2.2
Adam states: -0.2 0.004
Adam: 2.1
library match: True

Đọc library match để đối chiếu bản một bước với PyTorch, dùng cùng epsilon và beta mặc định. Đối chiếu một bước không thay cho kiểm một loop dài, nhưng giúp phát hiện quên bias correction hoặc dùng gradient chưa bình phương.

So đường loss với ngân sách bằng nhau

Ta chạy ba optimizer từ cùng điểm khởi tạo, cùng 40 bước, trên hàm bậc hai ở đầu bài. Phép so này đo tiến trình tối ưu theo số bước; nó không đo thời gian chạy và không đo validation. Mỗi optimizer có thang learning rate riêng, nên mỗi cái dùng một rate hợp với nó thay vì ép chung một con số.

def run_optimizer(kind, lr, steps=40):
    p = torch.nn.Parameter(torch.tensor([-2., 1.]))
    if kind == 'momentum':
        optimizer = torch.optim.SGD([p], lr=lr, momentum=.8)
    elif kind == 'adam':
        optimizer = torch.optim.Adam([p], lr=lr)
    else:
        optimizer = torch.optim.SGD([p], lr=lr)
    start = (p[0] ** 2 + 8 * p[1] ** 2).item()
    for _ in range(steps):
        optimizer.zero_grad(set_to_none=True)
        loss = p[0] ** 2 + 8 * p[1] ** 2
        loss.backward()
        optimizer.step()
    end = (p[0] ** 2 + 8 * p[1] ** 2).item()
    return start, end

for kind, lr in [('sgd', .03), ('momentum', .03), ('adam', .15), ('sgd', .2)]:
    start, end = run_optimizer(kind, lr)
    print(kind, lr, f'{start:.4g}', f'{end:.4g}')

Kết quảsgd 0.03 12 0.02833
momentum 0.03 12 0.0005109
adam 0.15 12 0.01883
sgd 0.2 12 1.981e+28

Dòng SGD rate 0.2 vượt ngưỡng $1/8$ đã tính: chiều $w_2$ bị nhân $-2.2$ mỗi bước nên loss lên tới cỡ $10^{28}$. Ba dòng còn lại đều hội tụ; ở bước 40 momentum thấp nhất, rồi đến Adam, rồi SGD. Thứ hạng ấy gắn với đúng ngân sách 40 bước. Cắt ở bước 30 thì sao?

for steps in (30, 40):
    _, sgd_end = run_optimizer('sgd', .03, steps)
    _, adam_end = run_optimizer('adam', .15, steps)
    print(steps, 'bước:', f'SGD {sgd_end:.3g}', f'Adam {adam_end:.3g}')

Kết quả30 bước: SGD 0.0977 Adam 0.314
40 bước: SGD 0.0283 Adam 0.0188

Ở bước 30, SGD thấp hơn Adam hơn ba lần; mười bước sau thứ tự đảo lại. Hình cho thấy nguyên nhân: SGD với rate nhỏ đi xuống đều, còn Adam và momentum vừa xuống vừa dao động, nên thứ hạng tại một bước cụ thể phụ thuộc pha dao động lúc dừng.

Trong mạng thật, train loss giảm nhanh chỉ tiết kiệm thời gian tìm nghiệm; validation mới cho biết nghiệm ấy có tốt trên dữ liệu mới không.

AdamW và weight decay

L2 regularization (bài 2.2) thêm $\lambda\lVert w\rVert^2/2$ vào loss, tạo gradient phụ $\lambda w$. Với SGD thường, bước $w\leftarrow w-\eta(g+\lambda w)=(1-\eta\lambda)w-\eta g$ đúng là co trọng số một hệ số $1-\eta\lambda$ rồi bước theo gradient dữ liệu. Phép co này gọi là weight decay; với SGD, L2 và weight decay trùng nhau.

Với Adam thì khác: $\lambda w$ đi vào $m$, $v$ và bị chia cho $\sqrt{\hat v}$ như mọi gradient khác, nên phép co không còn là một hệ số chung cho mọi tham số. AdamW (Loshchilov và Hutter) tách hai phần: bước Adam chỉ dùng gradient dữ liệu, còn phép co $1-\eta\lambda$ áp thẳng lên trọng số. Trong PyTorch, torch.optim.Adam(..., weight_decay=λ) vẫn làm theo kiểu L2; muốn kiểu tách phải dùng AdamW.

Để thấy riêng phần co, dùng AdamW vừa khởi tạo với các moment bằng 0 và cho gradient dữ liệu bằng 0. Nếu optimizer đã có lịch sử, moment cũ vẫn có thể tạo bước cập nhật dù gradient hiện tại bằng 0. Lưu ý grad=None có thể khiến optimizer bỏ qua tham số; ta đặt gradient tensor toàn 0 để kiểm decay:

q = torch.nn.Parameter(torch.tensor([2.]))
adamw = torch.optim.AdamW([q], lr=.1, weight_decay=.2)
q.grad = torch.zeros_like(q)
adamw.step()
print('after decay:', round(q.item(), 4))

Kết quảafter decay: 1.96

$2(1-0.1\cdot0.2)=1.96$, khớp output. Với SGD, weight_decay cộng $\lambda w$ vào gradient: $w = 2$, gradient 0.5, $\eta = 0.1$, $\lambda = 0.01$ cho $2 - 0.1(0.5 + 0.02) = 1.948$. Decay quá lớn ép model underfit, nên $\lambda$ cũng được chọn bằng validation. Mô hình lớn thường tách nhóm tham số để bias và tham số của các lớp normalization không bị decay.

Scheduler thay learning rate theo thời gian

Đầu quá trình, bước lớn giúp tiến nhanh; về sau, bước nhỏ giúp đi sâu vào đáy thay vì dao động quanh nó. Với SGD minibatch còn một lý do nữa: ngay cả khi tham số đã ở đúng nghiệm, gradient của từng minibatch vẫn khác 0, nên bước cố định tiếp tục đẩy tham số đi lệch; chỉ khi $\eta$ giảm thì vùng dao động mới thu nhỏ.

StepLR nhân rate với gamma sau mỗi step_size lần gọi. Cosine giảm trơn từ $\eta_{\max}$ về $\eta_{\min}$ trong $T$ bước:

nên ở giữa lịch, $t=T/2$, rate đúng bằng trung bình hai đầu. Các lịch khác:

• Warmup làm ngược lại ở đầu train: tăng tuyến tính từ gần 0 lên rate mục tiêu trong một số bước, rồi mới giảm. Với Adam, ở những bước đầu $\hat v$ mới tổng hợp từ vài gradient nên chưa phản ánh thang thật, trong khi bước Adam dài cỡ $\eta$ ở mọi tọa độ cùng lúc. Mô hình lớn khởi tạo ngẫu nhiên như Transformer dễ nổ loss nếu nhận loạt bước dài như vậy ngay từ đầu.
• OneCycle gộp hai pha trong một chu kỳ: tăng lên đỉnh rồi giảm sâu.
• ReduceLROnPlateau nhận một metric và chỉ giảm rate khi metric ngừng cải thiện. Với patience=2, metric không cải thiện quá hai lần gọi liên tiếp thì rate bị nhân với factor.

Mọi lịch chỉ điều chỉnh độ dài bước; chúng không sửa gradient sai hoặc dữ liệu rò rỉ.

r = torch.nn.Parameter(torch.tensor(2.))
sgd = torch.optim.SGD([r], lr=.1)
scheduler = torch.optim.lr_scheduler.StepLR(sgd, step_size=3, gamma=.5)
rates = []
for epoch in range(9):
    rates.append(sgd.param_groups[0]['lr'])
    sgd.zero_grad(set_to_none=True)
    ((r - 3) ** 2).backward()
    sgd.step()
    scheduler.step()
print([round(value, 4) for value in rates])

Kết quả[0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.025, 0.025, 0.025]

Rate được đọc trước khi cập nhật; scheduler.step() chạy sau optimizer.step(), một lần mỗi epoch. Ví dụ rate 0.1 với StepLR(step_size=10, gamma=0.1) gọi mỗi epoch: epoch 0-9 dùng 0.1, epoch 10-19 dùng 0.01.

Lịch theo step đếm số batch, lịch theo epoch đếm số epoch. Đổi batch size mà giữ nguyên lịch theo step thì số batch mỗi epoch đổi, nên rate giảm ở một thời điểm khác tính theo lượng dữ liệu đã học. Ngược lại, OneCycle tạo với total_steps=len(loader) * 10 mà chỉ gọi step() mỗi epoch thì sau 10 epoch lịch mới đi 10 bước, rate còn chưa lên tới đỉnh.