Transformer Thị giác và Sinh Ảnh

Transformer thị giác áp dụng cơ chế chú ý (self-attention) lên các mảnh ảnh, thách thức sự thống trị của CNN với khả năng học không gian dựa trên dữ liệu. File này đề cập đến ViT, DeiT, Swin Transformer, sinh ảnh với GAN (StyleGAN), VAE, và các mô hình khuếch tán (DDPM, Stable Diffusion), cùng với siêu phân giải và chuyển phong cách neural.

• CNN (file 02) xây dựng các inductive bias không gian mạnh: kết nối cục bộ, chia sẻ trọng số, và bất biến tịnh tiến. Vision Transformer (ViT) đặt một câu hỏi táo bạo: điều gì sẽ xảy ra nếu chúng ta loại bỏ hoàn toàn các bias này và để mô hình tự học cấu trúc không gian từ dữ liệu, chỉ sử dụng cơ chế chú ý từ chương 06?

• Vision Transformer (ViT) (Dosovitskiy và cộng sự, 2021) áp dụng bộ mã hoá Transformer tiêu chuẩn trực tiếp lên ảnh. Ý tưởng chính là coi một ảnh như một chuỗi các mảnh (patch), giống như NLP coi văn bản như một chuỗi các token.

• Quá trình này hoạt động như sau:

• Chia ảnh (chiều cao \(H\), chiều rộng \(W\), số kênh \(C\)) thành một lưới các mảnh không chồng lấn kích thước \(P \times P\). Điều này tạo ra \(N = HW / P^2\) mảnh.
• Làm phẳng mỗi mảnh thành một vector độ dài \(P^2 \cdot C\) và chiếu nó lên chiều mô hình \(D\) qua một phép embedding tuyến tính có thể học (một phép nhân ma trận đơn, chương 02).
• Thêm một embedding token [CLS] có thể học (tương tự [CLS] của BERT, chương 07). Token này chú ý đến tất cả các mảnh và biểu diễn cuối cùng của nó được dùng cho phân loại.
• Thêm position embeddings (một vector có thể học cho mỗi vị trí) để cung cấp thông tin không gian, vì cơ chế chú ý là bất biến với hoán vị.
• Đưa chuỗi \((N + 1)\) token embedding qua một bộ mã hoá Transformer tiêu chuẩn (multi-head self-attention + FFN, chương 06).
• Biểu diễn cuối cùng của token [CLS] được đưa qua đầu phân loại (một MLP nhỏ).

• Patch embedding tương đương với một tích chập với kích thước kernel \(P\) và stride \(P\) (không chồng lấn). ViT theo nghĩa đen chuyển đổi ảnh 2D thành chuỗi 1D, sau đó xử lý nó với cùng kiến trúc được dùng cho ngôn ngữ.

• ViT có ít inductive bias hơn CNN: nó không áp đặt kết nối cục bộ hay bất biến tịnh tiến. Điều này có nghĩa là nó cần nhiều dữ liệu huấn luyện hơn để học cấu trúc không gian từ đầu. Trên tập dữ liệu nhỏ, CNN vượt trội hơn ViT. Nhưng khi được huấn luyện trên các tập dữ liệu rất lớn (JFT-300M, 300 triệu ảnh), ViT đạt hoặc vượt các CNN tốt nhất, cho thấy rằng inductive bias của CNN hữu ích cho hiệu quả dữ liệu nhưng không cần thiết cho hiệu suất cuối cùng.

• Self-attention của ViT có độ phức tạp \(O(N^2)\) theo số lượng mảnh. Đối với ảnh 224x224 với mảnh 16x16, \(N = 196\), ở mức quản lý được. Nhưng đối với ảnh độ phân giải cao hơn hoặc mảnh nhỏ hơn, chi phí bậc hai trở nên cấm kỵ.

• DeiT (Data-efficient Image Transformer, Touvron và cộng sự, 2021) cho thấy ViT có thể được huấn luyện hiệu quả chỉ trên ImageNet (không cần tập JFT khổng lồ) bằng cách sử dụng tăng cường dữ liệu mạnh, chính quy hoá (stochastic depth, label smoothing, dropout), và chưng cất kiến thức (knowledge distillation): một CNN giáo viên đã tiền huấn luyện cung cấp các nhãn mềm mà ViT học sinh học để khớp. DeiT thêm một distillation token bên cạnh token [CLS], được huấn luyện để dự đoán đầu ra của giáo viên.

• Swin Transformer (Liu và cộng sự, 2021) giải quyết hai hạn chế chính của ViT: chi phí bậc hai với kích thước ảnh và thiếu feature map phân cấp (cần thiết cho phát hiện và phân đoạn).

• Swin giới thiệu cửa sổ dịch chuyển (shifted windows): thay vì self-attention toàn cục trên tất cả các mảnh, attention được tính toán trong các cửa sổ cục bộ (ví dụ: mảnh 7x7). Điều này làm chi phí tuyến tính với kích thước ảnh: \(O(N)\) thay vì \(O(N^2)\). Nhưng chỉ dùng cửa sổ cục bộ sẽ ngăn luồng thông tin giữa các vùng.

• Dịch chuyển cửa sổ (window shifting) giải quyết vấn đề này: trong các tầng xen kẽ, phân vùng cửa sổ được dịch chuyển một nửa kích thước cửa sổ. Điều này tạo ra các kết nối chéo giữa các cửa sổ, cho phép thông tin chảy giữa tất cả các phần của ảnh qua các tầng mà không tốn chi phí attention toàn cục.

• Swin cũng xây dựng một biểu diễn phân cấp (hierarchical representation) bằng cách gộp các mảnh qua các giai đoạn. Sau mỗi giai đoạn, các mảnh 2x2 lân cận được nối và chiếu để tăng gấp đôi chiều kênh và giảm một nửa độ phân giải không gian. Điều này tạo ra các feature map đa tỷ lệ tương tự như trong CNN và FPN (file 03), giúp Swin tương thích trực tiếp với các đầu phát hiện như Faster R-CNN và đầu phân đoạn như U-Net.

• PVT (Pyramid Vision Transformer) thực hiện cách tiếp cận phân cấp tương tự với spatial-reduction attention: tại mỗi giai đoạn, các keys và values được giảm mẫu về mặt không gian trước khi tính attention, giảm chi phí bậc hai trong khi duy trì trường tiếp nhận toàn cục.

• Học tự giám sát thị giác (self-supervised visual learning) huấn luyện các biểu diễn từ ảnh không gán nhãn. Nhãn rất đắt để thu thập, nhưng ảnh thì có sẵn. Mục tiêu là học các đặc trưng có thể chuyển giao tốt sang các tác vụ hạ nguồn mà không cần bất kỳ chú thích nào của con người.

• Học đối lập (contrastive learning) huấn luyện mô hình nhận ra rằng hai góc nhìn tăng cường khác nhau của cùng một ảnh (một "cặp dương tính") nên có biểu diễn tương tự nhau, trong khi các góc nhìn của các ảnh khác nhau ("cặp âm tính") nên có biểu diễn khác biệt.

• SimCLR (Chen và cộng sự, 2020) tạo hai góc nhìn tăng cường của mỗi ảnh trong một batch, mã hoá cả hai với một backbone + projection head dùng chung, và áp dụng mất mát NT-Xent (normalised temperature-scaled cross-entropy):

\[ \ell_{i,j} = -\log \frac{\exp(\text{sim}(z_i, z_j) / \tau)}{\sum_{k \neq i} \exp(\text{sim}(z_i, z_k) / \tau)} \]

• trong đó \(\text{sim}\) là độ tương tự cosine (chương 01) và \(\tau\) là tham số nhiệt độ. Tử số đẩy các cặp dương tính lại gần nhau; mẫu số đẩy các cặp âm tính xa nhau. SimCLR yêu cầu kích thước batch lớn (4,096+) để cung cấp đủ mẫu âm tính.

• MoCo (Momentum Contrast, He và cộng sự, 2020) giải quyết yêu cầu batch lớn bằng cách duy trì một hàng đợi cập nhật động lượng (momentum-updated queue) các embedding âm tính. Bộ mã hoá truy vấn được cập nhật bằng hạ gradient; bộ mã hoá khoá được cập nhật như một trung bình động luỹ thừa (EMA, chương 04) của bộ mã hoá truy vấn: \(\theta_k \leftarrow m \theta_k + (1 - m) \theta_q\), với \(m = 0.999\). Hàng đợi lưu trữ các key embedding gần đây, cung cấp một tập lớn và nhất quán các mẫu âm tính mà không cần batch khổng lồ.

• BYOL (Bootstrap Your Own Latent, Grill và cộng sự, 2020) loại bỏ hoàn toàn các cặp âm tính. Nó sử dụng hai mạng: mạng "online" và mạng "target" (EMA của online). Mạng online dự đoán biểu diễn của mạng target từ một góc nhìn tăng cường khác. Không có âm tính, BYOL tránh vấn đề suy sụp (collapse — khi mô hình xuất ra cùng vector cho mọi thứ) thông qua sự bất đối xứng của đầu dự đoán và target EMA.

• DINO (Self-Distillation with No Labels, Caron và cộng sự, 2021) áp dụng tự chưng cất (self-distillation) cho ViT. Mạng học sinh dự đoán đầu ra của mạng giáo viên (EMA của học sinh) qua các góc nhìn tăng cường khác nhau. Giáo viên sử dụng các crop lớn hơn; học sinh sử dụng các crop nhỏ hơn. DINO tạo ra các đặc trưng chứa thông tin tường minh về bố cục cảnh: các bản đồ self-attention của ViT được huấn luyện bằng DINO tự nhiên phân đoạn các vật thể mà không cần bất kỳ giám sát phân đoạn nào.

• Mô hình hoá ảnh bị che (masked image modelling) là tương tự thị giác của mô hình hoá ngôn ngữ bị che của BERT (chương 07). Một tỷ lệ lớn các mảnh đầu vào bị che, và mô hình học để tái tạo chúng.

• MAE (Masked Autoencoders, He và cộng sự, 2022) che 75% số mảnh và huấn luyện một bộ mã hoá-giải mã ViT để tái tạo các giá trị pixel bị thiếu. Chỉ các mảnh không bị che được xử lý bởi bộ mã hoá (tiết kiệm 4x tính toán trong quá trình tiền huấn luyện), và bộ giải mã nhẹ tái tạo ảnh đầy đủ từ các mảnh đã mã hoá cộng với các mask token có thể học.

• BEiT (BERT Pre-training of Image Transformers, Bao và cộng sự, 2022) che các mảnh và dự đoán các token thị giác rời rạc (thu được từ một bộ token hoá dVAE đã tiền huấn luyện) thay vì pixel thô. Điều này tương tự với việc BERT dự đoán các token từ rời rạc và tránh các chi tiết cấp thấp của tái tạo pixel.

• Sinh ảnh (image generation) nhằm tạo ra các ảnh mới, chân thực không tồn tại trong tập huấn luyện. Thách thức cốt lõi là mô hình hoá phân bố xác suất nhiều chiều của ảnh tự nhiên.

• Mạng Đối Sinh (Generative Adversarial Networks - GAN) (Goodfellow và cộng sự, 2014) sử dụng hai mạng cạnh tranh: một bộ sinh (generator) \(G\) tạo ảnh giả từ nhiễu ngẫu nhiên, và một bộ phân biệt (discriminator) \(D\) cố gắng phân biệt ảnh thật và ảnh giả. Chúng được huấn luyện đối kháng: \(G\) cố gắng đánh lừa \(D\), và \(D\) cố gắng bắt quả tang \(G\).

\[ \min_G \max_D \; \mathbb{E}_{x \sim p_{\text{data}}}[\log D(x)] + \mathbb{E}_{z \sim p(z)}[\log(1 - D(G(z)))] \]

• Bộ sinh nhận một vector tiềm ẩn ngẫu nhiên \(z\) (được lấy mẫu từ một phân bố đơn giản như Gaussian) và ánh xạ nó qua một chuỗi các tích chập chuyển vị để tạo ra một ảnh. Bộ phân biệt là một bộ phân loại CNN tiêu chuẩn. Tại trạng thái cân bằng, \(G\) tạo ra các ảnh không thể phân biệt với dữ liệu thật, và \(D\) xuất ra 0.5 cho mọi đầu vào.

• Sụp đổ mode (mode collapse) là chế độ thất bại chính của GAN: bộ sinh học chỉ tạo ra một vài loại ảnh đánh lừa được bộ phân biệt, bỏ qua sự đa dạng của dữ liệu huấn luyện. Bộ sinh tìm một tập nhỏ các đầu ra "an toàn" thay vì bao phủ toàn bộ phân bố.

• Các thủ thuật huấn luyện giúp ổn định GAN bao gồm: chuẩn hoá phổ (spectral normalisation — hạn chế hằng số Lipschitz của discriminator), progressive growing (huấn luyện ở độ phân giải thấp trước, sau đó tăng dần), feature matching (khớp thống kê của các đặc trưng discriminator trung gian thay vì đầu ra cuối cùng), và sử dụng khoảng cách Wasserstein thay vì mục tiêu phân kỳ JS gốc.

• StyleGAN (Karras và cộng sự, 2019) là kiến trúc GAN có ảnh hưởng nhất cho tổng hợp ảnh chất lượng cao. Cải tiến chính của nó là bộ sinh dựa trên phong cách (style-based generator): thay vì đưa vector tiềm ẩn \(z\) trực tiếp vào bộ sinh, nó trước hết được ánh xạ qua một mạng ánh xạ (mapping network) (MLP 8 tầng) để tạo ra vector phong cách \(w\). Vector phong cách này được tiêm vào mỗi tầng của bộ sinh qua adaptive instance normalisation (AdaIN), điều chỉnh các thống kê của feature map:

\[ \text{AdaIN}(x, y) = y_{s} \cdot \frac{x - \mu(x)}{\sigma(x)} + y_{b} \]

• trong đó \(y_s\) và \(y_b\) là tỷ lệ và độ dịch được suy ra từ \(w\). Các tầng khác nhau kiểm soát các khía cạnh khác nhau: tầng đầu kiểm soát các đặc trưng thô (tư thế, hình dạng khuôn mặt), tầng giữa kiểm soát các đặc trưng trung bình (kiểu tóc, mắt), và tầng cuối kiểm soát các chi tiết tinh tế (tàn nhang, kết cấu tóc). StyleGAN có thể tạo ra khuôn mặt chân thực ở độ phân giải 1024x1024.

• Bộ mã hoá tự động biến phân (Variational Autoencoders - VAE) (chương 06) cung cấp một cách tiếp cận sinh thay thế. Không giống GAN, VAE có một khung xác suất có nguyên tắc với mục tiêu huấn luyện rõ ràng (ELBO). Chúng có xu hướng tạo ra ảnh mờ hơn GAN nhưng cung cấp không gian tiềm ẩn mượt mà hơn, có cấu trúc hơn. VAE là cặp mã hoá-giải mã được dùng trong các mô hình khuếch tán tiềm ẩn để nén ảnh vào và ra khỏi không gian tiềm ẩn.

• Mô hình khuếch tán (diffusion models) đã trở thành mô hình thống trị cho sinh ảnh, vượt qua GAN về cả chất lượng và độ đa dạng. Ý tưởng khái niệm đơn giản: dần dần thêm nhiễu vào dữ liệu cho đến khi nó trở thành nhiễu Gaussian thuần tuý (quá trình xuôi - forward process), sau đó học cách đảo ngược quá trình này từng bước (quá trình ngược - reverse process).

• Quá trình xuôi (forward process) thêm nhiễu Gaussian qua \(T\) bước thời gian:

\[ q(x_t | x_{t-1}) = \mathcal{N}(x_t; \sqrt{1 - \beta_t} \, x_{t-1}, \beta_t I) \]

• trong đó \(\beta_t\) là lịch nhiễu tăng dần theo thời gian. Sau đủ số bước, \(x_T\) xấp xỉ là nhiễu Gaussian thuần tuý bất kể ảnh gốc \(x_0\). Sử dụng thủ thuật tham số hoá lại (reparametrisation trick, chương 06) và đặt \(\alpha_t = 1 - \beta_t\), \(\bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s\), ta có thể lấy mẫu \(x_t\) trực tiếp từ \(x_0\):

\[ x_t = \sqrt{\bar{\alpha}_t} \, x_0 + \sqrt{1 - \bar{\alpha}_t} \, \epsilon, \quad \epsilon \sim \mathcal{N}(0, I) \]

• Quá trình ngược (reverse process) học cách khử nhiễu: bắt đầu từ nhiễu thuần tuý \(x_T\), mô hình dự đoán nhiễu \(\epsilon\) được thêm tại mỗi bước và trừ nó đi để khôi phục \(x_{t-1}\). Điều này được tham số hoá bởi một mạng nơ-ron \(\epsilon_\theta\) (thường là một U-Net, từ file 03), được huấn luyện với một hàm mất mát MSE đơn giản:

\[ \mathcal{L} = \mathbb{E}_{t, x_0, \epsilon}\left[\|\epsilon - \epsilon_\theta(x_t, t)\|^2\right] \]

• DDPM (Denoising Diffusion Probabilistic Models, Ho và cộng sự, 2020) thiết lập khung sườn này. Lấy mẫu yêu cầu lặp qua tất cả \(T\) bước (thường là 1,000), rất chậm. DDIM (Denoising Diffusion Implicit Models, Song và cộng sự, 2021) công thức hoá lại quá trình lấy mẫu như một ánh xạ xác định (deterministic), cho phép bỏ qua các bước lớn (ví dụ: 50 bước thay vì 1,000) với mất mát chất lượng tối thiểu.

• Mô hình dựa trên điểm số (score-based models) (Song và Ermon, 2019) cung cấp một góc nhìn thay thế. Thay vì dự đoán nhiễu \(\epsilon\), mô hình ước lượng hàm điểm số (score function) \(\nabla_{x_t} \log p(x_t)\), gradient của log-xác suất theo ảnh nhiễu. Gradient này chỉ về phía các vùng có xác suất cao hơn (sạch hơn) của phân bố dữ liệu. Việc lấy mẫu đi theo gradient này bằng động lực học Langevin. Mô hình dựa trên điểm số và DDPM đã được thống nhất trong khung của phương trình vi phân ngẫu nhiên (stochastic differential equations - SDE): quá trình xuôi là một SDE thêm nhiễu, và quá trình ngược là SDE đảo ngược thời gian.

• Hướng dẫn không dùng bộ phân biệt (Classifier-free guidance) (Ho và Salimans, 2022) kiểm soát sự đánh đổi giữa chất lượng và độ đa dạng mẫu. Mô hình được huấn luyện cả có điều kiện (với một prompt văn bản hoặc nhãn lớp) và không điều kiện (với điều kiện bị bỏ ngẫu nhiên). Tại thời điểm lấy mẫu, dự đoán là một tổ hợp có trọng số:

\[ \hat{\epsilon} = \epsilon_\theta(x_t, \varnothing) + s \cdot (\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \varnothing)) \]

• trong đó \(c\) là điều kiện, \(\varnothing\) là điều kiện rỗng, và \(s > 1\) là thang hướng dẫn (guidance scale). \(s\) cao hơn tạo ra ảnh khớp với điều kiện mạnh hơn nhưng kém đa dạng hơn. \(s = 1\) cho mô hình không được hướng dẫn; \(s = 7.5\) là một giá trị mặc định phổ biến.

• Khuếch tán tiềm ẩn (latent diffusion) (Rombach và cộng sự, 2022; Stable Diffusion) chuyển quá trình khuếch tán từ không gian pixel sang một không gian tiềm ẩn được học. Một bộ mã hoá VAE đã tiền huấn luyện nén ảnh thành một biểu diễn tiềm ẩn chiều thấp hơn (thường giảm mẫu không gian 4x hoặc 8x), khuếch tán hoạt động trong không gian nén này, và bộ giải mã VAE tái tạo pixel từ tiềm ẩn đã khử nhiễu. Điều này hiệu quả hơn nhiều: khuếch tán một ảnh 512x512 trong không gian pixel nghĩa là xử lý một tensor \(512 \times 512 \times 3\), nhưng trong không gian tiềm ẩn chỉ là một tensor \(64 \times 64 \times 4\).

• Mạng U-Net khử nhiễu trong khuếch tán tiềm ẩn nhận tiềm ẩn nhiễu, bước thời gian (được mã hoá như một embedding sin, tương tự như positional encoding trong transformer), và một tín hiệu điều kiện hoá (embedding văn bản từ một bộ mã hoá CLIP hoặc T5 bị đóng băng). Điều kiện văn bản đi vào qua các tầng cross-attention trong U-Net: các embedding văn bản đóng vai trò là keys và values, và các đặc trưng ảnh đóng vai trò là queries. Điều này cho phép mô hình chú ý đến các phần liên quan của prompt văn bản tại mỗi vị trí không gian.

• Flow matching là một lựa chọn thay thế đang nổi lên so với khuếch tán, học một đường vận chuyển trực tiếp giữa nhiễu và dữ liệu, thay vì khử nhiễu lặp đi lặp lại của DDPM.

• Một dòng chuẩn hoá liên tục (continuous normalising flow - CNF) xác định một trường vận tốc phụ thuộc thời gian \(v_\theta(x, t)\) đẩy các mẫu từ một phân bố đơn giản \(p_0\) (nhiễu) đến phân bố dữ liệu \(p_1\) dọc theo các quỹ đạo mượt. Phép biến đổi tuân theo một phương trình vi phân thường (ODE):

\[ \frac{dx}{dt} = v_\theta(x, t), \quad t \in [0, 1] \]

• Bắt đầu từ \(x_0 \sim \mathcal{N}(0, I)\), tích phân ODE về phía trước đến \(t = 1\) tạo ra một mẫu từ phân bố dữ liệu. Trường vận tốc được tham số hoá bởi một mạng nơ-ron và được huấn luyện để khớp với một dòng có điều kiện mục tiêu.

• Flow matching tối ưu vận chuyển (optimal transport - OT) (Lipman và cộng sự, 2023) sử dụng các đường thẳng giữa nhiễu và dữ liệu làm dòng mục tiêu: đường có điều kiện từ mẫu nhiễu \(x_0\) đến mẫu dữ liệu \(x_1\) đơn giản là \(x_t = (1 - t) x_0 + t x_1\), và vận tốc mục tiêu là \(v = x_1 - x_0\). Hàm mất mát huấn luyện trở thành:

\[ \mathcal{L} = \mathbb{E}_{t, x_0, x_1} \left[\|v_\theta(x_t, t) - (x_1 - x_0)\|^2\right] \]

• Dòng được hiệu chỉnh (Rectified flows) (Liu và cộng sự, 2022) lặp đi lặp lại làm thẳng các đường dòng đã học. Sau một lượt huấn luyện ban đầu, mô hình được dùng để tạo ra các cặp (nhiễu, dữ liệu) bằng cách mô phỏng ODE. Các cặp này, được căn chỉnh chặt chẽ hơn so với các ghép ngẫu nhiên, được dùng để huấn luyện lại mô hình. Lặp lại quá trình này tạo ra các đường ngày càng thẳng, có thể đi qua với ít bước ODE hơn (thậm chí một bước duy nhất), cho phép sinh cực nhanh.

• Flow matching có một số ưu điểm so với khuếch tán: mục tiêu huấn luyện đơn giản hơn (hồi quy vận tốc trực tiếp, không cần lịch nhiễu), ODE lấy mẫu mượt hơn (yêu cầu ít bước tích phân hơn), và sự kết nối với tối ưu vận chuyển cung cấp nền tảng lý thuyết. Stable Diffusion 3 và Flux sử dụng flow matching thay vì DDPM truyền thống.

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

1. Cài đặt patch embedding của ViT từ đầu. Chia ảnh thành các mảnh, làm phẳng chúng, chiếu lên chiều mô hình, thêm position embeddings, và thêm token [CLS] vào đầu.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   def create_patch_embedding(image, patch_size, d_model, params):
       """Convert an image into a sequence of patch embeddings."""
       H, W, C = image.shape
       n_patches_h = H // patch_size
       n_patches_w = W // patch_size
       n_patches = n_patches_h * n_patches_w
   
       # Extract patches
       patches = []
       for i in range(n_patches_h):
           for j in range(n_patches_w):
               patch = image[i*patch_size:(i+1)*patch_size,
                             j*patch_size:(j+1)*patch_size, :]
               patches.append(patch.ravel())
       patches = jnp.stack(patches)  # (N, P*P*C)
   
       # Linear projection to d_model
       embeddings = patches @ params['proj_w'] + params['proj_b']  # (N, d_model)
   
       # Prepend CLS token
       cls_token = params['cls_token']  # (1, d_model)
       embeddings = jnp.concatenate([cls_token, embeddings], axis=0)  # (N+1, d_model)
   
       # Add position embeddings
       embeddings = embeddings + params['pos_embed']  # (N+1, d_model)
   
       return embeddings, patches
   
   # Setup
   H, W, C = 32, 32, 3
   patch_size = 8
   d_model = 64
   n_patches = (H // patch_size) * (W // patch_size)  # 16
   
   key = jax.random.PRNGKey(42)
   keys = jax.random.split(key, 5)
   
   # Create a synthetic image with distinct quadrants
   image = jnp.zeros((H, W, C))
   image = image.at[:16, :16, 0].set(1.0)   # red top-left
   image = image.at[:16, 16:, 1].set(1.0)   # green top-right
   image = image.at[16:, :16, 2].set(1.0)   # blue bottom-left
   image = image.at[16:, 16:, :2].set(1.0)  # yellow bottom-right
   
   params = {
       'proj_w': jax.random.normal(keys[0], (patch_size**2 * C, d_model)) * 0.02,
       'proj_b': jnp.zeros(d_model),
       'cls_token': jax.random.normal(keys[1], (1, d_model)) * 0.02,
       'pos_embed': jax.random.normal(keys[2], (n_patches + 1, d_model)) * 0.02,
   }
   
   embeddings, patches = create_patch_embedding(image, patch_size, d_model, params)
   
   print(f"Image shape: {image.shape}")
   print(f"Patch size: {patch_size}x{patch_size}")
   print(f"Number of patches: {n_patches}")
   print(f"Patch vector length: {patch_size**2 * C}")
   print(f"Embedding shape: {embeddings.shape}  (CLS + {n_patches} patches)")
   
   # Visualise patches
   fig, axes = plt.subplots(2, 5, figsize=(14, 6))
   axes[0, 0].imshow(image); axes[0, 0].set_title('Full Image'); axes[0, 0].axis('off')
   for idx in range(min(9, n_patches)):
       ax = axes[(idx+1) // 5, (idx+1) % 5]
       patch_img = patches[idx].reshape(patch_size, patch_size, C)
       ax.imshow(patch_img); ax.set_title(f'Patch {idx}'); ax.axis('off')
   plt.suptitle('ViT Patch Decomposition')
   plt.tight_layout(); plt.show()
   

2. Cài đặt một vòng lặp huấn luyện GAN đơn giản. Huấn luyện một bộ sinh và bộ phân biệt trên dữ liệu 2D và trực quan hoá phân bố được sinh hội tụ về phân bố thật.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   def generator(z, params):
       h = jnp.tanh(z @ params['g_w1'] + params['g_b1'])
       h = jnp.tanh(h @ params['g_w2'] + params['g_b2'])
       return h @ params['g_w3'] + params['g_b3']
   
   def discriminator(x, params):
       h = jax.nn.leaky_relu(x @ params['d_w1'] + params['d_b1'], 0.2)
       h = jax.nn.leaky_relu(h @ params['d_w2'] + params['d_b2'], 0.2)
       return jax.nn.sigmoid(h @ params['d_w3'] + params['d_b3'])
   
   def init_params(key):
       keys = jax.random.split(key, 6)
       z_dim, h_dim, data_dim = 2, 32, 2
       scale = 0.1
       return {
           'g_w1': jax.random.normal(keys[0], (z_dim, h_dim)) * scale,
           'g_b1': jnp.zeros(h_dim),
           'g_w2': jax.random.normal(keys[1], (h_dim, h_dim)) * scale,
           'g_b2': jnp.zeros(h_dim),
           'g_w3': jax.random.normal(keys[2], (h_dim, data_dim)) * scale,
           'g_b3': jnp.zeros(data_dim),
           'd_w1': jax.random.normal(keys[3], (data_dim, h_dim)) * scale,
           'd_b1': jnp.zeros(h_dim),
           'd_w2': jax.random.normal(keys[4], (h_dim, h_dim)) * scale,
           'd_b2': jnp.zeros(h_dim),
           'd_w3': jax.random.normal(keys[5], (h_dim, 1)) * scale,
           'd_b3': jnp.zeros(1),
       }
   
   def d_loss(params, real_data, fake_data):
       real_score = discriminator(real_data, params)
       fake_score = discriminator(fake_data, params)
       return -jnp.mean(jnp.log(real_score + 1e-7) + jnp.log(1 - fake_score + 1e-7))
   
   def g_loss(params, fake_data):
       fake_score = discriminator(fake_data, params)
       return -jnp.mean(jnp.log(fake_score + 1e-7))
   
   # Real data: ring distribution
   key = jax.random.PRNGKey(42)
   theta = jax.random.uniform(key, (512,)) * 2 * jnp.pi
   real_data = jnp.stack([jnp.cos(theta), jnp.sin(theta)], axis=1)
   real_data = real_data + jax.random.normal(key, real_data.shape) * 0.05
   
   params = init_params(jax.random.PRNGKey(0))
   d_grad = jax.grad(d_loss)
   g_grad = jax.grad(g_loss)
   lr = 0.001
   
   snapshots = []
   for step in range(3000):
       key, k1 = jax.random.split(key)
       z = jax.random.normal(k1, (512, 2))
       fake_data = generator(z, params)
   
       # Update discriminator
       grads = d_grad(params, real_data, fake_data)
       for k in ['d_w1', 'd_b1', 'd_w2', 'd_b2', 'd_w3', 'd_b3']:
           params[k] = params[k] - lr * grads[k]
   
       # Update generator
       fake_data = generator(z, params)
       grads = g_grad(params, fake_data)
       for k in ['g_w1', 'g_b1', 'g_w2', 'g_b2', 'g_w3', 'g_b3']:
           params[k] = params[k] - lr * grads[k]
   
       if step in [0, 500, 1500, 2999]:
           snapshots.append((step, fake_data.copy()))
   
   fig, axes = plt.subplots(1, 4, figsize=(16, 4))
   for ax, (step, fake) in zip(axes, snapshots):
       ax.scatter(real_data[:, 0], real_data[:, 1], s=5, alpha=0.3, c='#3498db', label='Real')
       ax.scatter(fake[:, 0], fake[:, 1], s=5, alpha=0.3, c='#e74c3c', label='Generated')
       ax.set_title(f'Step {step}'); ax.set_xlim(-2, 2); ax.set_ylim(-2, 2)
       ax.set_aspect('equal'); ax.legend(markerscale=3)
   plt.suptitle('GAN Training: Generator Learns the Ring Distribution')
   plt.tight_layout(); plt.show()
   

3. Cài đặt quá trình khuếch tán xuôi: thêm nhiễu vào một ảnh tại các bước thời gian tăng dần và trực quan hoá sự phá huỷ dần dần. Sau đó cài đặt một bước khử nhiễu đơn.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   def noise_schedule(T, beta_start=0.0001, beta_end=0.02):
       """Linear noise schedule."""
       betas = jnp.linspace(beta_start, beta_end, T)
       alphas = 1.0 - betas
       alpha_bars = jnp.cumprod(alphas)
       return betas, alphas, alpha_bars
   
   def forward_diffusion(x0, t, alpha_bars, key):
       """Add noise to x0 at timestep t."""
       alpha_bar_t = alpha_bars[t]
       noise = jax.random.normal(key, x0.shape)
       xt = jnp.sqrt(alpha_bar_t) * x0 + jnp.sqrt(1 - alpha_bar_t) * noise
       return xt, noise
   
   # Create a simple 2D "image" (checkerboard)
   img = jnp.zeros((32, 32))
   for i in range(4):
       for j in range(4):
           if (i + j) % 2 == 0:
               img = img.at[i*8:(i+1)*8, j*8:(j+1)*8].set(1.0)
   
   T = 1000
   betas, alphas, alpha_bars = noise_schedule(T)
   
   # Visualise forward process
   timesteps = [0, 50, 200, 500, 999]
   key = jax.random.PRNGKey(42)
   
   fig, axes = plt.subplots(1, len(timesteps), figsize=(16, 3.5))
   for ax, t in zip(axes, timesteps):
       key, subkey = jax.random.split(key)
       xt, noise = forward_diffusion(img, t, alpha_bars, subkey)
       ax.imshow(xt, cmap='gray', vmin=-2, vmax=2)
       ax.set_title(f't={t}\n$\\bar{{\\alpha}}$={alpha_bars[t]:.3f}')
       ax.axis('off')
   plt.suptitle('Diffusion Forward Process: Progressive Noise Addition')
   plt.tight_layout(); plt.show()
   
   # Simple denoising: train a tiny network to predict noise at t=200
   t_denoise = 200
   key, k1 = jax.random.split(key)
   xt, true_noise = forward_diffusion(img, t_denoise, alpha_bars, k1)
   
   # Tiny "denoiser": just learn a constant noise estimate (for illustration)
   noise_estimate = jnp.zeros_like(img)
   lr = 0.01
   for step in range(100):
       residual = noise_estimate - true_noise
       noise_estimate = noise_estimate - lr * residual
   
   # Reverse one step
   alpha_bar_t = alpha_bars[t_denoise]
   x_denoised = (xt - jnp.sqrt(1 - alpha_bar_t) * noise_estimate) / jnp.sqrt(alpha_bar_t)
   
   fig, axes = plt.subplots(1, 3, figsize=(12, 4))
   axes[0].imshow(img, cmap='gray'); axes[0].set_title('Original $x_0$'); axes[0].axis('off')
   axes[1].imshow(xt, cmap='gray', vmin=-2, vmax=2)
   axes[1].set_title(f'Noisy $x_{{200}}$'); axes[1].axis('off')
   axes[2].imshow(x_denoised, cmap='gray')
   axes[2].set_title('Denoised (one step)'); axes[2].axis('off')
   plt.tight_layout(); plt.show()
   
   mse = jnp.mean((x_denoised - img)**2)
   print(f"Denoising MSE: {mse:.4f}")