Đã đăng vào thg 1 26, 9:49 SA

trong

10 phút đọc

[Papers Notes] RMSNorm: Giải quyết một nửa vấn đề của Layer Norm nhưng lại được ưa chuộng hơn

Mình đang trong quá trình tìm hiểu lại các Foudation Language Model, thì đọc đến Llama 1 thì có đề cập tới kỹ thuật này. Dù ban đầu tính là sẽ nhét hết vào một bài viết về Llama, nhưng thấy rằng có khá nhiều kiến thức bị bỏ qua nên đánh riêng một bài luôn.

I. Layer Normalization

Layer Normalization (LayerNorm) là một kỹ thuật chuẩn hóa được giới thiệu vào năm 2016 bởi Jimmy Lei Ba và các cộng sự. Khác với các phương pháp chuẩn hóa khác như Batch Normalization (BatchNorm), LayerNorm hoạt động dựa trên việc chuẩn hóa \textit{các giá trị bên trong một lớp duy nhất} cho từng mẫu dữ liệu riêng lẻ. Điều này giúp LayerNorm thích nghi tốt với các loại dữ liệu có độ dài thay đổi hoặc các mô hình phụ thuộc thứ tự như Recurrent Neural Networks (RNNs) và Transformers.

1. Nguyên lý hoạt động của LayerNorm

Trong LayerNorm, tất cả các đầu vào đến neuron của một lớp tại một thời điểm sẽ được chuẩn hóa bằng cách sử dụng trung bình (mean) và phương sai (variance) của chính chúng. Cụ thể:

Cho ( a ) là vector chứa tổng đầu vào (weighted sum) của các neuron trong một lớp:

$mu = \frac{1}{H} \sum_{i=1}^H a_i, \quad \sigma^2 = \frac{1}{H} \sum_{i=1}^H (a_i - \mu)^2$

Vector chuẩn hóa sẽ được tính như sau:

$\hat{a}_i = \frac{a_i - \mu}{\sqrt{\sigma^2 + \epsilon}}$

Sau đó, LayerNorm áp dụng các tham số học được ( \gamma ) (gain) và ( \beta ) (bias) để tái định dạng dữ liệu:

$y_i = \gamma \hat{a}_i + \beta$

Trong đó:

$mu$ : Trung bình, giúp tái định tâm (re-centering).
$igma$ : Độ lệch chuẩn, đảm bảo tái định tỷ lệ (re-scaling).
$epsilon$ : Giá trị rất nhỏ (e.g., (10^{-5})) để tránh chia cho (0).
$gamma$ và $beta$ : Các tham số học được để tăng khả năng biểu diễn của mô hình.

2. Các đặc tính chính của LayerNorm

a. Re-centering (Tái định tâm)

LayerNorm giúp mạng không bị ảnh hưởng bởi độ lệch trung bình của dữ liệu, làm cho mô hình trở nên bất biến với các thay đổi nhỏ trong dữ liệu đầu vào và trọng số.

b. Re-scaling (Tái định tỷ lệ)

Tái định tỷ lệ giữ cho đầu ra của mỗi lớp có độ lớn ổn định, bất kể các giá trị đầu vào hoặc trọng số có bị phóng đại hay thu nhỏ. Đây là yếu tố quan trọng trong việc ổn định gradient và giúp huấn luyện hiệu quả hơn.

c. Không phụ thuộc vào batch size

Khác với BatchNorm, LayerNorm không phụ thuộc vào kích thước batch. Điều này giúp LayerNorm hoạt động tốt trong các mô hình sử dụng batch nhỏ hoặc kích thước batch bằng (1) (e.g., trong RNNs).

3. Vấn đề của LayerNorm

Chi phí tính toán cao: LayerNorm cần tính trung bình (( \mu )) và phương sai (( \sigma^2 )) trên toàn bộ neuron trong một lớp. Điều này làm tăng độ trễ, đặc biệt trong các mạng lớn như Transformers.
Re-centering không cần thiết: Nghiên cứu chỉ ra rằng re-centering không đóng vai trò then chốt trong việc ổn định hoặc cải thiện hiệu quả huấn luyện. Điều này gây ra chi phí không đáng có.
Không tối ưu cho RNNs: Khi được áp dụng trong RNNs, LayerNorm tăng chi phí tính toán tại mỗi bước thời gian, làm giảm hiệu suất tổng thể.
Hạn chế tốc độ hội tụ: Sự dư thừa trong việc điều chỉnh thông qua cả re-centering và re-scaling có thể làm giảm tốc độ hội tụ của mô hình.

II. RMSNorm

4. Dẫn đến RMSNorm – Tối ưu hóa LayerNorm

Từ những vấn đề trên, Root Mean Square Normalization (RMSNorm) được đề xuất nhằm giảm chi phí tính toán trong khi vẫn duy trì các lợi ích chính của LayerNorm. RMSNorm loại bỏ hoàn toàn re-centering và chỉ tập trung vào re-scaling bằng cách sử dụng Root Mean Square (RMS) để chuẩn hóa:

$\hat{a}_i = \frac{a_i}{\sqrt{\frac{1}{H} \sum_{i=1}^H a_i^2 + \epsilon}}$

Ở đây:

RMS thay thế vai trò của cả trung bình $mu$ và độ lệch chuẩn $sigma$ , giúp giảm bớt bước tính toán.
RMSNorm duy trì tính bất biến về tỷ lệ, hỗ trợ gradient ổn định và tốc độ hội tụ nhanh hơn.

Kết quả là RMSNorm không chỉ giảm chi phí tính toán mà còn đạt hiệu quả tương đương hoặc vượt trội LayerNorm trong nhiều tác vụ như dịch máy, phân loại hình ảnh, và các ứng dụng khác.

5. Root Mean Square Normalization (RMSNorm)

Root Mean Square Normalization (RMSNorm) được giới thiệu như một biến thể của LayerNorm nhằm loại bỏ những phần không cần thiết, đặc biệt là bước tái định tâm (re-centering), từ đó giảm thiểu chi phí tính toán mà vẫn duy trì các lợi ích quan trọng của quá trình chuẩn hóa.

a. Khác biệt cơ bản so với LayerNorm

Thay vì sử dụng cả trung bình $mu$ và phương sai $sigma^2$ để chuẩn hóa như LayerNorm, RMSNorm chỉ tập trung vào root mean square (RMS) – một phép đo phản ánh độ lớn của các đầu vào, mà không cần tái định tâm giá trị đầu vào. Công thức chuẩn hóa trong RMSNorm được định nghĩa như sau:

$\hat{a}_i = \frac{a_i}{\sqrt{\frac{1}{H} \sum_{i=1}^H a_i^2 + \epsilon}}$

Trong đó:

$H $: Số lượng neuron trong một lớp.
$a_i$ : Đầu vào (summed input) của neuron ( i ).
$epsilon$ : Một giá trị nhỏ để tránh chia cho 0.

Giống như LayerNorm, RMSNorm cũng sử dụng các tham số học được ( \gamma ) (gain) và ( \beta ) (bias) để tăng khả năng biểu diễn của mô hình:

$y_i = \gamma \hat{a}_i + \beta$

b. Lý do loại bỏ tái định tâm (Re-centering)

Nghiên cứu từ tác giả của RMSNorm chỉ ra rằng:

Re-centering không cải thiện ổn định huấn luyện: Tái định tâm (( \mu )) không có tác động đáng kể trong việc giảm phương sai của hidden states hay gradient.
Tái định tâm không cần thiết với re-scaling: Trong nhiều trường hợp, việc tái định tỷ lệ (re-scaling) đã đủ để ổn định quá trình huấn luyện mà không cần thêm bước tái định tâm.
Hiệu quả tính toán: Bỏ qua bước tính toán ( \mu ) giúp giảm đáng kể chi phí xử lý, đặc biệt trong các mô hình lớn hoặc các mạng yêu cầu hiệu suất cao như RNNs và Transformers.

c. Ưu điểm của RMSNorm

1. Giảm chi phí tính toán:

RMSNorm không yêu cầu tính toán trung bình (( \mu )), dẫn đến giảm độ phức tạp tính toán so với LayerNorm. Điều này làm tăng tốc độ xử lý trong mỗi bước huấn luyện.

2. Ổn định và hội tụ nhanh hơn:

Bằng cách duy trì tính bất biến theo tỷ lệ (re-scaling invariance), RMSNorm đảm bảo gradient ổn định và hỗ trợ tốc độ hội tụ nhanh hơn mà không làm giảm hiệu năng của mô hình.

3. Hiệu quả trong các mạng RNNs và Transformers:

Do RMSNorm chỉ tập trung vào việc tái định tỷ lệ, nó đặc biệt hiệu quả trong các mạng như RNNs và các mô hình tự hồi tiếp khác, nơi mỗi bước thời gian yêu cầu tính toán nhanh gọn.

4. Tính tổng quát và khả năng mở rộng:

RMSNorm có thể dễ dàng tích hợp vào các kiến trúc hiện đại mà không cần điều chỉnh lớn. Thử nghiệm cho thấy RMSNorm đạt được hiệu suất tương đương hoặc vượt trội LayerNorm trên các tác vụ như dịch máy và phân loại hình ảnh.

d. Biến thể: Partial RMSNorm

Tác giả của RMSNorm cũng đề xuất một biến thể gọi là Partial RMSNorm, trong đó chỉ sử dụng một phần (( p% )) các đầu vào để tính giá trị RMS. Điều này làm giảm thêm chi phí tính toán mà vẫn giữ được hiệu năng. Ví dụ:

Khi sử dụng ( p = 6.25% ), hiệu năng của Partial RMSNorm gần tương đương RMSNorm, nhưng tốc độ tăng lên đáng kể.

e. So sánh hiệu năng

1. Tốc độ huấn luyện:

RMSNorm cho thấy cải thiện tốc độ từ 7% đến 64% so với LayerNorm, tùy thuộc vào kiến trúc mô hình và kích thước dữ liệu.

2. Hiệu suất:

RMSNorm đạt hiệu suất tương đương LayerNorm trong các tác vụ như:

Dịch máy (Machine Translation)
Phân loại hình ảnh (Image Classification)
Các bài toán xử lý ngôn ngữ tự nhiên (NLP)

3. Ứng dụng thực tế:

RMSNorm được sử dụng rộng rãi trong các mô hình yêu cầu tối ưu hóa hiệu năng, chẳng hạn như Transformers trong xử lý ngôn ngữ tự nhiên (NLP) và các ứng dụng có thời gian thực.

Conclusions

RMSNorm là một bước cải tiến quan trọng so với LayerNorm, đặc biệt khi yêu cầu giảm chi phí tính toán mà không làm giảm hiệu năng. Với việc loại bỏ tái định tâm, RMSNorm giữ lại tính bất biến về tỷ lệ, hỗ trợ huấn luyện nhanh và ổn định hơn. Đây là một lựa chọn hiệu quả cho các mô hình hiện đại, đặc biệt trong các bài toán yêu cầu xử lý nhanh và hiệu quả.

References

Root Mean Square Layer Normalization

LLM NLP