Batch Normalization: Tăng tốc huấn luyện mạng sâu bằng cách giảm Internal Covariate Shift

Hướng dẫn đọc

Vấn đề và động lực

Trong quá trình tối ưu hóa mạng sâu bằng SGD, sự thay đổi tham số ở các tầng trước sẽ làm thay đổi phân phối đầu vào của các tầng tiếp theo. Các tác giả gọi hiện tượng này là Internal Covariate Shift.

Internal covariate shift tạo ra hai thách thức chính:

1. Các mạng con phải liên tục thích ứng với các phân phối kích hoạt luôn dịch chuyển.
2. Các kích hoạt dễ dàng rơi vào vùng bão hòa của các hàm phi tuyến (chẳng hạn như hàm sigmoid), gây ra hiện tượng triệt tiêu gradient (vanishing gradient) và đòi hỏi tốc độ học nhỏ cùng khâu khởi tạo tham số hết sức cẩn trọng.

Trong khi việc làm trắng (whitening - khử tương quan và chuẩn hóa) đầu vào của tầng trên toàn bộ tập dữ liệu có thể ổn định phân phối, thì việc tính ma trận hiệp phương sai đầy đủ và căn bậc hai nghịch đảo của nó ở mỗi bước cập nhật tham số lại quá tốn kém về mặt tính toán, đồng thời có thể khiến gradient bùng nổ nếu chuẩn hóa không được tích hợp vào quá trình lan truyền ngược (backpropagation).

Ý tưởng then chốt

• Tầng chuẩn hóa khả vi (Differentiable Normalization Layer): Thay vì chuẩn hóa các kích hoạt bên ngoài quá trình tối ưu, chuẩn hóa được tích hợp trực tiếp vào kiến trúc mô hình, đảm bảo lan truyền ngược xét đến sự phụ thuộc giữa các kích hoạt tầng và thống kê của mini-batch.
• Chuẩn hóa theo từng chiều (Per-Dimension Normalization): Thay vì làm trắng toàn phần (joint whitening), từng chiều đặc trưng vô hướng được chuẩn hóa độc lập về giá trị trung bình bằng 0 và phương sai bằng 1.
• Khôi phục khả năng biểu diễn (Restoring Representational Capacity): Để đảm bảo chuẩn hóa không làm mất đi khả năng biểu đạt của một tầng (như việc ép sigmoid hoạt động hoàn toàn tuyến tính), các tham số có thể học $\gamma$ (co giãn) và $\beta$ (dịch chuyển) được đưa vào để cho phép tầng khôi phục lại phép biến đổi đồng nhất (identity transformation) nếu cần.
• Suy luận xác định (Deterministic Inference): Trong quá trình huấn luyện, các giá trị thống kê được tính trên từng mini-batch $\mathcal{B}$. Khi suy luận (test time), các thống kê cố định trên toàn quần thể (được tính qua trung bình động hoặc trên toàn bộ các batch huấn luyện) sẽ thay thế cho thống kê của mini-batch, đảm bảo suy luận mang tính xác định cho từng mẫu đơn lẻ.
• Bất biến với tỷ lệ tham số (Parameter Scale Invariance): Phép biến đổi tỷ lệ của trọng số tầng thỏa mãn $\text{BN}((aW)\mathrm{u})=\text{BN}(W\mathrm{u})$, nghĩa là trọng số lớn hơn sẽ dẫn đến gradient hiệu dụng nhỏ hơn, giúp ngăn chặn hiện tượng bùng nổ gradient.

Phương pháp (với các phương trình chính)

Với một mini-batch $\mathcal{B}=\{x_{1\dots m}\}$ tương ứng với một chiều kích hoạt đơn lẻ:

1. Tính trung bình của mini-batch: $${\mu }_{\mathcal{B}}\leftarrow \frac{1}{m}\sum _{i=1}^m{x}_i$$
2. Tính phương sai của mini-batch: $${\sigma }_{\mathcal{B}}^2\leftarrow \frac{1}{m}\sum _{i=1}^m(x_i-{\mu }_{\mathcal{B}}{)}^2$$
3. Chuẩn hóa kích hoạt (với hằng số ổn định $ϵ$): $${\hat{x}}_i\leftarrow \frac{x_i-{\mu }_{\mathcal{B}}}{\sqrt{{\sigma }_{\mathcal{B}}^2+ϵ}}$$
4. Co giãn và dịch chuyển bằng các tham số có thể học $\gamma$ và $\beta$: $$y_i\leftarrow \gamma {\hat{x}}_i+\beta \equiv {\text{BN}}_{\gamma ,\beta }(x_i)$$

Áp dụng vào các tầng

• Vị trí: Được áp dụng ngay trước các hàm phi tuyến: $z=g(\text{BN}(Wu))$. Hệ số điều chỉnh bias $b$ ban đầu của tầng bị loại bỏ vì phép trừ trung bình đã triệt tiêu nó (vai trò của nó được thay thế bởi $\beta$).
• Các tầng tích chập (Convolutional Layers): Để bảo toàn tính tương đương tịnh tiến (translation equivariance), chuẩn hóa được áp dụng trên mỗi feature map thay vì trên từng kích hoạt không gian riêng biệt. Với batch size $m$ và độ phân giải không gian $p\times q$, kích thước mini-batch hiệu dụng trở thành $m^{\prime }=m\cdot p\cdot q$, với một cặp $(\gamma ,\beta )$ duy nhất được chia sẻ trên mỗi feature map.

Phép biến đổi khi suy luận

Khi kiểm thử, phép biến đổi trở thành một phép toán tuyến tính tĩnh sử dụng thống kê quần thể $\mathrm{E}[x]$ và $\mathrm{Var}[x]$: $$y=\frac{\gamma }{\sqrt{\mathrm{Var}[x]+ϵ}}\cdot x+\left(\beta -\frac{\gamma \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x]+ϵ}}\right)$$

Kết quả (với các số liệu quan trọng nhất)

• MNIST với hàm Sigmoid: Các mạng trang bị Batch Normalization duy trì được phân phối kích hoạt ổn định trong suốt quá trình huấn luyện, tăng tốc hội tụ và đạt độ chính xác kiểm thử cao hơn so với các mạng không chuẩn hóa.
• Tốc độ huấn luyện ImageNet (Mạng đơn lẻ):
  • Inception không chuẩn hóa đạt độ chính xác validation 72,2% sau $31,0\times 10^6$ bước.
  • Baseline Batch-Normalized Inception đạt 72,2% chỉ sau $13,3\times 10^6$ bước.
  • Với tốc độ học cao gấp 5 lần (BN-x5), mô hình đạt 72,2% sau $2,1\times 10^6$ bước—giảm 14 lần số bước huấn luyện cần thiết, cuối cùng đạt đỉnh ở mức 73,0%.
  • Với tốc độ học cao gấp 30 lần (BN-x30), mô hình đạt độ chính xác tối đa 74,8% sau $6,0\times 10^6$ bước.
• Hàm phi tuyến bão hòa: Trong khi Inception không chuẩn hóa dùng kích hoạt sigmoid hoàn toàn thất bại khi huấn luyện (chỉ duy trì ở mức ngẫu nhiên 0,1%), thì BN-x5-Sigmoid đạt độ chính xác top-1 là 69,8%.
• SOTA Ensemble trên ImageNet: Một ensemble gồm 6 mạng có chuẩn hóa theo batch đạt tỷ lệ lỗi top-5 validation là 4,9% và tỷ lệ lỗi test là 4,82% trên test server của ILSVRC, vượt qua hiệu suất của con người.

Hạn chế và các câu hỏi mở

• Phụ thuộc vào kích thước Batch: Do việc chuẩn hóa dựa trên các giá trị thống kê mẫu của mini-batch, kích thước mini-batch quá nhỏ có thể làm suy giảm chất lượng ước lượng phương sai và giá trị trung bình.
• Mạng hồi quy (Recurrent Networks): Bài báo mới chỉ thử nghiệm trên các mô hình truyền thẳng (feedforward) và tích chập, để ngỏ câu hỏi về mạng nơ-ron hồi quy (RNNs) như một hướng nghiên cứu chưa được kiểm chứng, nơi internal covariate shift cũng có thể đóng vai trò then chốt.
• Cơ sở lý thuyết: Các tác giả phỏng đoán rằng Batch Normalization làm cho ma trận Jacobian của các tầng có các giá trị kỳ dị (singular values) gần bằng 1, nhưng lưu ý rằng việc kiểm chứng lý thuyết chính xác về hành vi gradient vẫn còn là vấn đề mở cho các phân tích trong tương lai.