Học Sâu (Deep Learning)

Học sâu xếp chồng các lớp phi tuyến để xây dựng các biểu diễn phân cấp tự động biến đổi đầu vào thô thành các đặc trưng hữu ích. File này bao gồm MLP, các hàm kích hoạt, lan truyền ngược, CNN, RNN, LSTM, cơ chế chú ý, transformer, GAN, VAE, mô hình khuếch tán và các kỹ thuật chuẩn hóa.

• Điều gì làm cho một mạng "sâu"? Một mạng nông có một lớp ẩn; một mạng sâu có nhiều lớp. Độ sâu cho phép mạng xây dựng các biểu diễn phân cấp, với các lớp đầu học các đặc trưng đơn giản (cạnh, tông màu) và các lớp sau kết hợp chúng thành các khái niệm phức tạp (khuôn mặt, câu văn). Tính tổ hợp (compositionality) này là điều mang lại sức mạnh cho học sâu.

• Mạng sâu đơn giản nhất là mạng perceptron nhiều lớp (multi-layer perceptron - MLP), còn gọi là mạng kết nối đầy đủ (fully connected) hay dày đặc (dense). Mỗi lớp tính toán:

\[h = \sigma(Wx + b)\]

• Ở đây \(W\) là một ma trận trọng số (chương 02), \(b\) là một vector độ chệch, và \(\sigma\) là một hàm kích hoạt phi tuyến. Đầu ra của một lớp trở thành đầu vào của lớp tiếp theo. Nếu không có phi tuyến, việc xếp chồng các lớp sẽ vô nghĩa: \(W_2(W_1 x) = (W_2 W_1)x\), chỉ là một phép biến đổi tuyến tính khác. Đây chính xác là sự sụp đổ phép nhân ma trận từ chương 02.

• Các hàm kích hoạt (Activation functions) đưa vào tính phi tuyến làm cho độ sâu trở nên có ý nghĩa.

• ReLU (Rectified Linear Unit): \(\text{ReLU}(x) = \max(0, x)\). Đây là hàm kích hoạt được dùng rộng rãi nhất. Nó tính nhanh, không bão hòa đối với đầu vào dương, và tạo ra các kích hoạt thưa (nhiều nơ-ron xuất ra chính xác 0). Mặt trái: các nơ-ron có đầu vào âm luôn xuất ra 0, và nếu chúng bị kẹt ở đó vĩnh viễn, chúng "chết" và ngừng học.

• Sigmoid: \(\sigma(x) = \frac{1}{1+e^{-x}}\), ép đầu vào vào \((0, 1)\). Hữu ích cho các lớp đầu ra trong phân lớp nhị phân, nhưng có vấn đề ở các lớp ẩn vì gradient biến mất khi đầu vào xa 0 (đường cong gần như phẳng).

• Tanh: \(\tanh(x) = \frac{e^x - e^{-x}}{e^x + e^{-x}}\), ép vào \((-1, 1)\). Có tâm tại 0 (khác với sigmoid), giúp dòng gradient chảy tốt hơn, nhưng vẫn chịu hiện tượng gradient biến mất ở các cực trị.

• GELU (Gaussian Error Linear Unit): \(\text{GELU}(x) = x \cdot \Phi(x)\), trong đó \(\Phi\) là hàm phân bố tích lũy (CDF) chuẩn. Nó là một xấp xỉ mượt của ReLU cho phép các giá trị âm nhỏ đi qua. GELU là mặc định trong GPT và BERT.

• Swish: \(\text{Swish}(x) = x \cdot \sigma(x)\), một cổng mượt khác. Trong thực tế tương tự GELU.

• Một lớp dày đặc với \(d_{\text{in}}\) đầu vào và \(d_{\text{out}}\) đầu ra có \(d_{\text{in}} \times d_{\text{out}} + d_{\text{out}}\) tham số (trọng số cộng với độ chệnh). Phép nhân ma trận \(Wx\) chỉ là phép nhân ma trận-vector từ chương 02. Trong thiết lập batch, đầu vào là một ma trận \(X\) có kích thước \((B, d_{\text{in}})\) và đầu ra là \(XW^T + b\) có kích thước \((B, d_{\text{out}})\).

• Định lý xấp xỉ phổ quát (universal approximation theorem) phát biểu rằng một lớp ẩn đơn với đủ nơ-ron có thể xấp xỉ bất kỳ hàm liên tục nào trên một miền compact với độ chính xác tùy ý. Điều này nghe có vẻ độ sâu không quan trọng, nhưng cái bẫy là "đủ nơ-ron." Trong thực tế, các mạng sâu có thể biểu diễn cùng các hàm với số tham số ít hơn theo cấp số mũ so với các mạng nông. Độ sâu cho bạn hiệu quả, không chỉ là khả năng biểu diễn.

• Khi các mạng sâu hơn, hai bệnh lý gradient xuất hiện. Gradient biến mất (Vanishing gradients): khi gradient đi qua nhiều lớp (qua quy tắc dẫn hàm hàm hợp, chương 03), chúng bị nhân với nhiều nhân tử. Nếu các nhân tử này nhất quán nhỏ hơn 1 (như xảy ra với sigmoid và tanh bão hòa), gradient thu nhỏ theo cấp số mũ về 0. Các lớp đầu hầu như không học được gì. Gradient bùng nổ (Exploding gradients): nếu các nhân tử nhất quán lớn hơn 1, gradient tăng theo cấp số mũ, gây tràn số (overflow) và huấn luyện không ổn định.

• Các giải pháp cho gradient biến mất/bùng nổ:

• Dùng các hàm kích hoạt ReLU hoặc GELU (gradient bằng 1 với đầu vào dương, không bão hòa)
• Khởi tạo trọng số cẩn thận
• Các lớp chuẩn hóa (normalisation layers)
• Các kết nối thặng dư (skip connections)

• Cắt gradient (gradient clipping) (cho gradient bùng nổ): giới hạn chuẩn gradient tại một giá trị cực đại

• Khởi tạo trọng số (Weight initialisation) quan trọng vì nó xác định quy mô của các kích hoạt và gradient lúc bắt đầu huấn luyện. Nếu trọng số quá lớn, các kích hoạt bùng nổ; quá nhỏ, chúng biến mất.

• Khởi tạo Xavier (Glorot) đặt trọng số từ một phân bố với phương sai \(\frac{2}{d_{\text{in}} + d_{\text{out}}}\). Cách này giữ phương sai của các kích hoạt xấp xỉ không đổi qua các lớp, giả định các hàm kích hoạt tuyến tính hoặc tanh.

• Khởi tạo He (Kaiming) dùng phương sai \(\frac{2}{d_{\text{in}}}\), được hiệu chỉnh cho các hàm kích hoạt ReLU (vì ReLU triệt tiêu một nửa các kích hoạt, bạn cần gấp đôi phương sai để bù đắp).

• Các lớp chuẩn hóa (Normalisation layers) ổn định huấn luyện bằng cách đảm bảo các đầu vào của mỗi lớp có thống kê nhất quán (xấp xỉ trung bình 0, phương sai 1).

• Chuẩn hóa theo batch (Batch Normalisation - BatchNorm) chuẩn hóa theo chiều batch: với mỗi kênh/đặc trưng, tính trung bình và phương sai trên tất cả các mẫu trong mini-batch, sau đó chuẩn hóa. Nó thêm các tham số tỷ lệ \(\gamma\) và dịch chuyển \(\beta\) có thể học được để mạng có thể đảo ngược việc chuẩn hóa nếu cần:

\[\hat{x} = \frac{x - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}, \quad y = \gamma \hat{x} + \beta\]

• BatchNorm có một vấn đề: nó phụ thuộc vào kích thước batch. Với các batch rất nhỏ, thống kê bị nhiễu. Tại thời điểm suy luận, bạn dùng các trung bình chạy (running averages) thay vì thống kê batch, tạo ra sự khác biệt train/test.

• Chuẩn hóa theo lớp (Layer Normalisation - LayerNorm) chuẩn hóa theo chiều đặc trưng cho mỗi mẫu riêng lẻ. Nó không phụ thuộc vào các mẫu khác trong batch, làm cho nó là lựa chọn tiêu chuẩn cho transformer và các mạng hồi quy.

• Chuẩn hóa theo thực thể (Instance Normalisation) chuẩn hóa theo các chiều không gian cho mỗi mẫu và mỗi kênh độc lập. Nó phổ biến trong chuyển phong cách (style transfer).

• Chuẩn hóa theo nhóm (Group Normalisation) chia các kênh thành các nhóm và chuẩn hóa trong mỗi nhóm. Nó là sự thỏa hiệp giữa LayerNorm và InstanceNorm.

• Dropout là một kỹ thuật chính quy hóa ngẫu nhiên triệt tiêu một tỷ lệ \(p\) các nơ-ron trong quá trình huấn luyện. Cách này buộc mạng không dựa vào bất kỳ nơ-ron đơn lẻ nào, khuyến khích các biểu diễn dự phòng. Tại thời điểm kiểm tra, tất cả nơ-ron đều hoạt động. Dropout đảo ngược (Inverted dropout) tỷ lệ các kích hoạt với \(\frac{1}{1-p}\) trong huấn luyện để không cần tỷ lệ hóa tại thời điểm kiểm tra. Đây là triển khai tiêu chuẩn.

• Mạng nơ-ron tích chập (Convolutional Neural Networks - CNNs) khai thác cấu trúc không gian. Thay vì kết nối mọi đầu vào với mọi đầu ra (như trong các lớp dày đặc), một lớp tích chập trượt một bộ lọc nhỏ (kernel) qua đầu vào, tính tích vô hướng tại mỗi vị trí. Cùng một trọng số bộ lọc được chia sẻ trên tất cả các vị trí, giúp giảm mạnh tham số và xây dựng tính bất biến dịch chuyển (translation invariance).

• Phép tích chập (convolution operation) cho một đầu vào 2D với bộ lọc \(K\) kích thước \(k \times k\):

\[(\text{input} * K)[i,j] = \sum_{m=0}^{k-1} \sum_{n=0}^{k-1} \text{input}[i+m, j+n] \cdot K[m, n]\]

• Kích thước đầu ra phụ thuộc vào ba siêu tham số. Stride kiểm soát số pixel bộ lọc di chuyển giữa các vị trí (stride 2 giảm một nửa các chiều không gian). Padding thêm số 0 quanh biên đầu vào (padding "same" giữ nguyên kích thước không gian, "valid" thì không). Công thức kích thước đầu ra: \(\text{out} = \lfloor (\text{in} - k + 2p) / s \rfloor + 1\).

• Các lớp pooling giảm kích thước bản đồ đặc trưng. Max pooling lấy giá trị lớn nhất trong mỗi cửa sổ; average pooling lấy trung bình. Pooling giảm các chiều không gian trong khi giữ lại thông tin quan trọng nhất.

• Tích chập giãn nở (Dilated convolutions) chèn các khoảng trống giữa các phần tử bộ lọc, tăng trường tiếp nhận (receptive field) mà không tăng tham số. Tỷ lệ giãn nở (dilation rate) 2 nghĩa là bộ lọc 3x3 bao phủ một vùng 5x5.

• Tích chập 1x1 là các tích chập với bộ lọc 1x1. Chúng không nhìn vào các hàng xóm không gian; thay vào đó, chúng trộn thông tin qua các kênh. Hãy nghĩ về chúng như áp dụng một lớp dày đặc tại mỗi vị trí không gian. Chúng được dùng để thay đổi số kênh với chi phí rẻ.

• Kết nối thặng dư (Skip connections) (kết nối thặng dư) cho phép đầu vào bỏ qua một hoặc nhiều lớp: \(\text{output} = F(x) + x\). Lớp chỉ cần học phần dư \(F(x) = \text{output} - x\), điều dễ dàng hơn khi phép biến đổi tối ưu gần với hàm đồng nhất (identity). ResNets (Residual Networks) xếp chồng hơn 100 lớp dùng thủ thuật này, giải quyết vấn đề suy thoái (degradation) nơi các mạng sâu hơn hoạt động tệ hơn các mạng nông hơn.

• CNN xây dựng một phân cấp đặc trưng (feature hierarchy). Các lớp đầu phát hiện cạnh và họa tiết. Các lớp giữa kết hợp chúng thành các phần (mắt, bánh xe). Các lớp cuối nhận biết toàn bộ vật thể. Trường tiếp nhận của mỗi lớp (vùng đầu vào mà nó có thể "thấy") tăng theo độ sâu.

• Embedding ánh xạ các token rời rạc (từ, ký tự, ID mục) thành các vector dày đặc. Một lớp embedding chỉ là một bảng tra cứu: một ma trận \(E\) có kích thước (kích thước từ vựng, số chiều embedding). Tra cứu token \(i\) nghĩa là chọn hàng \(i\) của \(E\). Điều này tương đương với nhân với một vector one-hot, chỉ là một trường hợp đặc biệt của phép nhân ma trận-vector (chương 02). Embedding được học trong huấn luyện, do đó các token tương tự sẽ có các vector tương tự.

• Tokenisation là quá trình chuyển văn bản thô thành một chuỗi các token. Tokenisation cấp từ (word-level) tách theo khoảng trắng nhưng không xử lý được các từ chưa thấy. Tokenisation cận từ (Subword tokenisation) (BPE, WordPiece, SentencePiece) bẻ văn bản thành các đơn vị cận từ thường xuyên, cân bằng kích thước từ vựng và độ phủ. Từ "unhappiness" có thể trở thành ["un", "happiness"] hoặc ["un", "happ", "iness"].

• Mạng nơ-ron hồi quy (Recurrent Neural Networks - RNNs) xử lý chuỗi từng phần tử một, duy trì một trạng thái ẩn (hidden state) mang thông tin về phía trước:

\[h_t = \tanh(W_h h_{t-1} + W_x x_t + b)\]

• Trạng thái ẩn \(h_t\) là một bản tóm tắt nén của mọi thứ mạng đã thấy cho đến thời điểm \(t\). Cùng các trọng số \(W_h\) và \(W_x\) được chia sẻ qua tất cả các bước thời gian (chia sẻ trọng số, như CNN chia sẻ trọng số không gian).

• RNN nguyên bản gặp khó khăn với các chuỗi dài vì gradient biến mất: tín hiệu gradient từ bước \(t\) đến bước \(t - k\) đi qua \(k\) phép nhân với \(W_h\), và nó thu nhỏ (hoặc bùng nổ) theo cấp số mũ.

• LSTM (Long Short-Term Memory) giải quyết bằng cách đưa vào một trạng thái ô (cell state) \(c_t\) riêng biệt chảy qua thời gian với sự can thiệp tối thiểu. Ba cổng kiểm soát thông tin nào đi vào, ra và tồn tại:

• Cổng quên (forget gate) quyết định xóa gì khỏi trạng thái ô: \(f_t = \sigma(W_f [h_{t-1}, x_t] + b_f)\)
• Cổng vào (input gate) quyết định thông tin mới nào được ghi: \(i_t = \sigma(W_i [h_{t-1}, x_t] + b_i)\), với các giá trị ứng cử viên \(\tilde{c}_t = \tanh(W_c [h_{t-1}, x_t] + b_c)\)
• Trạng thái ô cập nhật: \(c_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t\)
• Cổng ra (output gate) quyết định cái gì được phơi bày: \(o_t = \sigma(W_o [h_{t-1}, x_t] + b_o)\), và \(h_t = o_t \odot \tanh(c_t)\)

• Trạng thái ô hoạt động như một băng truyền: thông tin có thể chảy không đổi qua nhiều bước thời gian (cổng quên giữ gần 1), giải quyết vấn đề gradient biến mất cho các phụ thuộc tầm xa.

• GRU (Gated Recurrent Unit) đơn giản hóa LSTM bằng cách gộp trạng thái ô và trạng thái ẩn thành một, và dùng hai cổng thay vì ba: một cổng cập nhật (kết hợp quên và vào) và một cổng đặt lại (reset). GRU có ít tham số hơn và thường hoạt động tương đương LSTM.

• Giới hạn cơ bản của RNN (kể cả LSTM) là xử lý tuần tự: bạn phải xử lý token 1 trước token 2 trước token 3. Điều này ngăn song song hóa và tạo ra một nút thắt thông tin, vì mọi ngữ cảnh phải chèn qua trạng thái ẩn có kích thước cố định.

• Cơ chế chú ý (Attention) giải quyết cả hai vấn đề. Thay vì nén toàn bộ đầu vào vào một vector cố định, chú ý cho phép mô hình nhìn lại tất cả các vị trí đầu vào và quyết định vị trí nào liên quan cho đầu ra hiện tại.

• Công thức hiện đại dùng truy vấn, khóa và giá trị (queries, keys, and values - Q, K, V). Hãy nghĩ về nó như một tìm kiếm thư viện: bạn có một truy vấn (cái bạn đang tìm), các khóa (nhãn trên mỗi cuốn sách), và các giá trị (nội dung sách thực tế). Bạn so sánh truy vấn của mình với tất cả các khóa để tìm ra giá trị nào cần truy xuất.

• Chú ý tích vô hướng có tỷ lệ (Scaled dot-product attention):

\[\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) V\]

• \(QK^T\) tính độ tương đồng giữa mọi truy vấn và mọi khóa. Đây là một phép nhân ma trận (chương 02), và các phần tử là các tích vô hướng, đo độ tương đồng cosin (chương 01). Chia cho \(\sqrt{d_k}\) ngăn các tích vô hướng trở nên quá lớn (điều sẽ làm softmax bão hòa và tạo ra các phân bố gần one-hot với gradient biến mất). Softmax chuyển các độ tương đồng thành một phân bố xác suất. Nhân với \(V\) tạo ra một tổ hợp có trọng số của các giá trị.

• Chú ý đa đầu (Multi-head attention) chạy \(h\) phép chú ý song song, mỗi cái với các phép chiếu Q, K, V đã học khác nhau. Cách này cho phép mô hình chú ý thông tin từ các không gian con biểu diễn khác nhau đồng thời. Một đầu có thể chú ý đến các quan hệ cú pháp trong khi đầu khác chú ý đến ngữ nghĩa. Các đầu ra được nối (concatenate) và chiếu:

\[\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O\]

• Kiến trúc Transformer (Vaswani et al., 2017) được xây dựng hoàn toàn từ chú ý và các lớp truyền xuôi (feed-forward), không có sự hồi quy. Khối encoder lặp lại: chú ý đa đầu tự chú ý, cộng và chuẩn hóa lớp, mạng truyền xuôi, cộng và chuẩn hóa lớp. Khối decoder thêm một tự chú ý có mặt nạ (ngăn mô hình thấy các token tương lai) và một lớp chú ý chéo (cross-attention) chú ý đến đầu ra của encoder.

• Mã hóa vị trí (Positional encoding) là cần thiết vì chú ý là tương đẳng hoán vị (permutation-equivariant), nghĩa là nó coi đầu vào như một tập hợp, không phải một chuỗi. Nếu không có thông tin vị trí, "the cat sat on the mat" và "the mat sat on the cat" sẽ giống hệt nhau. Transformer nguyên bản dùng các mã hóa vị trí sin:

\[PE_{(pos, 2i)} = \sin\!\left(\frac{pos}{10000^{2i/d}}\right), \quad PE_{(pos, 2i+1)} = \cos\!\left(\frac{pos}{10000^{2i/d}}\right)\]

• Mỗi vị trí nhận một vector duy nhất mà mô hình có thể dùng để phân biệt các vị trí. Các mô hình hiện đại thường dùng các embedding vị trí đã học hoặc mã hóa vị trí tương đối (RoPE, ALiBi) thay thế.

• Transformer xử lý tất cả các token song song (ma trận tự chú ý \(QK^T\) được tính trong một phép nhân ma trận), làm cho chúng huấn luyện nhanh hơn nhiều so với RNN trên phần cứng hiện đại. Sự đánh đổi là tự chú ý có độ phức tạp \(O(n^2)\) theo độ dài chuỗi (mọi token chú ý đến mọi token khác), trong khi RNN là \(O(n)\). Đây là lý do các mô hình ngữ cảnh dài cần các biến thể chú ý đặc biệt (chú ý thưa, chú ý tuyến tính, flash attention).

• Vision Transformers (ViT) áp dụng Transformer cho ảnh bằng cách chia ảnh thành các mảnh (patch) có kích thước cố định (ví dụ 16x16), làm phẳng mỗi mảnh thành một vector, và coi các mảnh như một chuỗi các token. Một token [CLS] có thể học được được thêm vào đầu, và biểu diễn cuối của nó được dùng cho phân lớp. Dù không có các thiên kiến quy nạp tích chập (convolutional inductive biases), ViT sánh ngang hoặc vượt CNN khi được huấn luyện với đủ dữ liệu.

• MLP-Mixer là một kiến trúc thậm chí còn đơn giản hơn, thay thế cả chú ý và tích chập bằng MLP. Nó luân phiên giữa các MLP "trộn token" (áp dụng qua các vị trí không gian) và các MLP "trộn kênh" (áp dụng qua các đặc trưng). Nó hoạt động cạnh tranh, gợi ý rằng ý tưởng then chốt của các kiến trúc hiện đại không phải là chú ý bản thân, mà là sự trộn thông tin hiệu quả qua các token và đặc trưng.

• Bộ tự mã hóa (Autoencoders) học các biểu diễn nén bằng cách huấn luyện một mạng tái tạo lại chính đầu vào của nó. Bộ mã hóa (encoder) ánh xạ đầu vào vào một nút thắt (bottleneck) có số chiều thấp hơn (mã tiềm ẩn - latent code), và bộ giải mã (decoder) ánh xạ nó trở lại:

\[z = f_{\text{enc}}(x), \quad \hat{x} = f_{\text{dec}}(z), \quad \mathcal{L} = \|x - \hat{x}\|^2\]

• Nút thắt buộc mạng phải học các đặc trưng quan trọng nhất. Bộ tự mã hóa được dùng cho giảm chiều, khử nhiễu (huấn luyện trên đầu vào nhiễu, tái tạo đầu ra sạch) và phát hiện bất thường (lỗi tái tạo cao báo hiệu một đầu vào bất thường).

• Bộ tự mã hóa biến phân (Variational Autoencoders - VAE) thêm một bước ngoặt xác suất. Thay vì mã hóa thành một điểm đơn \(z\), bộ mã hóa xuất ra các tham số của một phân bố (trung bình \(\mu\) và phương sai \(\sigma^2\) của một Gaussian). Mã tiềm ẩn được lấy mẫu từ phân bố này: \(z = \mu + \sigma \odot \epsilon\), với \(\epsilon \sim \mathcal{N}(0, I)\). Thủ thuật tái tham số hóa (reparameterisation trick) này làm cho việc lấy mẫu có thể vi phân (differentiable) để gradient có thể chảy qua.

• Mất mát VAE có hai số hạng:

\[\mathcal{L} = \underbrace{\|x - \hat{x}\|^2}_{\text{reconstruction}} + \underbrace{D_{\text{KL}}(q(z|x) \| p(z))}_{\text{regularisation}}\]

• Số hạng phân kỳ KL (từ chương 05) đẩy hậu nghiệm đã học \(q(z|x)\) về tiên nghiệm \(p(z) = \mathcal{N}(0, I)\), đảm bảo không gian tiềm ẩn mượt và có cấu trúc tốt. Bạn sau đó có thể lấy mẫu từ tiên nghiệm và giải mã để sinh dữ liệu mới. Đây là điều làm cho VAE là các mô hình sinh (generative models).

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

1. Xây dựng một MLP đơn giản từ đầu trong JAX. Huấn luyện nó trên một bài toán phân lớp 2D (ví dụ các đường tròn đồng tâm) và trực quan hóa ranh giới quyết định.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   from sklearn.datasets import make_circles
   
   # Data
   X, y = make_circles(n_samples=500, noise=0.1, factor=0.5, random_state=42)
   X, y = jnp.array(X), jnp.array(y, dtype=jnp.float32)
   
   # Initialise a 2-layer MLP: 2 -> 16 -> 16 -> 1
   def init_params(key):
       k1, k2, k3 = jax.random.split(key, 3)
       return {
           'W1': jax.random.normal(k1, (2, 16)) * 0.5,
           'b1': jnp.zeros(16),
           'W2': jax.random.normal(k2, (16, 16)) * 0.5,
           'b2': jnp.zeros(16),
           'W3': jax.random.normal(k3, (16, 1)) * 0.5,
           'b3': jnp.zeros(1),
       }
   
   def forward(params, x):
       h = jnp.maximum(0, x @ params['W1'] + params['b1'])  # ReLU
       h = jnp.maximum(0, h @ params['W2'] + params['b2'])   # ReLU
       logit = (h @ params['W3'] + params['b3']).squeeze()
       return jax.nn.sigmoid(logit)
   
   def loss_fn(params, X, y):
       pred = forward(params, X)
       return -jnp.mean(y * jnp.log(pred + 1e-7) + (1 - y) * jnp.log(1 - pred + 1e-7))
   
   grad_fn = jax.jit(jax.grad(loss_fn))
   params = init_params(jax.random.PRNGKey(0))
   lr = 0.1
   
   for step in range(2000):
       grads = grad_fn(params, X, y)
       params = {k: params[k] - lr * grads[k] for k in params}
   
   # Plot decision boundary
   xx, yy = jnp.meshgrid(jnp.linspace(-2, 2, 200), jnp.linspace(-2, 2, 200))
   grid = jnp.column_stack([xx.ravel(), yy.ravel()])
   zz = forward(params, grid).reshape(xx.shape)
   
   plt.figure(figsize=(7, 6))
   plt.contourf(xx, yy, zz, levels=[0, 0.5, 1], alpha=0.3, colors=['#e74c3c', '#3498db'])
   plt.scatter(X[y==0,0], X[y==0,1], c='#e74c3c', s=10, label='Class 0')
   plt.scatter(X[y==1,0], X[y==1,1], c='#3498db', s=10, label='Class 1')
   plt.title("MLP Decision Boundary on Concentric Circles")
   plt.legend(); plt.grid(alpha=0.3); plt.show()
   
   acc = jnp.mean((forward(params, X) > 0.5) == y)
   print(f"Accuracy: {acc:.2%}")
   

2. Cài đặt tích chập 1D từ đầu. Áp dụng một bộ lọc phát hiện cạnh đơn giản cho một tín hiệu và so sánh với jnp.convolve có sẵn.

   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   def conv1d(signal, kernel):
       """1D convolution (valid mode) from scratch."""
       n, k = len(signal), len(kernel)
       output = jnp.zeros(n - k + 1)
       for i in range(n - k + 1):
           output = output.at[i].set(jnp.sum(signal[i:i+k] * kernel))
       return output
   
   # Create a signal with a step function
   t = jnp.linspace(0, 4, 200)
   signal = jnp.where(t < 1, 0.0, jnp.where(t < 2, 1.0, jnp.where(t < 3, 0.5, 1.5)))
   
   # Edge detection kernel
   edge_kernel = jnp.array([-1.0, 0.0, 1.0])
   
   # Our implementation vs built-in
   our_output = conv1d(signal, edge_kernel)
   jnp_output = jnp.convolve(signal, edge_kernel, mode='valid')
   
   fig, axes = plt.subplots(3, 1, figsize=(10, 6), sharex=True)
   axes[0].plot(t, signal, color='#3498db', linewidth=1.5)
   axes[0].set_title("Original Signal"); axes[0].set_ylabel("Value")
   
   axes[1].plot(t[:len(our_output)], our_output, color='#e74c3c', linewidth=1.5)
   axes[1].set_title("After Edge Detection (our conv1d)"); axes[1].set_ylabel("Value")
   
   axes[2].plot(t[:len(jnp_output)], jnp_output, color='#27ae60', linewidth=1.5, linestyle='--')
   axes[2].set_title("After Edge Detection (jnp.convolve)"); axes[2].set_ylabel("Value")
   axes[2].set_xlabel("t")
   
   plt.tight_layout(); plt.show()
   print(f"Outputs match: {jnp.allclose(our_output, jnp_output)}")
   

3. Cài đặt chú ý tích vô hướng có tỷ lệ từ đầu. Tính trọng số chú ý cho một ví dụ nhỏ và trực quan hóa ma trận chú ý như một bản đồ nhiệt.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   def scaled_dot_product_attention(Q, K, V):
       """Scaled dot-product attention."""
       d_k = Q.shape[-1]
       scores = Q @ K.T / jnp.sqrt(d_k)
       weights = jax.nn.softmax(scores, axis=-1)
       output = weights @ V
       return output, weights
   
   # Example: 4 tokens, embedding dim 8
   key = jax.random.PRNGKey(42)
   k1, k2, k3 = jax.random.split(key, 3)
   seq_len, d_model = 4, 8
   
   Q = jax.random.normal(k1, (seq_len, d_model))
   K = jax.random.normal(k2, (seq_len, d_model))
   V = jax.random.normal(k3, (seq_len, d_model))
   
   output, weights = scaled_dot_product_attention(Q, K, V)
   
   print(f"Q shape: {Q.shape}")
   print(f"Attention weights shape: {weights.shape}")
   print(f"Output shape: {output.shape}")
   print(f"\nAttention weights (rows sum to 1):")
   print(weights)
   print(f"Row sums: {weights.sum(axis=-1)}")
   
   # Visualise attention
   fig, ax = plt.subplots(figsize=(5, 4))
   im = ax.imshow(weights, cmap='Blues', vmin=0, vmax=1)
   ax.set_xlabel("Key position"); ax.set_ylabel("Query position")
   ax.set_title("Attention Weights")
   tokens = ['tok 0', 'tok 1', 'tok 2', 'tok 3']
   ax.set_xticks(range(4)); ax.set_xticklabels(tokens)
   ax.set_yticks(range(4)); ax.set_yticklabels(tokens)
   for i in range(4):
       for j in range(4):
           ax.text(j, i, f"{weights[i,j]:.2f}", ha='center', va='center', fontsize=10)
   plt.colorbar(im); plt.tight_layout(); plt.show()
   

4. Xây dựng một bộ tự mã hóa đơn giản nén dữ liệu 2D qua một nút thắt 1D và tái tạo nó. Trực quan hóa không gian tiềm ẩn và các phép tái tạo.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   from sklearn.datasets import make_moons
   
   # Data
   X, _ = make_moons(n_samples=500, noise=0.05, random_state=42)
   X = jnp.array(X)
   
   # Autoencoder: 2 -> 8 -> 1 -> 8 -> 2
   def init_ae(key):
       k1, k2, k3, k4 = jax.random.split(key, 4)
       return {
           'enc_W1': jax.random.normal(k1, (2, 8)) * 0.5, 'enc_b1': jnp.zeros(8),
           'enc_W2': jax.random.normal(k2, (8, 1)) * 0.5, 'enc_b2': jnp.zeros(1),
           'dec_W1': jax.random.normal(k3, (1, 8)) * 0.5, 'dec_b1': jnp.zeros(8),
           'dec_W2': jax.random.normal(k4, (8, 2)) * 0.5, 'dec_b2': jnp.zeros(2),
       }
   
   def encode(p, x):
       h = jnp.tanh(x @ p['enc_W1'] + p['enc_b1'])
       return h @ p['enc_W2'] + p['enc_b2']
   
   def decode(p, z):
       h = jnp.tanh(z @ p['dec_W1'] + p['dec_b1'])
       return h @ p['dec_W2'] + p['dec_b2']
   
   def ae_loss(p, X):
       z = encode(p, X)
       X_hat = decode(p, z)
       return jnp.mean((X - X_hat) ** 2)
   
   grad_fn = jax.jit(jax.grad(ae_loss))
   params = init_ae(jax.random.PRNGKey(0))
   lr = 0.01
   
   for step in range(3000):
       grads = grad_fn(params, X)
       params = {k: params[k] - lr * grads[k] for k in params}
   
   z = encode(params, X)
   X_hat = decode(params, z)
   
   fig, axes = plt.subplots(1, 2, figsize=(12, 5))
   axes[0].scatter(X[:,0], X[:,1], c=z.squeeze(), cmap='viridis', s=10)
   axes[0].set_title("Original Data (coloured by latent code)")
   axes[1].scatter(X_hat[:,0], X_hat[:,1], c=z.squeeze(), cmap='viridis', s=10)
   axes[1].set_title("Reconstruction from 1D bottleneck")
   for ax in axes:
       ax.set_aspect('equal'); ax.grid(alpha=0.3)
   plt.tight_layout(); plt.show()
   
   print(f"Reconstruction MSE: {ae_loss(params, X):.4f}")