🔁 Lan truyền ngược qua lớp tuyến tính và lớp tích chập

By Better Wise

Vị trí của bài này: từ shape đúng đến gradient đúng

Hai bài trước đã cho bạn forward pass của tích chập và cách vector hóa nó bằng im2col: biến một phép conv thành một phép nhân ma trận. Bài này đi tiếp bước cuối của module: nếu forward là một chuỗi phép toán, thì backward chỉ là áp dụng chain rule ngược lại chuỗi đó, với một ràng buộc duy nhất và không thể thương lượng: mỗi gradient phải có đúng shape của tensor tương ứng ở forward.

Mục tiêu cụ thể, sau 45 phút bạn có thể:

  • Viết ba công thức backward của lớp linear: dX, dW, db.
  • Suy ra backward của lớp conv từ chính công thức linear nhờ im2col.
  • Giải thích vì sao dX cần scatter-add còn dW thì không.
  • Sanity check gradient bằng shape và bằng sai phân số.
  • Tự tay viết linear_backward và conv_backward chạy được.

Note: Nhắc lại nhanh từ bài trước: với batch X shape (N, C, H, W) và kernel (F, C, KH, KW), im2col duỗi mỗi receptive field thành một hàng, cho X_col shape (N*OH*OW, C*KH*KW). Khi đó forward conv chính xác là Y_col = X_col @ W_col với W_col shape (C*KH*KW, F).

Nguyên tắc duy nhất: gradient đi ngược, shape đi xuôi

Mỗi lớp trong mạng nhận một upstream gradient dY = ∂L/∂Y từ lớp phía sau, và phải trả về ba thứ:

  1. dX = ∂L/∂X, để chuyển tiếp cho lớp phía trước.
  2. dW = ∂L/∂W, để optimizer cập nhật.
  3. db = ∂L/∂b, nếu lớp có bias.
flowchart LR
    DY["⬅️<br/>dY"] --> DW["🧮<br/>dW = Xᵀ dY"]
    DY --> DB["➕<br/>db = Σ dY"]
    DY --> DX["↩️<br/>dX = dY Wᵀ"]
    X["📥<br/>X"] --> DW
    DX --> PREV["⬅️<br/>lớp trước"]

Quy tắc kiểm tra nhanh nhất: dX.shape ** X.shape, dW.shape ** W.shape, db.shape == b.shape. Nếu một công thức cho ra shape sai, nó sai, không cần chạy số.

Backward của lớp tuyến tính: ba công thức, ba shape

Với Y = X @ W + b, trong đó X là (N, D_in), W là (D_in, D_out), b là (D_out,):

$$ \frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y} W^{\top}, \qquad \frac{\partial L}{\partial W} = X^{\top} \frac{\partial L}{\partial Y}, \qquad \frac{\partial L}{\partial b} = \sum_{n=1}^{N} \frac{\partial L}{\partial Y_{n}} $$

Đọc theo shape để tự kiểm tra: dY là (N, D_out), nhân với Wᵀ là (D_out, D_in) cho ra (N, D_in) đúng bằng X. Tương tự Xᵀ là (D_in, N) nhân dY là (N, D_out) cho ra (D_in, D_out) đúng bằng W.

Điểm dễ sai nhất là db. Bias được broadcast cho mọi mẫu trong batch, nên mỗi mẫu đều đóng góp một phần gradient: db là tổng theo trục batch, không phải trung bình. Nếu bạn chia trung bình ở đây, bạn đang vô tình chia gradient của bias cho N trong khi dW thì không, và training sẽ lệch một cách rất khó debug.

def linear_backward(dY, X, W):
    dX = dY @ W.T          # (N, D_in)
    dW = X.T @ dY          # (D_in, D_out)
    db = dY.sum(axis=0)    # (D_out,)
    return dX, dW, db

Ba dòng, và mỗi dòng bạn nên đọc được shape của nó mà không cần chạy.

Bài tập 1

Một lớp linear có D_in = 6, D_out = 4, batch N = 3. Không chạy code, hãy cho biết shape của dY, dX, dW, db.

Xem đáp án

dY: (3, 4). dX: (3, 6). dW: (6, 4). db: (4,). Kiểm tra: dX khớp X, dW khớp W, db khớp b. Nếu bạn viết db là (3, 4) thì bạn quên phép tổng theo batch.

Backward của lớp tích chập qua lăng kính im2col

Đây là phần đẹp nhất của bài. Vì Y_col = X_col @ W_col đúng là một lớp linear, backward của conv chính là backward của linear, chỉ thêm hai bước đổi biểu diễn ở hai đầu:

  1. Forward lưu lại X_col (hoặc lưu X rồi dựng lại X_col, đổi bộ nhớ lấy thời gian).
  2. Backward tính dW_col và dX_col bằng công thức linear.
  3. dW_col reshape về (F, C, KH, KW).
  4. dX_col được scatter-add ngược về dX shape (N, C, H, W).

Cụ thể:

$$ \frac{\partial L}{\partial W_{col}} = X_{col}^{\top} \frac{\partial L}{\partial Y_{col}}, \qquad \frac{\partial L}{\partial X_{col}} = \frac{\partial L}{\partial Y_{col}} W_{col}^{\top} $$

Vì sao dX phải scatter-add chứ không phải gán? Vì các receptive field chồng lấn nhau. Khi stride nhỏ hơn kernel size, một pixel input xuất hiện trong nhiều hàng của X_col, tức là nó góp vào nhiều phần tử output. Đạo hàm của một biến được dùng nhiều lần là tổng các đóng góp. Nếu bạn gán đè, bạn chỉ giữ đóng góp cuối cùng và gradient sẽ sai một cách im lặng, loss vẫn giảm chậm nên rất khó phát hiện. Ngược lại dW không cần scatter vì mỗi trọng số kernel ứng với đúng một cột của W_col, và phép X_colᵀ @ dY_col đã tự cộng dồn qua mọi vị trí không gian và mọi mẫu trong batch.

def conv_backward(dY_col, X_col, W_col, F, C, KH, KW):
    dW_col = X_col.T @ dY_col              # (C*KH*KW, F)
    dX_col = dY_col @ W_col.T              # (N*OH*OW, C*KH*KW)
    dW = dW_col.T.reshape(F, C, KH, KW)    # khớp shape kernel
    dX = col2im_scatter_add(dX_col)        # cộng dồn các receptive field
    db = dY_col.sum(axis=0)                # (F,)
    return dX, dW, db

Với bias của conv: mỗi filter có một bias, và bias đó được cộng vào mọi vị trí output của mọi mẫu, nên db là tổng dY theo cả batch lẫn hai trục không gian.

  1. Giữ lại X_col và indices từ forward, đây là toàn bộ những gì backward cần
  2. Tính dW_col = X_col.T @ dY_col rồi reshape về shape kernel
  3. Tính dX_col = dY_col @ W_col.T
  4. Khởi tạo dX = zeros(X.shape) rồi scatter-add dX_col theo indices
  5. Tính db bằng tổng dY theo batch và theo không gian
  6. Assert ba shape khớp X, W, b trước khi trả về

Tip: Nếu bạn nhớ được một câu duy nhất: conv backward là linear backward cộng thêm hai phép đổi biểu diễn, im2col ở đầu vào và col2im (scatter-add) ở đầu ra.

Bài tập 2

Cho conv có C = 3, F = 8, KH = KW = 3, stride = 1, padding = 1, input (2, 3, 10, 10). Hãy cho biết shape của X_col, W_col, dW_col, và giải thích tại sao dX không thể lấy trực tiếp từ dX_col.

Xem đáp án

OH = OW = 10, nên X_col là (200, 27), W_col là (27, 8), dW_col là (27, 8). dX_col là (200, 27) nhưng dX phải là (2, 3, 10, 10). Không lấy trực tiếp được vì mỗi hàng của X_col là một receptive field và chúng chồng lấn (stride 1 < kernel 3), nên nhiều hàng cùng ánh xạ về một pixel của dX. Phải scatter-add để cộng dồn các đóng góp đó, rồi mới reshape về (2, 3, 10, 10).

Kiểm tra bằng số trước khi tin vào công thức

Công thức đúng về shape vẫn có thể sai về dấu hoặc hệ số. Trước khi train thật, hãy làm gradient check trên một mạng nhỏ với vài chiều: so dW giải tích với xấp xỉ sai phân trung tâm _____. Sai số tương đối dưới 1e-7 với float64 là tín hiệu tốt. Với conv, hãy test cả trường hợp stride 1 (chồng lấn mạnh, dễ lộ lỗi scatter-add) và trường hợp stride bằng kernel size (không chồng lấn, dX gần như chỉ là gán).

Vài con số nên nhớ: dX của lớp linear là _____, dW là _____, và db là _____. Trong conv, gradient theo input được tích lũy bằng _____ vì các receptive field chồng lấn.

Những lỗi làm gradient sai mà không báo lỗi

  • Quên cộng dồn bias trên batch và không gian. Đây là lỗi phổ biến nhất của conv. db luôn là tổng, không bao giờ là trung bình, trừ khi bạn chia trung bình cho toàn bộ dY một cách nhất quán.
  • Gán thay vì scatter-add cho dX. Loss vẫn giảm nên bạn tưởng mọi thứ ổn, nhưng mô hình học chậm và kém ổn định. Test bằng stride 1 là cách lộ lỗi nhanh nhất.
  • Nhầm dY với Y. Backward nhận gradient, không nhận output. Nếu bạn thấy mình đang dùng Y trong conv_backward, hãy dừng lại và kiểm tra biến.
  • Quên reset buffer dX về 0. Nếu tái sử dụng một mảng dX giữa các lần gọi, gradient sẽ cộng dồn qua các iteration và nổ ngay.
  • Reshape kernel sai thứ tự trục. W_col là (C*KH*KW, F), nên dW_col.T là (F, C*KH*KW) và reshape thành (F, C, KH, KW) mới khớp. Đảo C và KH sẽ không lỗi shape nếu tình cờ bằng nhau, và đó là kiểu bug tệ nhất.
  • Gradient check trên float32. Sai phân số cần độ chính xác cao. Hãy check trên float64 với vài chiều nhỏ, rồi mới chạy float32 cho training.

Chốt lại và bước sang checkpoint

Ba công thức của lớp linear (dX = dY Wᵀ, dW = Xᵀ dY, db = Σ dY) là toàn bộ nội dung bài này, phần conv chỉ là phiên bản đổi biểu diễn của chúng qua im2col và col2im. Nếu bạn viết được hai hàm backward đó và assert shape ở mỗi bước, bạn đã có đủ công cụ để debug bất kỳ mạng nơ-ron nào ở mức toán học.

Bài tiếp theo là checkpoint của module: bạn sẽ tự viết forward và backward cho một mạng nhỏ, đối chiếu với gradient check, và xem những lỗi ở trên xuất hiện trong code của chính mình như thế nào.

Front: Forward của lớp tuyến tính có dạng gì? Back: Y = X W + b, với X là (N, D_in), W là (D_in, D_out), b là (D_out).

Front: dX của lớp tuyến tính bằng gì? Back: dX = dY Wᵀ, shape (N, D_in).

Front: dW của lớp tuyến tính bằng gì? Back: dW = Xᵀ dY, shape (D_in, D_out).

Front: db của lớp tuyến tính bằng gì? Back: Tổng dY theo trục batch, shape (D_out,).

Front: Vì sao db là tổng chứ không phải trung bình? Back: Bias được broadcast cho mọi mẫu nên mọi mẫu đều đóng góp gradient cho nó.

Front: dW_col của lớp conv bằng gì? Back: dW_col = X_colᵀ dY_col, rồi reshape về (F, C, KH, KW).

Front: dX_col của lớp conv bằng gì? Back: dX_col = dY_col W_colᵀ, rồi scatter-add về (N, C, H, W).

Front: Vì sao dX của conv cần scatter-add? Back: Vì các receptive field chồng lấn, một pixel input góp vào nhiều phần tử output.

Front: Shape của gradient so với tensor tương ứng ở forward? Back: Luôn trùng khớp tuyệt đối: dX như X, dW như W, db như b.

Front: Gradient check dùng công thức nào? Back: Sai phân trung tâm: (f(θ + ε) - f(θ - ε)) / (2ε), chạy trên float64.

Front: im2col giúp gì cho backward của conv? Back: Biến conv thành matmul nên dùng lại được nguyên công thức backward của lớp linear.

Front: Trường hợp stride nào dễ lộ lỗi scatter-add nhất khi test? Back: Stride 1, vì mức chồng lấn của receptive field là lớn nhất. Gradient theo input của lớp tuyến tính là _____, có shape trùng với X. Gradient theo weight của lớp tuyến tính là _____, có shape trùng với W. Gradient theo bias của lớp tuyến tính là _____. Trong conv, dW_col bằng _____, sau đó reshape về (F, C, KH, KW). Để có dX từ dX_col, ta phải _____ các cột về vị trí pixel gốc. Nghiệm thu số học của backward nên dùng _____, không dùng float32.

Những chỗ dễ sai khi tự viết backward

  • Quên cộng dồn ở bias. db là reduce theo batch, không phải lấy một hàng của dY. Nếu shape ra (N, D_out) thì bạn đang thiếu bước sum.
  • Transpose sai chiều. Viết ra giấy shape của từng toán hạng trước khi gõ code. dY @ W.T hợp lệ chỉ khi W là (D_in, D_out). Nếu bạn lưu W theo kiểu (D_out, D_in) thì công thức đổi thành dY @ W.
  • Scatter-add biến thành gán. Dùng dX[..., h, w] += ... hoặc np.add.at, tuyệt đối không dùng =. Lỗi này không làm code crash, chỉ làm gradient nhỏ đi một cách âm thầm.
  • Quên zero-init cho dX_col. Nếu tái sử dụng buffer giữa các lần gọi, gradient sẽ cộng dồn qua nhiều batch.
  • Test bằng stride lớn. Stride 2 hoặc 3 che giấu lỗi chồng lấn. Luôn test stride 1 trước, vì đó là trường hợp receptive field chồng lấn nhiều nhất.
  • Gradient check trên float32. Sai số sai phân xấp xỉ 1e-3 sẽ lấn át sai số thật. Chuyển toàn bộ mạng sang float64 và dùng ε ≈ 1e-5.
  • Quy tắc ngón tay: nếu gradient check lệch ở một vài phần tử góc (biên ảnh), gần như chắc chắn là lỗi index khi map (h, w) sang cột X_col, không phải lỗi công thức.

Chốt lại và cầu nối sang checkpoint

Ba ý cần mang theo:

  1. Lớp tuyến tính là một phép matmul, nên backward của nó chỉ là ba phép matmul: dX = dY Wᵀ, dW = Xᵀ dY, db = sum(dY, axis=0).
  2. Lớp conv không có công thức backward riêng. Sau im2col nó chính là lớp tuyến tính, cộng thêm một phép scatter-add để trả gradient về đúng pixel.
  3. Bất biến duy nhất luôn đúng: shape của dX, dW, db trùng khít shape của X, W, b. Kiểm tra bất biến này trước khi kiểm tra giá trị.

Bài tiếp theo là checkpoint của module. Ở đó bạn sẽ không được xem trước công thức: bạn nhận một đoạn forward bị lỗi shape hoặc lỗi scatter-add và phải tự tìm ra, rồi chứng minh bằng gradient check rằng bản sửa của mình đúng. Hãy chuẩn bị sẵn một hàm grad_check chạy được trên float64, vì đó là công cụ chấm điểm chính của bài đó.

Front: Định nghĩa gradient của một tham số θ? Back: Vector (hoặc tensor) các đạo hàm riêng ∂L/∂θ, cùng shape với θ.

Front: Vì sao dX cần cả dY và W, không cần b? Back: Vì b không xuất hiện trong biểu thức của Y theo X, nên ∂Y/∂X chỉ phụ thuộc W.

Front: Trong conv, một pixel input ảnh hưởng tới bao nhiêu phần tử output? Back: Tới mọi vị trí output mà receptive field của nó phủ lên, nên cần scatter-add.

Front: Sai số của sai phân trung tâm theo bậc ε là bao nhiêu? Back: O(ε²), tốt hơn nhiều so với sai phân tiến O(ε).

Front: Vì sao phải kiểm tra shape gradient trước khi kiểm tra giá trị? Back: Vì lỗi shape thường là lỗi transpose hoặc thiếu reduce, và nó làm mọi so sánh số vô nghĩa.