Tìm hiểu Cơ chế Multi-Head Attention trong Transformers

Giới thiệu về cơ chế Self-Attention

Cơ chế attention, cụ thể hơn là cơ chế self-attention, lần đầu được giới thiệu trong bài báo nổi tiếng  “Attention is All You Need“ của Vaswani và các cộng sự vào năm 2017. 

Chính cơ chế này là động lực cho mô hình Transformers, từ đó vận hành một số mô hình mạnh mẽ nhất hiện nay như LLM và VLM. 

Vâng, nó thực sự rất quan trọng!

Nhưng có một điều: nhiều người (tôi trước đây cũng vậy!) tin rằng attention được phát minh vào năm 2017.

Không phải vậy.

Cơ chế attention nguyên thuỷ được giới thiệu lần đầu vào năm 2015, khi Bahdanau và cộng sự giới thiệu nó trong bối cảnh dịch máy. Ý tưởng của họ là cho phép decoder tập trung vào các phần khác nhau của câu đầu vào ở mỗi bước giải mã, thay vì dựa vào một vector ngữ cảnh cố định duy nhất.

Sự thay đổi nhỏ đó tạo ra khác biệt lớn — và gieo mầm cho self-attention và sau này là multi-head attention.

Nhưng lúc này bạn có thể có câu hỏi:

“Thực ra attention có nghĩa là gì?”

Hãy bắt đầu khái niệm này với hai câu đơn giản:

• “The cat sat on the mat.” → Mô hình có thể tập trung attention từ sat tới cả cat và mat.
• “She drank her tea because she was thirsty.” → Attention từ drank trỏ tới cả tea và thirsty, kết nối nguyên nhân và hành động.

Từ các ví dụ này, ta bắt đầu thấy attention cố gắng làm gì. Tất cả là về tập trung; giúp mô hình quyết định những từ nào liên quan đến từ đang xử lý, thường dựa trên nghĩa ngữ nghĩa hoặc quan hệ ngữ pháp.

Một điểm quan trọng cần lưu ý là tất cả các token đều được sử dụng khi tính trọng số và điểm attention (sẽ nói kỹ hơn sau), không chỉ một vài token cụ thể, như thể hiện trong sơ đồ trên. Điều này cũng có nghĩa là sẽ có attention mạnh hơn giữa một số từ so với các từ khác, thể hiện qua độ đậm của các mũi tên.

Đây là lý do transformers rất giỏi trong việc hiểu ngôn ngữ; chúng dùng attention để liên hệ động mọi từ với mọi từ khác.

Các bước trung gian

Giờ ta đã hiểu attention là gì, nhưng vẫn chưa biết xử lý nội bộ diễn ra thế nào.

Đây là lúc Query, Key và Value (thường viết tắt là Q, K và V) xuất hiện. Phần còn lại của bài, chúng ta sẽ gắn với câu ví dụ “the cat sat on the mat”. 

Trước hết, khi một câu được đưa vào Transformer, nó lập tức được token hóa bởi một tokenizer; một tokenizer nổi tiếng là Byte-pair encoding. 

Có thể coi tokenization là việc tách văn bản thành các đơn vị nhỏ hơn. Vì các từ trong câu ví dụ đều ngắn, nên ở đây ta giả định mỗi đơn vị tương ứng một từ. Trực quan, câu sau khi đi qua tokenizer sẽ trông như sau:

Lưu ý cả dấu chấm câu (dấu chấm kết câu) cũng được chuyển thành một Token. Một lần nữa, để đơn giản, ta giả định dấu chấm là một token riêng. 

Sau khi token hóa, ta dùng từ vựng của tokenizer để chuyển mỗi token thành Token ID tương ứng. Có thể hình dung từ vựng của tokenizer như một cơ sở dữ liệu/bảng lưu tất cả token mà mô hình có thể nhận diện và Token ID tương ứng.

Tiếp đó, mỗi token ID duy nhất này được ánh xạ sang vector nhúng (embedding) tương ứng thông qua một ma trận embedding. Quan trọng là tất cả vector embedding này đều có cùng số chiều; trong ví dụ, ta giả định là 512.

Bước cuối cùng là bổ sung positional embedding để mô hình có thông tin vị trí của từng token. Hai phương pháp phổ biến là:

1. Mã hoá hình sin (Sinusoidal) (dùng trong Transformer gốc)
2. Positional embedding học được (dùng trong BERT, GPT)

Cơ chế nền tảng của các positional embedding này vượt ngoài phạm vi hướng dẫn. Hiện tại, điều cần biết là ta thu được vector đầu vào qua công thức:

Vector đầu vào = Token Embedding + Positional Embedding

Query, key và value

Giờ chúng ta sẵn sàng giới thiệu các vector Q,K,V. Bắt đầu bằng việc lưu ý: với mỗi từ ở đầu vào, ta tạo ba vector:

• Query (Q) – điều mà từ đó đang tìm kiếm
• Key (K) – điều mà từ đó cung cấp cho các từ khác
• Value (V) – chính là thông tin mà từ đó nắm giữ

Đây là các biến đổi có thể học được và cũng là các ma trận tương ứng. Về mặt toán học, các biến đổi này được viết như sau:

Từ đó, mỗi token có thể được biểu diễn bằng ba vector mới. Tôi đã tạo sơ đồ dưới đây để trực quan hoá quá trình này:

Tính toán scaled dot-product 

Trong bài “Attention is all you need”, một công thức toán học được đưa ra, cho thấy cách các vector mới được dùng để tính Trọng số Attention. Hãy đi từng bước, giả sử ta đang tính Trọng số Attention khi ở từ “sat”:

Đầu tiên, ta dùng vector Query của “sat” để so sánh với các vector Key của mọi từ khác — bao gồm cả chính nó. Kết quả là một tập các điểm tương đồng (các vô hướng) cho biết mỗi từ liên quan tới “sat” như thế nào. Chúng còn được gọi chính thức là Điểm Attention.

Các điểm này sau đó được chia tỷ lệ qua sqrt(d_k) để tránh giá trị quá lớn và giúp ổn định gradient khi huấn luyện bằng lan truyền ngược. Các điểm đã chia tỷ lệ được đưa qua softmax, biến chúng thành các trọng số attention. 

Softmax cho ta một xác suất cho biết mô hình nên phân bổ bao nhiêu trọng số cho mỗi từ, do đó tổng trọng số attention cho một token cụ thể sẽ bằng 1. 

Cuối cùng, ta lấy tổng có trọng số của các vector value, sử dụng các trọng số attention này. 

Điều này cho ta một vector mới biểu diễn “sat” trong ngữ cảnh, được làm giàu bởi thông tin của tất cả token, nhưng chủ yếu từ cat, mat và các token liên quan khác.

Quá trình này diễn ra cho mọi token trong câu. Mỗi token chú ý tới mọi token khác, thậm chí cả chính nó. Đó là lý do ta gọi là self-attention.

Cũng lưu ý rằng ta vừa tính vector làm giàu cho một token, nhưng trên thực tế, ta làm điều này cho nhiều, nếu không muốn nói là tất cả, token cùng lúc, vì phép nhân ma trận cực kỳ hiệu quả. 

Ta làm được vì kiến trúc này hoàn toàn có thể song song hoá. Khác với mạng nơ-ron hồi quy (RNN) xử lý tuần tự, cơ chế này cho phép Transformers xử lý tất cả token cùng lúc, khiến chúng nhanh và mạnh.

Để trực quan hơn, hãy xem sơ đồ trên. Bạn có thể thấy ma trận đầu vào X được biến đổi thành các ma trận Query, Key và Value tương ứng, rồi công thức scaled dot-product attention được áp dụng để thu được các vector embedding làm giàu đầy đủ của mọi token trong đầu vào — A.

Hiện thực hoá scaled dot-product attention trong PyTorch

Giờ hãy viết mã những gì ta vừa thảo luận trong PyTorch. Bất ngờ là khá dễ:

import torchimport torch.nn.functional as Fdef scaled_dot_product_attention(Q, K, V):    d_k = Q.size(-1)  # the dimension of the key vectors    scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))    attn_weights = F.softmax(scores, dim=-1)    output = torch.matmul(attn_weights, V)    return output, attn_weightsĐược hỗ trợ bởi

Bằng cách hiện thực một hàm đơn giản, ta có thể tái tạo phần thảo luận trên! Dưới đây là những điểm quan trọng nhất:

• Q, K và V là các tensor có dạng [batch_size, seq_len, dim]
• Đầu ra là một tập vector mới trong đó mỗi token được biểu diễn trong ngữ cảnh của toàn bộ chuỗi
• Attention Scores là các điểm tương đồng thô, chưa chuẩn hoá
• Attention Weights là các giá trị đã chuẩn hoá (qua softmax) để tập trung và có tổng bằng 1

Các loại Attention khác nhau

Tuy nhiên, có một chi tiết quan trọng tôi chưa nói. Không chỉ có một loại attention trong Transformer — thực ra có ba loại:

1. Encoder self-attention
2. Decoder self-attention
3. Encoder–decoder (cross) attention

Mỗi loại đóng một vai trò hơi khác nhau, hãy điểm qua nhanh.

Encoder self-attention

Đây là những gì ta đã khám phá, trong encoder, mỗi token chú ý tới mọi token khác trong cùng chuỗi đầu vào, như ví dụ “the cat sat on the mat.” 

Điểm quan trọng ở đây là không có hạn chế nào vì toàn bộ đầu vào đã biết trước, cho phép mô hình xây dựng biểu diễn ngữ cảnh sâu cho từng token, hiểu được trước và sau nó là gì. 

Cũng lưu ý rằng điều này diễn ra trong encoder. 

Decoder self-attention

Đây là nơi mọi thứ khó hơn một chút. Vì ta đang ở decoder, ta không muốn mỗi token nhìn thấy tương lai, vì như vậy là gian lận khi sinh văn bản.

Lấy ví dụ khác: “The dog barked very loudly.” Khi mô hình đang sinh “barked” (giả sử mỗi từ là một token cho đơn giản), nó không nên biết token tiếp theo là “very”, nếu không mô hình sẽ không học cách sinh token mới. 

Để giải quyết, ta áp dụng mặt nạ nhân quả (causal mask) (đôi khi gọi là mặt nạ tam giác) để đảm bảo mỗi token chỉ có thể chú ý tới chính nó và các token phía trước, không bao giờ là những token phía sau.

Từ sơ đồ trên, ta thấy các điểm attention hướng về tương lai được loại bỏ bằng cách masking, cho phép mô hình học cách sinh ra các token mới, chính xác và hợp lý.

Encoder-decoder (cross) attention

Khi decoder bắt đầu sinh, nó vẫn cần thông tin từ đầu vào gốc (thứ mà nó đang phản hồi). Đây là lúc cross attention xuất hiện.

Ở đây, decoder chú ý tới đầu ra của encoder. Vì vậy trong khi decoder self-attention được che mặt nạ, cross-attention thì không, và decoder được phép nhìn tất cả token của encoder cùng lúc.

Khi học khái niệm này, tôi rất thích ví dụ Dịch thuật. Giả sử bạn đang dịch “Bonjour tout le monde” sang tiếng Anh, decoder nên có thể nhìn tất cả từ tiếng Pháp khi quyết định mỗi từ tiếng Anh, đúng không? 

Đó chính xác là những gì cross attention làm. Tuy nhiên, chúng ta sẽ không đi sâu hơn về cross attention ở đây, vì nó không cần thiết trong phạm vi bài viết này.

Nhu cầu về Multi-Head Attention

Giờ ta đã hiểu self-attention hoạt động ra sao. Nhưng bạn có thể hỏi:

Nếu đã có attention, tại sao ta cần nhiều đầu attention?

Với điều này, tôi xin đưa ra phép ẩn dụ đèn pin, đã giúp tôi xây dựng trực giác khi mới học.

Hãy nghĩ một attention head như một chiếc đèn pin cố gắng soi rọi những phần quan trọng của câu. Nhưng một câu có nhiều loại quan hệ khác nhau cần hiểu: chủ ngữ-–động từ, tân ngữ-–đại từ, tính từ-–danh từ, phụ thuộc xa, v.v.

Chẳng phải sẽ tốt hơn nếu ta có thể nắm bắt nhiều quan hệ thay vì chỉ một quan hệ mỗi lần sao?

Vì vậy, giải pháp là đưa cho mô hình nhiều đèn pin, mỗi chiếc tập trung vào các phần khác nhau của câu. Đó là multi-head attention.

Thay vì tính một bộ Q, K, V và chạy một phép attention duy nhất, ta chia embedding thành các phần nhỏ hơn và chạy nhiều head độc lập song song. 

Như thấy trong sơ đồ trên, embedding đầu vào (trong trường hợp này cho một token, để đơn giản) được chia đều cho n head khác nhau. 

Trong ví dụ, giả sử số chiều là 512 và số head là 8, do đó mỗi head sẽ xử lý song song (512/8=64) chiều song song, cố gắng nắm bắt các quan hệ khác nhau. 

Để xây dựng trực giác hơn về những loại quan hệ mà các head có thể (tôi nói “có thể” vì ta vẫn chưa hiểu đầy đủ mỗi head nhìn vào quan hệ nào) xử lý, hãy đi qua một ví dụ đơn giản:

Xét câu sau:

“The quick brown fox jumps over the lazy dog.”

• Một head có thể tập trung vào cặp tính từ–danh từ: “quick → fox”, “lazy → dog”
• Head khác có thể tập trung vào chủ ngữ–động từ: “fox → jumps”
• Và head khác có thể nắm bắt quan hệ xa: “jumps → over → dog”

Lưu ý chính là tất cả các head này chạy đồng thời, và đầu ra của chúng được nối (concatenate) và chiếu trở lại kích thước embedding gốc (trong trường hợp ta là 512). 

Trực quan hoá nhiều head qua mã nguồn

Để củng cố hiểu biết về Multi-Head Attention, hãy dùng thư viện HuggingFace để xem cơ chế Multi-Head Attention hoạt động thế nào. 

!pip install transformers bertviz torchĐược hỗ trợ bởi

Trước tiên, chạy ô lệnh trên để tải các gói liên quan.

from transformers import BertTokenizer, BertModelfrom bertviz import head_viewimport torchmodel_name = 'bert-base-uncased' # We will be using the Bert modelmodel = BertModel.from_pretrained(model_name, output_attentions=True)tokenizer = BertTokenizer.from_pretrained(model_name)model.eval() # As we are not training, we have set the model to evaluation mode so no updates are made# This is our input sentencesentence = "The quick brown fox jumps over the lazy dog."# Using the tokenizer to convert the sentence into tokensinputs = tokenizer(sentence, return_tensors='pt')input_ids = inputs['input_ids']# Forward pass with attentionswith torch.no_grad():    outputs = model(**inputs)    attentions = outputs.attentions  # Tuple of (num_layers, batch, num_heads, seq_len, seq_len)# Here we decode the tokens tokens = tokenizer.convert_ids_to_tokens(input_ids[0])head_view(attentions, tokens) # VisualizationĐược hỗ trợ bởi

Đầu ra bạn nhận được như bên dưới. Hãy để ý có 12 màu khác nhau, nghĩa là có 12 Head trong mô hình này. Độ dày của các đường cũng tương ứng với độ lớn điểm attention giữa các token. 

Hiện thực Multi-Head Attention: Thực hành tốt

Trong phần này, ta sẽ viết giả mã (pseudocode) cho cơ chế Multi-Head Attention. Giả sử đầu vào: X có dạng [batch, seq_len, d_model]. 

Sau đó ta tiến hành giả mã: 

Q = X @ W_Q # Calculating the Query valuesK = X @ W_K # Calculating the Key valuesV = X @ W_V# Calculating the Value valuesQ_i, K_i, V_i = split_into_heads(Q, K, V, num_heads) # Now we split the Q,K,V into multiple different heads (by evenly splitting the dimensions)# Here we compute scaled dot-product attention for each headscores = (Q_i @ K_i.T) / sqrt(head_dim)scores += maskA_i = softmax(scores)A_i = dropout(A_i)Z_i = A_i @ V_iZ = concat_heads(Z_i) # We have to concatenate all of the Heads# This is our final linear projection matrixoutput = Z @ W_Ooutput = dropout(output)Được hỗ trợ bởi

Hãy điểm qua nhanh. 

Sau khi mọi head attention riêng lẻ hoàn tất, đầu ra của chúng được nối lại thành một vector duy nhất. Nhưng đầu ra kết hợp này vẫn cần được đưa về kích thước mô hình gốc và đây là nơi ma trận chiếu tuyến tính cuối cùng xuất hiện. Nó đơn giản là ánh xạ đầu ra đã nối về cùng kích thước với embedding đầu vào, để có thể truyền sang lớp tiếp theo trong Transformer.

Cùng thảo luận một vài tối ưu hoá và thực hành tốt:

1. Song song hoá: Vì mỗi head độc lập, ta có thể tính tất cả song song.
2. Nhân ma trận: Ta gộp theo lô mọi thứ bằng các phép đại số tuyến tính (ví dụ matmul), vốn được tối ưu cao.
3. FlashAttention (Nâng cao): Kỹ thuật mới hơn dùng thủ thuật CUDA cấp thấp để giảm bộ nhớ và tăng tốc attention — đặc biệt với chuỗi dài.

Về thực hành tốt, khi viết lớp attention, có 3 siêu tham số chính cần lưu ý:

1. num_heads: Số head attention chạy song song, cần đảm bảo chia hết vào d_model.
2. head_dim: Số chiều mỗi head.
3. dropout_rate: Để tránh overfitting.

Kết luận

Chúng ta đã bao quát nhiều nội dung. Bắt đầu với ý tưởng đơn giản: attention là về sự tập trung. Sau đó, ta khám phá cách Transformers dùng sự tập trung này bằng cách tính mối quan hệ giữa mỗi từ với mọi từ khác thông qua Query, Key và Value.

Ta đã thấy self-attention (và lướt qua các dạng Attention khác như Causal và Cross) giúp mỗi token có biểu diễn giàu ngữ cảnh. Rồi ta tiến thêm một bước với multi-head attention, nơi thay vì chỉ dựa vào một “đèn pin”, ta đưa cho mô hình nhiều “đèn pin”, mỗi cái nắm bắt các mẫu khác nhau trong câu, đồng thời.

Đây là điều tạo nên sức mạnh thực sự của Transformers. Chúng không chỉ đọc; chúng hiểu cấu trúc, ý nghĩa và các mối quan hệ xuyên suốt toàn bộ chuỗi. 

Trong bài này, ta đã đề cập một khía cạnh của Transformers, nhưng vẫn còn nhiều điều để học, hãy xem khoá học chi tiết đầy đủ này: Khoá học Transformer Models with PyTorch. 

Nếu bạn muốn đi vào ứng dụng, tôi cũng khuyên thử dự án Developing Multi-Input Models For OCR hoặc khoá học Working with HuggingFace. Nếu bạn thích bài viết, hãy xem thêm Hướng dẫn Hoàn chỉnh Xây dựng Transformer bằng PyTorch và bài viết khác của tôi về Vision Transformers.