Làm chủ các chiều Tensor trong Transformer
Làm chủ các chiều Tensor trong Transformer
- 10 min read
Làm chủ Chiều Tensor trong Transformer
Bài viết trên blog của Not Lain trên Hugging Face
Lời giới thiệu
Việc nắm vững các hình dạng và chiều của tensor là yếu tố cốt lõi để hiểu cách các mô hình ngôn ngữ lớn hoạt động. Trong bài viết này, chúng ta sẽ khám phá cách các chiều của tensor thay đổi trong quá trình xử lý tại mỗi lớp của mô hình Transformer, cung cấp cái nhìn sâu sắc về hoạt động của chúng.
Các Điều Kiện Tiên Quyết
Để có trải nghiệm tốt nhất, bạn nên làm quen với các khái niệm cơ bản về hình dạng tensor và phép nhân ma trận. Nếu bạn cần ôn lại, hãy tham khảo bài viết PyTorch Matrix Multiplication.
Thiết lập
Hầu hết các mô hình AI tạo sinh hiện đại đều được xây dựng dựa trên kiến trúc chỉ sử dụng bộ giải mã (decoder-only). Trong bài viết này, chúng ta sẽ đi sâu vào một mô hình tạo văn bản đơn giản, minh họa bằng sơ đồ dưới đây.
Hãy xem xét một ví dụ đầu vào: câu “Hello world !” có thể được tách thành các token: “Hello”, “world” và “!”. Để xử lý trình tự, hai token phụ trợ là <bos> (bắt đầu câu) và <eos> (kết thúc câu) được thêm vào. Điều này giúp dịch chuyển đầu vào một cách chính xác.
Sau khi token hóa, đầu vào trở thành một tensor dạng [12, 15496, 2159, 5145]. Khi đưa vào mô hình theo một batch, một chiều bổ sung sẽ được thêm vào, tạo thành [[12, 15496, 2159, 5145]]. Trong bài viết này, chúng ta sẽ tập trung vào các chiều của tensor, biểu diễn đầu vào dưới dạng [1, 4], trong đó 1 là kích thước batch và 4 là độ dài chuỗi.
Lớp Embedding
Khi tensor đầu vào đi qua lớp embedding, hình dạng của nó thay đổi thành [1, 4, 768]. Con số 768 ở đây đại diện cho chiều embedding.
Lớp embedding đóng vai trò quan trọng trong kiến trúc này vì:
- Chiều embedding sẽ lan truyền qua mạng nơ-ron và được sử dụng nhiều trong lớp attention.
- Nó chuyển đổi các token thành các vector có chiều cao, nắm bắt các mối quan hệ ngữ nghĩa giữa các từ. Ví dụ, các token như “king” và “man” có thể có các giá trị số khác nhau (ví dụ: 8848 và 9584), nhưng biểu diễn vector của chúng trong không gian chiều cao sẽ cho thấy sự tương đồng có ý nghĩa.
Mã hóa Vị trí (Positional Encoding)
Lớp này không làm thay đổi chiều của tensor mà tiêm thông tin vị trí vào đầu vào. Điều này rất quan trọng vì đầu vào sẽ được xử lý song song ở các giai đoạn sau của kiến trúc. Mã hóa vị trí đảm bảo mô hình giữ lại thông tin về thứ tự của các token trong chuỗi.
Lớp Bộ giải mã (Decoder Layer)
Một mô hình tạo sinh thường bao gồm nhiều lớp bộ giải mã liên tiếp, mỗi lớp chứa:
- Một lớp attention đa đầu có mặt nạ (masked multi-head attention).
- Một thao tác cộng và chuẩn hóa (add and normalize).
- Một mạng feed-forward.
Attention Đa Đầu Có Mặt Nạ (Masked Multi-Head Attention)
Lớp Multi-Head Attention cho phép mô hình tập trung vào các phần khác nhau của đầu vào, gán trọng số và biểu diễn từng token trong ngữ cảnh của toàn bộ chuỗi. Masked Multi-Head Attention giới hạn mỗi token chỉ chú ý đến chính nó và các token đứng trước, đạt được bằng cách áp dụng mặt nạ cho các token tương lai trong trọng số attention.
Dữ liệu đầu tiên đi qua ba lớp tuyến tính song song, mỗi lớp có nn.Linear(768, 768), giữ nguyên hình dạng đầu vào. Đầu ra được gọi là Query (Q), Key (K) và Value (V), mỗi loại có hình dạng tensor là [1, 4, 768].
Sau đó, chiều embedding được chia theo 768 = 8 * 96, trong đó 8 là số lượng đầu và 96 là kích thước đầu. Điều này định hình lại tensor thành [1, 4, 8, 96]. Để căn chỉnh các chiều cho phép nhân ma trận, độ dài chuỗi và kích thước đầu được hoán vị, tạo ra các hình dạng [1, 8, 4, 96] cho Q, K và V.
Cơ chế attention, như được định nghĩa trong bài báo gốc, được tính như sau:
$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right) V $$
- Đầu tiên, $QK^\top$ được tính toán:
- $K^\top$ (Key chuyển vị) có hình dạng
[1, 8, 96, 4]. - $QK^\top$ tạo ra tensor có hình dạng
[1, 8, 4, 4], được tính bằng[1, 8, 4, 96] × [1, 8, 96, 4].
- $K^\top$ (Key chuyển vị) có hình dạng
- MẶT NẠ (MASKING): Mặt nạ đảm bảo mỗi token chỉ chú ý đến chính nó và các token đứng trước, ngăn mô hình truy cập các token tương lai trong quá trình tạo sinh.
- Trọng số Attention: Các trọng số attention được tính bằng:
$$
\text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)
$$
trong đó $\sqrt{d_k} = \sqrt{96}$. Việc chia cho kích thước đầu sẽ ngăn chặn sự chênh lệch lớn về giá trị, và softmax đảm bảo tổng các biểu diễn vector của mỗi token bằng 1. Điều này sẽ làm cho các giá trị
-infbằng 0 trong khi vẫn giữ nguyên các giá trị dương, chỉ ảnh hưởng đến các giá trị của tensor chứ không phải hình dạng của nó. - Tính toán Attention: Các trọng số attention được nhân với các giá trị: $$ \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right) V = [1, 8, 6, 4] * [1, 8, 4, 96] \Rightarrow [1, 8, 6, 96] $$
- Nối (Concat): Số lượng đầu được khôi phục bằng cách hoán vị các chiều:
[1, 8, 6, 96] ⇒ [1, 6, 8, 96] ⇒ [1, 6, 768] - Lớp Chiếu (Projection Layer): Một lớp tuyến tính với
nn.Linear(768, 768)giữ nguyên hình dạng tensor tại[1, 6, 768].
Quan sát: Hình dạng tensor đã trở lại
[1, 6, 768], khớp với hình dạng đầu vào ban đầu trước khi qua cơ chế attention. Sự nhất quán này là cần thiết để đảm bảo khả năng tương thích với các lớp và phép tính tiếp theo trong mô hình.
Cộng và Chuẩn hóa (Add and Normalize)
Bước này bao gồm kết nối bỏ qua (skip connection), trong đó các tensor trước và sau lớp attention được cộng lại và chuẩn hóa. Phép cộng đảm bảo các giá trị tensor được cập nhật thay vì bị thay thế, trong khi chuẩn hóa ngăn chặn sự tăng trưởng theo cấp số nhân của các giá trị. Các thao tác cộng và chuẩn hóa này được áp dụng sau mỗi lớp để duy trì các đặc điểm ban đầu của tensor.
Mạng Feed-Forward
Mạng feed-forward thường bao gồm hai lớp tuyến tính liên tiếp: một lớp mở rộng tensor và một lớp thu hẹp tensor, thường đi kèm với một lớp dropout để điều chuẩn. Các lớp này giới thiệu các phép biến đổi phi tuyến, cho phép mô hình nắm bắt các mẫu phong phú và phức tạp hơn trong các chiều embedding (số lượng đầu và kích thước đầu).
Hệ số mở rộng thường là 1 ⇒ 3, và hệ số thu hẹp là 3 ⇒ 1. Điều này dẫn đến các lớp tuyến tính có cấu trúc như sau:
python
nn.Linear(768, 3 * 768)
nn.Linear(3 * 768, 768)
Hình dạng đầu ra cuối cùng trở về hình dạng đầu vào, [1, 4, 768]. Việc duy trì hình dạng này cho phép áp dụng lớp cộng và chuẩn hóa sau bước feed-forward.
Ngoài ra, vì hình dạng cuối cùng của bộ giải mã khớp với hình dạng đầu vào của nó, nhiều lớp bộ giải mã liên tiếp có thể được xếp chồng lên nhau một cách liền mạch.
Đầu Mô Hình Ngôn Ngữ (Language-Model Head)
Sau khi đi qua một loạt các lớp bộ giải mã, tensor đến một lớp tuyến tính cuối cùng chuyển đổi embed_dim ⇒ vocab_size.
Điều này tạo ra một tensor có hình dạng [1, 4, 9735] (với giả định kích thước từ vựng là 9735), trong đó:
- 1: kích thước batch
- 4: độ dài chuỗi
- 9735: kích thước từ vựng
Việc áp dụng hàm softmax và tính toán tổn thất giữa đầu ra của mô hình và dữ liệu thực tế sẽ tạo ra lỗi (hoặc tổn thất), mà bộ tối ưu hóa sử dụng để cập nhật trọng số của mô hình.
Trong Lớp Attention Đa Đầu Có Mặt Nạ, sự chú ý của mỗi token được tính toán chỉ dựa trên token ĐẦU VÀO hiện tại và các token đứng trước nó. Do đầu vào bị dịch sang phải, việc tạo một token mới đòi hỏi mô hình phải xem xét các biểu diễn vector của tất cả các token trước đó, bao gồm cả mối quan hệ của chúng với các token đứng trước. Cơ chế này là nền tảng cho cách các mô hình AI tạo sinh hoạt động, như minh họa dưới đây:
Transformer và Attention Chéo (Cross-Attention)
Kiến trúc Transformer thường bao gồm cấu trúc bộ mã hóa-bộ giải mã, thường được sử dụng trong các tác vụ mà ngữ cảnh và đầu ra không liên quan trực tiếp, chẳng hạn như dịch thuật. Tuy nhiên, ngày nay kiến trúc chỉ sử dụng bộ giải mã thường được ưa chuộng hơn cho các tác vụ này.
Trong các lớp bộ mã hóa, hình dạng tensor lan truyền tương tự như bộ giải mã, với một điểm khác biệt chính: lớp attention không áp dụng mặt nạ, cho phép mỗi token chú ý đến tất cả các token khác trong chuỗi, cả trước và sau.
Ví dụ, hãy xem xét ngữ cảnh “I am at home” và đích “<bos> je suis à la maison”. Các hình dạng tensor sẽ là:
- Ngữ cảnh (“I am at home”):
[1, 4] - Đích ("
<bos>je suis à la maison"):[1, 6]
Độ dài chuỗi của đầu vào và đầu ra khác nhau, điều này được xử lý bởi lớp attention chéo. Đầu ra của bộ mã hóa có hình dạng [1, 4, 768], trong khi đầu ra từ bộ giải mã của lớp attention đa đầu có mặt nạ là [1, 6, 768]. Ở đây, Key (K) và Value (V) đến từ bộ mã hóa, còn Query (Q) đến từ bộ giải mã.
Bạn nên tự tính toán các chiều tensor để hiểu rõ hơn, nhưng giải pháp được cung cấp dưới đây để xác minh.
$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right) V $$
- Query:
[1, 6, 768] - Key & Value:
[1, 4, 768]
Chia và hoán vị chiều:
- Query:
[1, 8, 6, 96] - Key & Value:
[1, 8, 4, 96]
Tính toán Attention:
- $K^\top$:
[1, 8, 96, 4] - $QK^\top$:
[1, 8, 6, 96] * [1, 8, 96, 4] = [1, 8, 6, 4] - $\text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right) = [1, 8, 6, 4]$
- $\text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right) V = [1, 8, 6, 4] * [1, 8, 4, 96] \Rightarrow [1, 8, 6, 96]$
Nối:
[1, 8, 6, 96] ⇒ [1, 6, 8, 96] ⇒ [1, 6, 768]
Quan sát: Hình dạng tensor đã trở lại
[1, 6, 768], khớp với hình dạng đầu ra của lớp attention đa đầu có mặt nạ. Sự nhất quán này đảm bảo khả năng tương thích với các lớp tiếp theo.
Lời kết
Bài viết này nhằm mục đích cung cấp sự hiểu biết rõ ràng hơn về cách thức hoạt động của các cơ chế attention và cách các hình dạng tensor lan truyền qua kiến trúc Transformer.
Nếu bạn thấy bài viết này hữu ích, hãy cân nhắc upvote 🤗. Để có phản hồi hoặc câu hỏi, vui lòng liên hệ qua các phương thức liên lạc được liệt kê trên danh mục đầu tư của tôi: https://not-lain.github.io/.
Cộng đồng
- p3nGu1nZz: Great write up. The part about feed forward is especially elegant.
- Not Lain: Thanks a lot for the feedback, Glad you loved it!
- Will3087: I’m trying to create a model specifically specialized in fiscal and fiscal issues in Brazil, but I’m having difficulty following a trail, there’s so much information that I don’t know where to start, could you help me with that? By the way, congratulations on the writing, very easy to read and comfortable, thank you for that.
- XuTianyi: Excellent writing — it helped me fully understand the Decoder architecture
- Not Lain: Thanks a lot for the nice words @nada2666 🤗 I also wrote another blogpost about KV Caching which you can find here that might interest you!
Đăng bởi Not Lain, ngày 12 tháng 1 năm 2025