Embedding và Mô hình Chuỗi¶

Word embedding nén văn bản tượng trưng thưa thớt thành các không gian vector dày đặc, nơi tương đồng ngữ nghĩa trở thành gần gũi hình học. File này bao gồm Word2Vec (CBOW, Skip-gram), GloVe, FastText, RNN, LSTM, GRU, seq2seq với cơ chế chú ý (attention), và mô hình mã hoá-giải mã (encoder-decoder) — sự tiến triển từ bag-of-words đến các biểu diễn ngữ cảnh.

Ở file 01, chúng ta đã giới thiệu giả thuyết phân bố: các từ xuất hiện trong những ngữ cảnh tương tự có xu hướng mang ý nghĩa tương tự. Ở file 02, chúng ta biểu diễn văn bản bằng các đặc trưng thưa, được xây dựng thủ công như vector TF-IDF. Các vector này sống trong không gian rất nhiều chiều (một chiều cho mỗi từ vựng) và hầu hết là số không. Word embedding nén thông tin này thành các vector dày đặc, ít chiều nắm bắt các quan hệ ngữ nghĩa, và chúng được học trực tiếp từ dữ liệu.
Word2Vec (Mikolov và cộng sự, 2013) học word embedding bằng cách huấn luyện một mạng nơ-ron nông trên một tác vụ dự đoán đơn giản. Có hai kiến trúc.
Continuous Bag of Words (CBOW) dự đoán một từ đích từ các từ ngữ cảnh xung quanh nó. Với một cửa sổ các từ ngữ cảnh (ví dụ: "the cat ___ on the"), mô hình lấy trung bình các vector embedding của chúng và đưa kết quả qua một lớp tuyến tính để dự đoán từ còn thiếu ("sat"). Mục tiêu huấn luyện là cực đại hoá:

\[P(w_t \mid w_{t-k}, \ldots, w_{t-1}, w_{t+1}, \ldots, w_{t+k})\]

Skip-gram làm ngược lại: từ một từ đích, dự đoán các từ ngữ cảnh xung quanh. Với từ đích "sat", mô hình cố gắng dự đoán "the", "cat", "on", "the" trong các dự đoán riêng biệt. Mục tiêu cực đại hoá:

\[P(w_{t+j} \mid w_t) \quad \text{với mỗi } j \in [-k, k], \; j \neq 0\]

Skip-gram thường hoạt động tốt hơn cho các từ hiếm vì mỗi từ tạo ra nhiều mẫu huấn luyện (một mẫu cho mỗi vị trí ngữ cảnh). CBOW nhanh hơn và tốt hơn một chút cho các từ phổ biến vì nó lấy trung bình nhiều tín hiệu ngữ cảnh.
Huấn luyện trên toàn bộ từ vựng rất tốn kém vì mẫu số softmax tính tổng trên tất cả \(V\) từ. Lấy mẫu âm (Negative sampling) xấp xỉ điều này bằng cách chuyển bài toán thành phân loại nhị phân: phân biệt từ ngữ cảnh thật (mẫu dương) với các từ nhiễu được lấy mẫu ngẫu nhiên (mẫu âm). Thay vì tính softmax đầy đủ, mô hình chỉ cập nhật embedding cho từ đích, từ ngữ cảnh thật và một số ít mẫu âm:

\[\mathcal{L} = \log \sigma(v_{w_O}^T v_{w_I}) + \sum_{i=1}^{k} \mathbb{E}_{w_i \sim P_n} [\log \sigma(-v_{w_i}^T v_{w_I})]\]

Ở đây \(v_{w_I}\) là embedding của từ đầu vào, \(v_{w_O}\) là embedding của từ đầu ra (ngữ cảnh), và \(P_n\) là phân bố nhiễu, thường là tần suất unigram được nâng lên luỹ thừa 3/4 (làm giảm trọng số của các từ rất phổ biến như "the").
Tại sao mục tiêu đơn giản này lại tạo ra các embedding có ý nghĩa? Levy và Goldberg (2014) đã chỉ ra rằng skip-gram với negative sampling ngầm phân rã một ma trận thông tin tương hỗ từng điểm (PMI) đã dịch chuyển. Tại hội tụ, tích vô hướng của hai vector từ xấp xỉ:

\[v_w^T v_c \approx \text{PMI}(w, c) - \log k\]

trong đó \(\text{PMI}(w, c) = \log \frac{P(w, c)}{P(w) P(c)}\) đo lường mức độ đồng xuất hiện của hai từ \(w\) và \(c\) nhiều hơn kỳ vọng ngẫu nhiên (lý thuyết thông tin chương 05), và \(k\) là số mẫu âm. Các từ đồng xuất hiện nhiều hơn ngẫu nhiên có PMI cao và do đó tích vô hướng cao (embedding tương tự). Các từ đồng xuất hiện ít hơn kỳ vọng có PMI âm và embedding không tương tự. Điều này cho thấy Word2Vec đang làm cùng một việc với các phương pháp ngữ nghĩa phân bố kinh điển như phân tích ngữ nghĩa tiềm ẩn (SVD trên ma trận đồng xuất hiện), nhưng theo cách trực tuyến và khả mở rộng hơn.
Tính chất đáng ngạc nhiên nhất của Word2Vec embedding là chúng nắm bắt được phép loại suy (analogies) thông qua số học vector. Vector \(v_{\text{king}} - v_{\text{man}} + v_{\text{woman}}\) gần nhất với \(v_{\text{queen}}\). Điều này hoạt động vì không gian embedding mã hoá các quan hệ ngữ nghĩa như các hướng xấp xỉ tuyến tính: hướng "hoàng gia" xấp xỉ là \(v_{\text{king}} - v_{\text{man}}\), và cộng nó với \(v_{\text{woman}}\) đưa ta đến gần \(v_{\text{queen}}\). Điều này kết nối với đại số tuyến tính ở chương 01: các quan hệ ngữ nghĩa là các phép tịnh tiến vector.
GloVe (Global Vectors for Word Representation, Pennington và cộng sự, 2014) có cách tiếp cận khác. Thay vì học từ các cửa sổ ngữ cảnh địa phương từng cái một, nó xây dựng một ma trận đồng xuất hiện từ toàn cục \(X\), trong đó \(X_{ij}\) đếm tần suất từ \(j\) xuất hiện trong ngữ cảnh của từ \(i\) trên toàn bộ kho ngữ liệu. Mô hình sau đó học các embedding có tích vô hướng xấp xỉ log của đồng xuất hiện:

\[w_i^T \tilde{w}_j + b_i + \tilde{b}_j = \log X_{ij}\]

Hàm mất mát gán trọng số cho mỗi cặp bằng một hàm giới hạn \(f(X_{ij})\) ngăn các đồng xuất hiện rất phổ biến lấn át:

\[\mathcal{L} = \sum_{i,j=1}^{V} f(X_{ij}) \left(w_i^T \tilde{w}_j + b_i + \tilde{b}_j - \log X_{ij}\right)^2\]

GloVe kết hợp lợi ích của phân rã ma trận toàn cục (như phân tích ngữ nghĩa tiềm ẩn) với học ngữ cảnh địa phương của Word2Vec. Trong thực tế, GloVe và Word2Vec tạo ra các embedding chất lượng tương tự.
FastText (Bojanowski và cộng sự, 2017) mở rộng skip-gram bằng cách biểu diễn mỗi từ như một túi các n-gram ký tự. Từ "where" với \(n = 3\) trở thành: "", cộng với token cả từ "". Embedding của từ là tổng của tất cả các embedding n-gram của nó.
Điều này có một lợi thế quan trọng: FastText có thể tạo ra embedding cho các từ chưa từng thấy trong huấn luyện. Từ "whereabouts" chia sẻ n-gram với "where", vì vậy embedding của nó sẽ hợp lý ngay cả khi "whereabouts" chưa bao giờ xuất hiện trong dữ liệu huấn luyện. Điều này đặc biệt hữu ích cho các ngôn ngữ giàu hình thái (file 01) nơi từ có nhiều dạng biến tố.
Đánh giá embedding thường dùng hai loại chuẩn. Bài toán loại suy (Analogy tasks) kiểm tra xem \(v_a - v_b + v_c \approx v_d\) (ví dụ: "Paris" \(-\) "France" \(+\) "Italy" \(\approx\) "Rome"). Chuẩn tương tự (Similarity benchmarks) so sánh độ tương tự cosine (chương 01) giữa các cặp từ với đánh giá của con người. Các tập dữ liệu phổ biến bao gồm WordSim-353, SimLex-999 và bộ kiểm tra loại suy Google. Một lưu ý thực tế: embedding giỏi loại suy có thể không tốt nhất cho các tác vụ hạ nguồn như phân loại cảm xúc. Cách đánh giá tốt nhất thường là chính tác vụ đó.
Trong chương 06, chúng ta đã giới thiệu RNN, LSTM và GRU như các kiến trúc cho dữ liệu chuỗi. Ở đây chúng ta tập trung vào cách chúng được áp dụng cụ thể cho các tác vụ ngôn ngữ.
RNN mô hình ngôn ngữ (Language model RNN) đọc từng token một và dự đoán token tiếp theo tại mỗi bước. Trạng thái ẩn \(h_t\) nén toàn bộ lịch sử \(w_1, \ldots, w_t\) thành một vector kích thước cố định, và một lớp tuyến tính cộng softmax ánh xạ \(h_t\) thành phân bố trên từ vựng. Huấn luyện dùng mất mát entropy chéo so với token đúng tiếp theo, tương đương với cực tiểu hoá độ rối (file 02). Hạn chế chính: vector kích thước cố định phải mã hoá toàn bộ lịch sử, và thông tin từ các token đầu dần bị ghi đè.
RNN hai chiều (Bidirectional RNN) xử lý chuỗi theo cả hai hướng: một RNN đọc từ trái sang phải, một RNN khác đọc từ phải sang trái. Tại mỗi vị trí \(t\), trạng thái ẩn tiến \(\overrightarrow{h}_t\) và trạng thái ẩn lùi \(\overleftarrow{h}_t\) được nối lại để tạo thành biểu diễn có nhận thức ngữ cảnh \(h_t = [\overrightarrow{h}_t ; \overleftarrow{h}_t]\). Điều này cho mô hình truy cập cả ngữ cảnh quá khứ và tương lai, rất mạnh cho các tác vụ như POS tagging và NER (file 02) nơi nhãn của một từ phụ thuộc vào các từ cả trước và sau nó. RNN hai chiều không thể dùng cho mô hình ngôn ngữ vì bạn không thể nhìn trước các token tương lai khi dự đoán chúng.
RNN xếp chồng sâu (Deep stacked RNN) đặt nhiều lớp RNN chồng lên nhau. Trạng thái ẩn của lớp \(l\) tại mọi bước thời gian trở thành chuỗi đầu vào cho lớp \(l + 1\). Xếp chồng 2-4 lớp thường cải thiện hiệu suất bằng cách xây dựng các biểu diễn phân cấp, tương tự cách CNN sâu hơn xây dựng hệ thống phân cấp đặc trưng (chương 06). Trên 4 lớp, gradient biến mất và quá khớp trở thành vấn đề trừ khi thêm các kết nối dư (residual connections) giữa các lớp.
Kiến trúc chuỗi-tới-chuỗi (Sequence-to-sequence - seq2seq) (Sutskever và cộng sự, 2014) ánh xạ một chuỗi đầu vào có độ dài thay đổi thành một chuỗi đầu ra có độ dài thay đổi. Nó gồm một RNN mã hoá (encoder) đọc đầu vào và nén nó thành một vector ngữ cảnh (trạng thái ẩn cuối cùng), và một RNN giải mã (decoder) sinh đầu ra từng token một, dựa trên vector ngữ cảnh này.
Seq2seq là kiến trúc đột phá cho dịch máy. Encoder đọc một câu tiếng Pháp, decoder tạo ra bản dịch tiếng Anh. Decoder bắt đầu với một token đặc biệt đánh dấu bắt đầu chuỗi và sinh token tự hồi quy (autoregressively) cho đến khi nó tạo ra token kết thúc chuỗi. Một mẹo thực tế: đảo ngược chuỗi đầu vào (đưa "chat le" thay vì "le chat") đã cải thiện kết quả vì nó đặt từ đầu vào đầu tiên gần hơn với từ đầu ra đầu tiên trong đồ thị tính toán, rút ngắn đường gradient.
Vấn đề nút thắt cổ chai: toàn bộ đầu vào phải được nén vào một vector kích thước cố định duy nhất. Với các câu dài, vector này không thể nắm bắt mọi thông tin, và hiệu suất suy giảm. Điều này đã thúc đẩy cơ chế chú ý (attention mechanisms).
Chương 06 đã giới thiệu công thức Q, K, V hiện đại của attention. Các cơ chế attention ban đầu cho NLP được xây dựng khác, như các mô hình căn chỉnh (alignment) giữa trạng thái encoder và decoder.
Bahdanau attention (additive attention, Bahdanau và cộng sự, 2015) tính điểm căn chỉnh giữa trạng thái ẩn decoder \(s_t\) và mỗi trạng thái ẩn encoder \(h_i\) bằng một mạng truyền thẳng được học:

\[e_{ti} = v^T \tanh(W_s s_{t-1} + W_h h_i)\]

Các điểm số được chuẩn hoá thành trọng số attention qua softmax, và vector ngữ cảnh là tổng có trọng số của các trạng thái encoder:

\[\alpha_{ti} = \frac{\exp(e_{ti})}{\sum_j \exp(e_{tj})}, \quad c_t = \sum_i \alpha_{ti} h_i\]

Decoder sau đó dùng cả \(s_{t-1}\) và \(c_t\) để tạo đầu ra tiếp theo. Ý tưởng chính: thay vì một vector ngữ cảnh cố định cho toàn bộ câu, mỗi bước decoder nhận một tổ hợp có trọng số khác nhau của các trạng thái encoder, cho phép mô hình "nhìn lại" các phần liên quan của đầu vào.
Luong attention (multiplicative attention, Luong và cộng sự, 2015) đơn giản hoá việc tính điểm. Biến thể dot dùng \(e_{ti} = s_t^T h_i\). Biến thể general dùng \(e_{ti} = s_t^T W h_i\). Các biến thể này nhanh hơn attention dạng cộng của Bahdanau vì chúng dùng phép nhân ma trận thay vì mạng truyền thẳng. Luong attention cũng tính vector ngữ cảnh từ trạng thái decoder hiện tại \(s_t\) (thay vì \(s_{t-1}\)), cho nó truy cập nhiều thông tin hơn nhưng cách tính hơi khác.
Trọng số attention thường được trực quan hoá như bản đồ nhiệt cho thấy token đầu vào nào decoder tập trung vào khi tạo mỗi token đầu ra. Trong dịch thuật, các bản đồ nhiệt này phác thảo căn chỉnh từ giữa ngôn ngữ nguồn và đích, với mẫu đường chéo bị phá vỡ bởi sự đổi trật tự (ví dụ: trật tự tính từ-danh từ khác nhau giữa tiếng Pháp và tiếng Anh).
Tại thời điểm suy luận, decoder phải chọn một token ở mỗi bước. Giải mã tham lam (Greedy decoding) chọn token có xác suất cao nhất tại mỗi vị trí, nhưng điều này có thể dẫn đến các chuỗi dưới tối ưu: một lựa chọn tốt cục bộ có thể buộc mô hình vào một câu tồi toàn cục. Tìm kiếm chùm (Beam search) duy trì \(k\) chuỗi bộ phận tốt nhất (độ rộng chùm) tại mỗi bước, mở rộng mỗi chuỗi bằng mọi token tiếp theo có thể và giữ lại \(k\) chuỗi tốt nhất tổng thể.
Với độ rộng chùm \(k = 1\), beam search trở về giải mã tham lam. Giá trị điển hình là \(k = 4\) đến \(k = 10\). Chùm lớn hơn tìm được chuỗi tốt hơn nhưng chậm hơn tương ứng. Beam search cũng cần chuẩn hoá độ dài (length normalisation) để tránh ưu tiên các chuỗi ngắn hơn vốn có tổng xác suất cao hơn một cách tự nhiên vì chúng nhân ít số hạng hơn. Điểm số đã chuẩn hoá là:

\[\text{score}(y) = \frac{1}{|y|^\alpha} \sum_{t=1}^{|y|} \log P(y_t \mid y_{<t})\]

trong đó \(|y|\) là độ dài chuỗi và \(\alpha\) (thường 0.6-0.7) kiểm soát độ mạnh của phạt độ dài. Với \(\alpha = 0\), không có chuẩn hoá độ dài. Với \(\alpha = 1\), điểm số là log-xác suất trên mỗi token (trung bình nhân). Giá trị trung gian cân bằng giữa ưu tiên đầu ra cô đọng và không cắt ngắn quá sớm.
Trong khi RNN xử lý văn bản tuần tự, CNN 1D xử lý nó song song bằng cách trượt các bộ lọc qua chuỗi token. Mỗi bộ lọc phát hiện một mẫu địa phương (một đặc trưng n-gram).
TextCNN (Kim, 2014) áp dụng nhiều bộ lọc tích chập 1D với các độ rộng khác nhau (ví dụ: 3, 4, 5 token) lên ma trận embedding đầu vào. Mỗi bộ lọc tạo ra một bản đồ đặc trưng (feature map), và gộp max theo thời gian (max-over-time pooling) lấy giá trị lớn nhất duy nhất từ mỗi bản đồ đặc trưng, nắm bắt liệu mẫu có được phát hiện ở đâu đó trong văn bản bất kể vị trí. Các đặc trưng đã được gộp từ mọi bộ lọc được nối lại và đưa đến bộ phân loại.
TextCNN nhanh và hiệu quả bất ngờ cho các tác vụ phân loại văn bản như phân tích cảm xúc. Nó nắm bắt các mẫu n-gram địa phương nhưng không thể mô hình hoá các phụ thuộc xa: một bộ lọc độ rộng 5 chỉ thấy 5 token liên tiếp. Tích chập nhân quả giãn (Dilated causal convolutions) giải quyết điều này bằng cách chèn các khoảng trống (dilation) giữa các phần tử bộ lọc. Xếp chồng các lớp với tỷ lệ giãn tăng theo cấp số nhân (1, 2, 4, 8, ...) mở rộng trường tiếp nhận theo cấp số nhân mà không tăng tham số, cho phép mô hình nắm bắt các phụ thuộc qua hàng trăm token.
Tất cả các embedding đã thảo luận đến nay (Word2Vec, GloVe, FastText) đều tạo ra một vector duy nhất cho mỗi loại từ bất kể ngữ cảnh. "Bank" nhận cùng một embedding dù nó có nghĩa là ngân hàng hay bờ sông. Đây là một hạn chế cơ bản mà embedding ngữ cảnh (contextual embeddings) giải quyết.
ELMo (Embeddings from Language Models, Peters và cộng sự, 2018) tạo ra biểu diễn từ ngữ cảnh bằng cách chạy một mô hình ngôn ngữ LSTM hai chiều sâu trên văn bản đầu vào. LSTM xuôi dự đoán từ tiếp theo tại mỗi vị trí; một LSTM ngược riêng biệt dự đoán từ trước đó. Cả hai được huấn luyện như mô hình ngôn ngữ trên các kho ngữ liệu lớn.
Tại mỗi vị trí \(k\), ELMo kết hợp các trạng thái ẩn từ tất cả \(L\) lớp bằng các trọng số học được cho từng tác vụ:

\[\text{ELMo}_k = \gamma \sum_{j=0}^{L} s_j \, h_{k,j}\]

Ở đây \(h_{k,j}\) là trạng thái ẩn tại vị trí \(k\) và lớp \(j\) (lớp 0 là embedding token thô), \(s_j\) là các trọng số vô hướng chuẩn hoá softmax, và \(\gamma\) là hệ số tỷ lệ theo tác vụ. Các lớp khác nhau nắm bắt thông tin khác nhau: các lớp dưới nắm bắt cú pháp (nhãn từ loại, hình thái từ), các lớp trên nắm bắt ngữ nghĩa (nghĩa từ, vai trò ngữ nghĩa). Bằng cách trộn tất cả các lớp với trọng số học được, ELMo embedding thích ứng với các tác vụ hạ nguồn đa dạng.
ELMo đánh dấu sự khởi đầu của mô hình tiền huấn luyện rồi tinh chỉnh (pre-train then fine-tune): huấn luyện một mô hình ngôn ngữ lớn trên văn bản không nhãn khổng lồ, sau đó dùng các biểu diễn của nó cho các tác vụ hạ nguồn. Cụ thể, ELMo dùng các biểu diễn đã tiền huấn luyện như các đặc trưng cố định hoặc được tinh chỉnh nhẹ và được nối với đầu vào theo tác vụ. BERT và GPT (file 04) đẩy xa hơn bằng cách tinh chỉnh toàn bộ mô hình từ đầu đến cuối, chứng tỏ hiệu quả mạnh mẽ hơn đáng kể.
Sự tiến triển từ Word2Vec đến ELMo minh hoạ một chủ đề lặp lại trong NLP: chuyển từ biểu diễn tĩnh sang động, từ ngữ cảnh địa phương sang toàn cục, và từ mô hình nông sang sâu. Mỗi bước đánh đổi chi phí tính toán lấy biểu diễn giàu hơn. Transformer (file 04) hoàn thiện sự tiến triển này bằng cách thay thế hoàn toàn tính tuần tự bằng cơ chế chú ý, cho phép cả ngữ cảnh hoá sâu và tính toán song song.

Bài tập Lập trình (dùng CoLab hoặc notebook)¶

Triển khai Word2Vec skip-gram với negative sampling từ đầu. Huấn luyện trên một kho ngữ liệu nhỏ và trực quan hoá các embedding đã học bằng PCA.

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

# Kho ngữ liệu nhỏ
corpus = """the king ruled the kingdom . the queen ruled the kingdom .
the prince is the son of the king . the princess is the daughter of the queen .
a man worked in the castle . a woman worked in the castle .
the king and queen lived in the castle . the prince and princess played outside .""".lower().split()

vocab = sorted(set(corpus))
word2idx = {w: i for i, w in enumerate(vocab)}
idx2word = {i: w for w, i in word2idx.items()}
V = len(vocab)

# Tạo cặp skip-gram với kích thước cửa sổ 2
window = 2
pairs = []
for i, word in enumerate(corpus):
    for j in range(max(0, i - window), min(len(corpus), i + window + 1)):
        if i != j:
            pairs.append((word2idx[word], word2idx[corpus[j]]))

pairs = jnp.array(pairs)
print(f"Từ vựng: {V} từ, Cặp huấn luyện: {len(pairs)}")

# Tham số mô hình
embed_dim = 16
key = jax.random.PRNGKey(42)
k1, k2 = jax.random.split(key)
W_in = jax.random.normal(k1, (V, embed_dim)) * 0.1    # embedding đầu vào
W_out = jax.random.normal(k2, (V, embed_dim)) * 0.1   # embedding đầu ra

# Hàm mất mát negative sampling cho một cặp
def neg_sampling_loss(W_in, W_out, target, context, neg_ids):
    v_in = W_in[target]      # (embed_dim,)
    v_out = W_out[context]   # (embed_dim,)
    v_neg = W_out[neg_ids]   # (k, embed_dim)

    pos_loss = -jax.nn.log_sigmoid(jnp.dot(v_in, v_out))
    neg_loss = -jnp.sum(jax.nn.log_sigmoid(-v_neg @ v_in))
    return pos_loss + neg_loss

# Vòng lặp huấn luyện
num_neg = 5
lr = 0.05

@jax.jit
def train_step(W_in, W_out, target, context, neg_ids):
    loss, (g_in, g_out) = jax.value_and_grad(neg_sampling_loss, argnums=(0, 1))(
        W_in, W_out, target, context, neg_ids)
    return loss, W_in - lr * g_in, W_out - lr * g_out

key = jax.random.PRNGKey(0)
for epoch in range(50):
    total_loss = 0.0
    for i in range(len(pairs)):
        key, subkey = jax.random.split(key)
        neg_ids = jax.random.randint(subkey, (num_neg,), 0, V)
        loss, W_in, W_out = train_step(W_in, W_out, pairs[i, 0], pairs[i, 1], neg_ids)
        total_loss += loss
    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}: mất mát tb = {total_loss / len(pairs):.4f}")

# Trực quan hoá với PCA (chương 01)
embeddings = W_in
mean = embeddings.mean(axis=0)
centered = embeddings - mean
U, S, Vt = jnp.linalg.svd(centered, full_matrices=False)
coords = centered @ Vt[:2].T  # chiếu lên 2 PC đầu

plt.figure(figsize=(10, 8))
for i, word in idx2word.items():
    plt.scatter(coords[i, 0], coords[i, 1], c='#3498db', s=40)
    plt.annotate(word, (coords[i, 0] + 0.02, coords[i, 1] + 0.02), fontsize=9)
plt.title("Word2Vec Skip-gram Embeddings (PCA chiếu)")
plt.grid(alpha=0.3); plt.show()

Xây dựng mô hình ngôn ngữ RNN cấp ký tự học cách sinh văn bản từ một chuỗi huấn luyện nhỏ.

import jax
import jax.numpy as jnp

# Văn bản huấn luyện nhỏ
text = "to be or not to be that is the question "
chars = sorted(set(text))
char2idx = {c: i for i, c in enumerate(chars)}
idx2char = {i: c for c, i in char2idx.items()}
V = len(chars)
data = jnp.array([char2idx[c] for c in text])

# Tham số RNN
hidden_dim = 64
key = jax.random.PRNGKey(0)
k1, k2, k3, k4, k5 = jax.random.split(key, 5)

params = {
    'Wx': jax.random.normal(k1, (V, hidden_dim)) * 0.1,
    'Wh': jax.random.normal(k2, (hidden_dim, hidden_dim)) * 0.05,
    'bh': jnp.zeros(hidden_dim),
    'Wy': jax.random.normal(k3, (hidden_dim, V)) * 0.1,
    'by': jnp.zeros(V),
}

def rnn_step(params, h, x_idx):
    x = jnp.eye(V)[x_idx]  # one-hot
    h = jnp.tanh(x @ params['Wx'] + h @ params['Wh'] + params['bh'])
    logits = h @ params['Wy'] + params['by']
    return h, logits

def loss_fn(params, inputs, targets):
    h = jnp.zeros(hidden_dim)
    total_loss = 0.0
    for t in range(len(inputs)):
        h, logits = rnn_step(params, h, inputs[t])
        log_probs = jax.nn.log_softmax(logits)
        total_loss -= log_probs[targets[t]]
    return total_loss / len(inputs)

grad_fn = jax.jit(jax.grad(loss_fn))

# Huấn luyện
inputs = data[:-1]
targets = data[1:]
lr = 0.01

for step in range(500):
    grads = grad_fn(params, inputs, targets)
    params = {k: params[k] - lr * grads[k] for k in params}
    if (step + 1) % 100 == 0:
        l = loss_fn(params, inputs, targets)
        print(f"Bước {step+1}: mất mát = {l:.4f}")

# Sinh văn bản
def generate(params, seed_char, length=60):
    h = jnp.zeros(hidden_dim)
    idx = char2idx[seed_char]
    result = [seed_char]
    key = jax.random.PRNGKey(42)
    for _ in range(length):
        h, logits = rnn_step(params, h, idx)
        key, subkey = jax.random.split(key)
        idx = jax.random.categorical(subkey, logits)
        result.append(idx2char[int(idx)])
    return ''.join(result)

print(f"\nSinh: {generate(params, 't')}")

Triển khai mô hình seq2seq đồ chơi với Bahdanau attention cho tác vụ đảo ngược chuỗi. Trực quan hoá ma trận căn chỉnh attention.

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

# Tác vụ: đảo ngược một chuỗi chữ số (ví dụ: [3, 1, 4] -> [4, 1, 3])
vocab_size = 10  # chữ số 0-9
SOS, EOS = 10, 11  # token đặc biệt
total_vocab = 12
embed_dim, hidden_dim = 16, 32
max_len = 5

key = jax.random.PRNGKey(42)
keys = jax.random.split(key, 8)

params = {
    'embed': jax.random.normal(keys[0], (total_vocab, embed_dim)) * 0.1,
    'enc_Wx': jax.random.normal(keys[1], (embed_dim, hidden_dim)) * 0.1,
    'enc_Wh': jax.random.normal(keys[2], (hidden_dim, hidden_dim)) * 0.05,
    'dec_Wx': jax.random.normal(keys[3], (embed_dim, hidden_dim)) * 0.1,
    'dec_Wh': jax.random.normal(keys[4], (hidden_dim, hidden_dim)) * 0.05,
    # Bahdanau attention
    'Ws': jax.random.normal(keys[5], (hidden_dim, hidden_dim)) * 0.1,
    'Wh_att': jax.random.normal(keys[6], (hidden_dim, hidden_dim)) * 0.1,
    'v_att': jax.random.normal(keys[7], (hidden_dim,)) * 0.1,
    # Chiếu đầu ra (từ hidden + context sang vocab)
    'Wo': jax.random.normal(keys[0], (hidden_dim * 2, total_vocab)) * 0.1,
}

def encode(params, seq):
    """Mã hoá chuỗi đầu vào, trả về tất cả trạng thái ẩn."""
    h = jnp.zeros(hidden_dim)
    states = []
    for t in range(len(seq)):
        x = params['embed'][seq[t]]
        h = jnp.tanh(x @ params['enc_Wx'] + h @ params['enc_Wh'])
        states.append(h)
    return jnp.stack(states), h

def bahdanau_attention(params, dec_state, enc_states):
    """Tính trọng số Bahdanau attention và vector ngữ cảnh."""
    scores = jnp.tanh(enc_states @ params['Wh_att'] + dec_state @ params['Ws'])
    e = scores @ params['v_att']  # (src_len,)
    alpha = jax.nn.softmax(e)
    context = alpha @ enc_states
    return context, alpha

def decode_step(params, dec_h, prev_token, enc_states):
    x = params['embed'][prev_token]
    dec_h = jnp.tanh(x @ params['dec_Wx'] + dec_h @ params['dec_Wh'])
    context, alpha = bahdanau_attention(params, dec_h, enc_states)
    combined = jnp.concatenate([dec_h, context])
    logits = combined @ params['Wo']
    return dec_h, logits, alpha

def seq2seq_loss(params, src, tgt):
    enc_states, enc_final = encode(params, src)
    dec_h = enc_final
    loss = 0.0
    prev_token = SOS
    for t in range(len(tgt)):
        dec_h, logits, _ = decode_step(params, dec_h, prev_token, enc_states)
        log_probs = jax.nn.log_softmax(logits)
        loss -= log_probs[tgt[t]]
        prev_token = tgt[t]
    return loss / len(tgt)

# Tạo dữ liệu huấn luyện: đảo ngược chuỗi
key = jax.random.PRNGKey(0)
train_srcs, train_tgts = [], []
for _ in range(200):
    key, subkey = jax.random.split(key)
    length = jax.random.randint(subkey, (), 3, max_len + 1)
    key, subkey = jax.random.split(key)
    seq = jax.random.randint(subkey, (int(length),), 0, vocab_size)
    train_srcs.append(seq)
    train_tgts.append(seq[::-1])  # đảo ngược

# Huấn luyện
grad_fn = jax.grad(seq2seq_loss)
lr = 0.01

for epoch in range(100):
    total_loss = 0.0
    for src, tgt in zip(train_srcs, train_tgts):
        grads = grad_fn(params, src, tgt)
        params = {k: params[k] - lr * grads[k] for k in params}
        total_loss += seq2seq_loss(params, src, tgt)
    if (epoch + 1) % 20 == 0:
        print(f"Epoch {epoch+1}: mất mát tb = {total_loss / len(train_srcs):.4f}")

# Trực quan hoá attention cho một ví dụ
test_src = jnp.array([3, 1, 4, 1, 5])
test_tgt = test_src[::-1]

enc_states, enc_final = encode(params, test_src)
dec_h = enc_final
attentions = []
prev_token = SOS
for t in range(len(test_tgt)):
    dec_h, logits, alpha = decode_step(params, dec_h, prev_token, enc_states)
    attentions.append(alpha)
    prev_token = test_tgt[t]

att_matrix = jnp.stack(attentions)
fig, ax = plt.subplots(figsize=(6, 5))
im = ax.imshow(att_matrix, cmap='Blues')
ax.set_xlabel("Vị trí nguồn"); ax.set_ylabel("Vị trí đích")
src_labels = [str(int(x)) for x in test_src]
tgt_labels = [str(int(x)) for x in test_tgt]
ax.set_xticks(range(len(src_labels))); ax.set_xticklabels(src_labels)
ax.set_yticks(range(len(tgt_labels))); ax.set_yticklabels(tgt_labels)
for i in range(len(tgt_labels)):
    for j in range(len(src_labels)):
        ax.text(j, i, f"{att_matrix[i,j]:.2f}", ha='center', va='center', fontsize=9)
ax.set_title("Bahdanau Attention Alignment (đảo ngược chuỗi)")
plt.colorbar(im); plt.tight_layout(); plt.show()