Token hóa Ảnh và Video

Token hóa ảnh và video chuyển đổi dữ liệu thị giác liên tục thành các chuỗi token rời rạc mà transformer có thể xử lý giống như văn bản. File này bàn về VQ-VAE, VQ-GAN, học codebook, dVAE của DALL-E, token hóa video và lượng tử không cần tra cứu (lookup-free quantisation).

Tại sao cần token hóa ảnh

• Hãy nghĩ về ngôn ngữ như một bảng chữ cái hữu hạn: tiếng Anh có khoảng 26 chữ cái, và các mô hình ngôn ngữ hiện đại chia văn bản thành 30.000-100.000 token. Mỗi câu trở thành một chuỗi các ký hiệu rời rạc mà transformer có thể dự đoán từng cái một. Mặt khác, ảnh tồn tại trong một không gian liên tục, có số chiều cao: một ảnh RGB 256x256 đơn lẻ là một điểm trong \(\mathbb{R}^{256 \times 256 \times 3} \approx \mathbb{R}^{196{,}608}\). Nếu bạn muốn một mô hình ngôn ngữ "nói" về ảnh với cùng cơ chế nó dùng để nói tiếng Anh, bạn cần chuyển đổi các mảng pixel liên tục đó thành một chuỗi token rời rạc có thể quản lý được, được rút ra từ một từ vựng hữu hạn. Sự chuyển đổi đó là token hóa ảnh (image tokenisation).

• Hãy tưởng tượng bạn là một nghệ nhân khảm. Bạn không có vô hạn các sắc thái gạch; bạn có một bảng màu cố định, chẳng hạn gồm 8192 màu gạch riêng biệt. Để tái tạo một bức ảnh dưới dạng tranh khảm, bạn phải (1) quyết định vùng nào của ảnh mỗi viên gạch đại diện, (2) chọn màu gạch gần nhất cho mỗi vùng, và (3) chấp nhận rằng một số chi tiết bị mất nhưng bức tranh tổng thể vẫn có thể nhận ra. Token hóa ảnh làm chính điều này: một bộ mã hóa nén các patch không gian thành các vector tiềm ẩn (latent), một codebook ánh xạ mỗi vector đến mục nhập gần nhất của nó, và kết quả là một lưới các chỉ số nguyên, một chỉ số cho mỗi patch, mà một mô hình rời rạc có thể xử lý.

• Lợi ích của token hóa là ba mặt. Thứ nhất, nó nén ảnh một cách đáng kể: một ảnh 256x256 có thể trở thành lưới token 16x16, giảm độ dài chuỗi từ 65.536 pixel xuống 256 token, khả thi cho các mô hình dựa trên attention vốn có chi phí tăng theo bậc hai với độ dài chuỗi. Thứ hai, nó thống nhất biểu diễn: token văn bản và token ảnh tồn tại trong cùng một từ vựng rời rạc, cho phép một transformer tự hồi quy duy nhất sinh ra văn bản và ảnh xen kẽ. Thứ ba, nó áp đặt một nút thắt cổ chai hữu ích buộc mô hình học các mã có ý nghĩa ngữ nghĩa thay vì ghi nhớ nhiễu pixel.

• Nhớ lại từ Chương 8 cách mạng nơ-ron tích chập trích xuất các bản đồ đặc trưng phân cấp từ ảnh, và từ Chương 7 cách bộ token hóa văn bản chuyển đổi chuỗi ký tự thành các chuỗi số nguyên. Token hóa ảnh nằm ở giao điểm: nó dùng một CNN hoặc vision transformer encoder (Chương 8) để tạo ra các đặc trưng không gian, sau đó mượn ý tưởng về một từ vựng rời rạc (Chương 7) để chuyển đổi các đặc trưng đó thành các chỉ số token.

VQ-VAE: Lượng tử hóa vector (Vector Quantisation)

• Như ta đã thấy ở Chương 6, một autoencoder biến phân (variational autoencoder, VAE) chuẩn mã hóa đầu vào thành một phân bố tiềm ẩn liên tục và giải mã các mẫu từ phân bố đó trở lại thành tái tạo. Không gian tiềm ẩn là liên tục, điều này khiến nó khó để đưa vào các mô hình chuỗi rời rạc. Vector Quantised Variational Autoencoder (VQ-VAE), được giới thiệu bởi van den Oord et al. (2017), thay thế không gian tiềm ẩn liên tục bằng một không gian rời rạc bằng cách đưa vào một codebook học được gồm các vector embedding và gắn (snap) mỗi đầu ra của bộ mã hóa đến mục nhập codebook gần nhất.

• Hãy hình dung một thư viện có chính xác \(K\) kệ được dán nhãn. Khi một cuốn sách mới (đầu ra của bộ mã hóa) đến, thủ thư đặt nó lên kệ mà những cuốn sách hiện có (các vector codebook) trên đó trông giống nhất, và ghi lại số kệ. Sau đó, để lấy lại cuốn sách, bạn chỉ cần số kệ: mục nhập codebook trên kệ đó là một đại diện đủ tốt. Đây chính là lượng tử hóa vector.

• Về mặt hình thức, VQ-VAE có ba thành phần:

• Một bộ mã hóa (encoder) \(E\) ánh xạ ảnh đầu vào \(\mathbf{x} \in \mathbb{R}^{H \times W \times 3}\) thành một lưới không gian các vector tiềm ẩn liên tục \(\mathbf{z}_e = E(\mathbf{x}) \in \mathbb{R}^{h \times w \times d}\), trong đó \(h \times w\) là độ phân giải không gian được giảm mẫu và \(d\) là số chiều embedding.

• Một codebook \(\mathcal{C} = \{\mathbf{e}_1, \mathbf{e}_2, \ldots, \mathbf{e}_K\} \subset \mathbb{R}^d\) chứa \(K\) vector embedding học được. Kích thước codebook điển hình từ 512 đến 16.384 mục.

• Một bộ giải mã (decoder) \(D\) tái tạo ảnh từ các biểu diễn tiềm ẩn đã lượng tử hóa.

• Bước lượng tử hóa thay thế mỗi đầu ra của bộ mã hóa \(\mathbf{z}_e(\mathbf{x})\) tại vị trí không gian \((i, j)\) bằng mục nhập codebook gần nhất của nó:

\[\mathbf{z}_q(i,j) = \mathbf{e}_{k^\ast} \quad \text{với} \quad k^\ast = \arg\min_k \|\mathbf{z}_e(i,j) - \mathbf{e}_k\|_2\]

• Đây là phép tra cứu láng giềng gần nhất trong không gian embedding, chính xác là thao tác gán trong k-means (Chương 6). Chỉ số \(k^\ast\) là token rời rạc cho vị trí không gian \((i,j)\), và toàn bộ ảnh được biểu diễn dưới dạng một lưới \(h \times w\) các số nguyên từ \(\{1, \ldots, K\}\).

• Thách thức là \(\arg\min\) không khả vi: bạn không thể lan truyền ngược (backpropagate) thông qua một lựa chọn rời rạc. VQ-VAE giải quyết vấn đề này bằng bộ ước lượng đường thẳng (straight-through estimator): trong lượt truyền tiến, bộ giải mã nhận \(\mathbf{z}_q\) (vector đã lượng tử hóa); trong lượt truyền ngược, gradient của mất mát tái tạo đối với \(\mathbf{z}_q\) được sao chép trực tiếp đến \(\mathbf{z}_e\), như thể bước lượng tử hóa là hàm đồng nhất. Điều này được viết gọn là:

\[\mathbf{z}_q = \mathbf{z}_e + \text{sg}(\mathbf{z}_q - \mathbf{z}_e)\]

• trong đó \(\text{sg}(\cdot)\) là toán tử dừng-gradient (stop-gradient). Trong lượt truyền tiến, biểu thức này cho kết quả \(\mathbf{z}_q\); trong lượt truyền ngược, gradient chỉ chảy qua số hạng \(\mathbf{z}_e\).

• Mất mát VQ-VAE đầy đủ có ba số hạng:

\[\mathcal{L} = \underbrace{\|\mathbf{x} - D(\mathbf{z}_q)\|_2^2}_{\text{tái tạo}} + \underbrace{\|\text{sg}(\mathbf{z}_e) - \mathbf{e}\|_2^2}_{\text{codebook (VQ)}} + \underbrace{\beta \|\mathbf{z}_e - \text{sg}(\mathbf{e})\|_2^2}_{\text{cam kết (commitment)}}\]

• Mất mát tái tạo (reconstruction loss) huấn luyện bộ mã hóa và bộ giải mã tái tạo trung thực đầu vào. Mất mát codebook (còn gọi là mất mát VQ) kéo các vector codebook về phía đầu ra của bộ mã hóa; lưu ý rằng \(\text{sg}(\mathbf{z}_e)\) có nghĩa là bộ mã hóa không nhận gradient từ số hạng này, vì vậy nó chỉ cập nhật codebook. Mất mát cam kết (commitment loss) làm điều ngược lại: nó khuyến khích đầu ra bộ mã hóa ở gần các vector codebook, ngăn bộ mã hóa "chạy xa" khỏi codebook. Siêu tham số \(\beta\) (thường là 0,25) kiểm soát sự cân bằng giữa hai số hạng codebook và cam kết.

• Trong thực tế, codebook thường được cập nhật bằng trung bình động hàm số mũ (exponential moving average, EMA) thay vì hạ gradient, vốn ổn định hơn. Gọi \(\mathbf{n}_k\) là số lượng đầu ra bộ mã hóa được gán cho mục codebook \(k\) và \(\mathbf{s}_k\) là tổng của chúng. Cập nhật EMA là:

\[\mathbf{n}_k \leftarrow \gamma \mathbf{n}_k + (1 - \gamma) |\{(i,j) : k^\ast_{ij} = k\}|\]

\[\mathbf{s}_k \leftarrow \gamma \mathbf{s}_k + (1 - \gamma) \sum_{(i,j) : k^\ast_{ij} = k} \mathbf{z}_e(i,j)\]

\[\mathbf{e}_k \leftarrow \frac{\mathbf{s}_k}{\mathbf{n}_k}\]

• trong đó \(\gamma\) là tốc độ suy giảm (thường 0,99). Điều này tương đương với việc chạy thuật toán k-means trực tuyến trên các đầu ra của bộ mã hóa.

Sụp đổ codebook (Codebook Collapse)

• Một chế độ thất bại khét tiếng của VQ-VAE là sụp đổ codebook (codebook collapse) (còn gọi là sụp đổ chỉ số): mô hình học chỉ dùng một phần nhỏ trong \(K\) mục codebook, khiến hầu hết các mục bị "chết." Hãy tưởng tượng một thư viện nơi 90% kệ trống vì thủ thư luôn đặt sách vào một vài kệ phổ biến giống nhau. Điều này lãng phí khả năng biểu diễn.

• Sụp đổ codebook xảy ra vì bộ mã hóa, codebook và bộ giải mã cùng thích nghi trong quá trình huấn luyện. Nếu một mục không được chọn trong nhiều batch, nó trôi dạt khỏi đa tạp của bộ mã hóa, khiến nó càng ít có khả năng được chọn hơn, tạo ra một vòng phản hồi tích cực.

• Một số kỹ thuật giảm thiểu sụp đổ codebook:

  • Đặt lại codebook (Codebook reset): định kỳ khởi tạo lại các mục chết bằng cách sao chép các đầu ra bộ mã hóa được lấy mẫu ngẫu nhiên. Điều này cho các mục chết một khởi đầu mới gần vùng hoạt động của không gian tiềm ẩn.
  • Cập nhật EMA với làm mượt Laplace (Laplace smoothing): thêm một hằng số nhỏ vào \(\mathbf{n}_k\) để ngăn bất kỳ mục nào có số đếm bằng không, đảm bảo tất cả mục đều nhận được tín hiệu gradient.
  • Điều chỉnh mất mát cam kết: tăng \(\beta\) buộc các đầu ra bộ mã hóa cụm lại chặt chẽ hơn quanh các mục codebook, phân phối các gán đều hơn.
  • Mã phân tích (Factorised codes): phân rã tra cứu codebook thành tích của các tra cứu nhỏ hơn (ví dụ, hai codebook mỗi cái kích thước \(\sqrt{K}\)), cải thiện tỷ lệ sử dụng bằng cách giảm kích thước codebook hiệu quả cho mỗi tra cứu.
  • Chính quy hóa entropy (Entropy regularisation): thêm một hình phạt khuyến khích phân bố đồng đều trên việc sử dụng codebook, tối đa hóa entropy \(H = -\sum_k p_k \log p_k\) trong đó \(p_k\) là xác suất gán thực nghiệm.

VQ-GAN: Huấn luyện đối nghịch để đạt độ trung thực cao hơn

• VQ-VAE tạo ra các tái tạo khá tốt, nhưng mất mát \(\ell_2\) ở mức pixel có xu hướng sinh ra các đầu ra mờ vì nó phạt mọi độ lệch pixel như nhau, lấy trung bình trên các chi tiết hợp lý thay vì chọn ra những chi tiết sắc nét. Hãy tưởng tượng yêu cầu ai đó vẽ một khuôn mặt sao cho giảm thiểu độ khác biệt trung bình so với mọi khuôn mặt có thể — họ sẽ vẽ một khuôn mặt trung bình mờ, chứ không phải một cá nhân sắc nét.

• VQ-GAN (Esser et al., 2021) giải quyết vấn đề này bằng cách kết hợp khung VQ-VAE với một bộ phân biệt (discriminator) từ mạng đối nghịch sinh (generative adversarial network, GAN — Chương 6). Bộ phân biệt là một mạng tích chập theo vùng (patch-based) phán đoán xem một patch ảnh cục bộ là thật (từ dữ liệu huấn luyện) hay giả (từ bộ giải mã). Mất mát đối nghịch này khuyến khích bộ giải mã tạo ra các kết cấu sắc nét, chân thực về mặt nhận thức thị giác thay vì các trung bình từng pixel.

• Mục tiêu của VQ-GAN thêm hai số hạng vào mất mát VQ-VAE:

\[\mathcal{L}_\text{VQ-GAN} = \mathcal{L}_\text{VQ-VAE} + \lambda_\text{adv} \mathcal{L}_\text{adv} + \lambda_\text{perc} \mathcal{L}_\text{perc}\]

• Mất mát đối nghịch (adversarial loss) \(\mathcal{L}_\text{adv}\) là mục tiêu GAN chuẩn áp dụng cho đầu ra của bộ giải mã. Bộ phân biệt \(\mathcal{D}\) cố gắng phân biệt các patch thật khỏi các patch được giải mã, còn bộ giải mã (bộ sinh, generator) cố gắng đánh lừa nó. Công thức không bão hòa (non-saturating) là:

\[\mathcal{L}_\text{adv} = -\mathbb{E}[\log \mathcal{D}(D(\mathbf{z}_q))]\]

• Mất mát nhận thức (perceptual loss) \(\mathcal{L}_\text{perc}\) so sánh các kích hoạt đặc trưng từ một mạng đã huấn luyện trước (thường là VGG hoặc LPIPS) giữa ảnh gốc và ảnh tái tạo:

\[\mathcal{L}_\text{perc} = \sum_l \|\phi_l(\mathbf{x}) - \phi_l(D(\mathbf{z}_q))\|_2^2\]

• trong đó \(\phi_l\) ký hiệu bản đồ đặc trưng tại lớp \(l\) của mạng đã huấn luyện trước. Mất mát này nắm bắt sự tương đồng cấu trúc cấp cao thay vì độ chính xác từng pixel.

• Trọng số \(\lambda_\text{adv}\) được thiết lập thích ứng sao cho gradient đối nghịch và gradient tái tạo được cân bằng, ngăn mất mát đối nghịch chiếm ưu thế trong giai đoạn đầu huấn luyện khi các tái tạo còn kém.

• Kết quả là một bộ token hóa tạo ra các tái tạo sắc nét hơn hẳn so với VQ-VAE ở cùng kích thước codebook. VQ-GAN là bộ token hóa nền tảng đứng sau nhiều hệ thống sinh ảnh lớn, bao gồm DALL-E bản gốc, Parti và vô số mô hình văn bản-sinh-ảnh. Nó biến một ảnh 256x256 thành một lưới 16x16 hoặc 32x32 các token rời rạc từ một codebook kích thước 1024-16384, đạt tỷ lệ nén 16x đến 64x trên mỗi chiều không gian.

Lượng tử hóa thặng dư và codebook đa tỷ lệ

• Một codebook duy nhất áp đặt một trần cứng lên chất lượng tái tạo: mỗi vị trí không gian được biểu diễn bằng đúng một vector codebook, và mọi chi tiết chi tiết hơn mức codebook có thể biểu đạt đều bị mất. Hãy nghĩ về việc mô tả một màu bằng một từ đơn duy nhất từ một bảng màu cố định: "xanh ngọc" (teal) là gần đúng nhưng không chính xác. Nếu bạn có thể thêm một sự tinh chỉnh — "xanh ngọc, nhưng hơi xanh hơn một chút và sáng hơn một chút" — bạn sẽ tiến gần hơn nhiều.

• Lượng tử hóa thặng dư (residual quantisation, RQ) áp dụng ý tưởng này một cách lặp lại. Sau khi bước lượng tử hóa đầu tiên sinh ra \(\mathbf{z}_q^{(1)}\), tính thặng dư \(\mathbf{r}^{(1)} = \mathbf{z}_e - \mathbf{z}_q^{(1)}\), rồi lượng tử hóa thặng dư đó đối với một codebook thứ hai để có \(\mathbf{z}_q^{(2)}\), và cứ thế cho \(T\) mức:

\[\mathbf{r}^{(0)} = \mathbf{z}_e\]

\[\mathbf{z}_q^{(t)} = \text{Quantise}(\mathbf{r}^{(t-1)}, \mathcal{C}^{(t)})\]

\[\mathbf{r}^{(t)} = \mathbf{r}^{(t-1)} - \mathbf{z}_q^{(t)}\]

• Biểu diễn đã lượng tử hóa cuối cùng là \(\hat{\mathbf{z}} = \sum_{t=1}^{T} \mathbf{z}_q^{(t)}\). Với \(T\) mức, mỗi mức dùng một codebook kích thước \(K\), kích thước từ vựng hiệu dụng là \(K^T\), nhưng bạn chỉ cần lưu \(T \times K\) vector thay vì \(K^T\). Ví dụ, 8 mức với \(K = 1024\) cho kích thước hiệu dụng \(1024^8 \approx 10^{24}\) mục trong khi chỉ lưu 8192 vector.

• Mỗi mức liên tiếp nắm bắt các chi tiết tinh tế hơn: codebook đầu tiên nắm bắt cấu trúc thô, codebook thứ hai nắm bắt các hiệu chỉnh tần số trung bình, và cứ thế tiếp tục. Điều này tương tự như xấp xỉ lần lượt (successive approximation) trong JPEG hoặc kết xuất tiến dần (progressive rendering) trong ảnh web, nơi một phiên bản thô xuất hiện trước và chi tiết được lấp đầy dần dần.

• Codebook đa tỷ lệ (multi-scale codebooks) mở rộng ý tưởng này bằng cách hoạt động ở các độ phân giải không gian khác nhau. Thay vì lượng tử hóa cùng một lưới không gian nhiều lần, bạn lượng tử hóa ở nhiều tỷ lệ: một lưới thô nắm bắt cấu trúc toàn cục, các lưới tinh hơn nắm bắt chi tiết cục bộ. Điều này liên quan đến ý tưởng hình tháp đặc trưng (feature pyramid) từ phần phát hiện đối tượng của Chương 8, nơi các đặc trưng ở các tỷ lệ khác nhau nắm bắt các mức độ chi tiết khác nhau.

• Lượng tử hóa tích (product quantisation) là một kỹ thuật liên quan, trong đó vector tiềm ẩn \(d\) chiều được chia thành \(M\) vector con (sub-vector) mỗi cái có số chiều \(d/M\), và mỗi vector con được lượng tử hóa độc lập với codebook riêng của nó. Điều này cho một từ vựng hiệu dụng \(K^M\) trong khi chỉ lưu \(M \times K\) vector. Lượng tử hóa tích được dùng rộng rãi trong tìm kiếm láng giềng gần xấp xỉ (approximate nearest-neighbour search, Chương 13) và đã được thích ứng cho token hóa ảnh.

• Lượng tử hóa vô hướng hữu hạn (finite scalar quantisation, FSQ), được giới thiệu bởi Mentzer et al. (2023), tiếp cận hoàn toàn khác: thay vì học một codebook, nó đơn giản làm tròn mỗi chiều của vector tiềm ẩn về một trong một tập các mức nguyên cố định (ví dụ, \(\{-2, -1, 0, 1, 2\}\)). Với \(L\) mức mỗi chiều và \(d\) chiều, kích thước codebook ngầm định là \(L^d\). FSQ tránh sụp đổ codebook hoàn toàn vì không có vector codebook học được nào, chỉ có các đầu ra bộ mã hóa học được rồi được làm tròn một cách tất định. Bộ ước lượng đường thẳng xử lý sự không khả vi của phép làm tròn.

Các bộ token hóa ảnh trong thực tiễn

• Sự tiến triển từ VQ-VAE đến VQ-GAN rồi đến lượng tử hóa thặng dư đã sinh ra một họ các bộ token hóa ảnh thực tế được dùng trong các mô hình sinh tiên tiến nhất.

Bộ token hóa DALL-E (dVAE)

• DALL-E bản gốc (Ramesh et al., 2021) dùng một VAE rời rạc (dVAE) để token hóa các ảnh 256x256 thành lưới token 32x32 từ một codebook kích thước 8192. dVAE thay thế lượng tử hóa \(\arg\min\) cứng bằng một sự nới lỏng Gumbel-Softmax (Gumbel-Softmax relaxation), làm cho lượt truyền tiến khả vi trong quá trình huấn luyện. Tại thời điểm suy luận (inference), \(\arg\max\) được dùng để sinh ra các gán token cứng. dVAE được huấn luyện với sự kết hợp của mất mát tái tạo, phân kỳ KL (KL divergence) so với một tiên nghiệm (prior) đều, và một lịch trình nhiệt độ (temperature) học được cho Gumbel-Softmax. Sau đó DALL-E huấn luyện một transformer tự hồi quy 12 tỷ tham số để mô hình hóa phân bố đồng thời của 256 token văn bản và 1024 token ảnh (32x32).

LlamaGen

• LlamaGen (Sun et al., 2024) cho thấy bạn có thể tái sử dụng một kiến trúc mô hình ngôn ngữ kiểu Llama chuẩn (Chương 7) cho việc sinh ảnh tự hồi quy, miễn là bạn có một bộ token hóa ảnh tốt. LlamaGen dùng một bộ token hóa VQ-GAN cải tiến với codebook lớn (16.384 mục) và huấn luyện một transformer tự hồi quy thuần túy (không có bất kỳ sửa đổi đặc thù cho ảnh nào ngoài bộ token hóa) để dự đoán các token ảnh từ trái sang phải theo thứ tự quét raster (raster scan order). Điểm mấu chốt là: một khi ảnh được token hóa thành các chuỗi rời rạc, cùng mô hình dự đoán token tiếp theo (next-token-prediction) vốn hiệu quả cho ngôn ngữ cũng hiệu quả cho ảnh, khẳng định ý tưởng rằng token hóa thực sự thu hẹp khoảng cách giữa các phương thức (modality gap).

Cosmos Tokeniser

• Bộ token hóa Cosmos (NVIDIA, 2024) được thiết kế cho cả ảnh và video trong một khung thống nhất. Nó dùng một kiến trúc 3D nhân quả (causal 3D) coi ảnh như những video một khung hình, cho phép cùng một bộ token hóa xử lý cả hai phương thức. Cosmos hỗ trợ cả chế độ token hóa liên tục và rời rạc: chế độ liên tục xuất ra các vector tiềm ẩn thực (cho các backend mô hình khuếch tán), trong khi chế độ rời rạc áp dụng lượng tử hóa vô hướng hữu hạn để sinh ra các token nguyên (cho các backend mô hình tự hồi quy). Bộ mã hóa dùng các tích chập 3D nhân quả (causal 3D convolutions) sao cho token của mỗi khung hình chỉ phụ thuộc vào khung hình hiện tại và trước đó, cho phép token hóa video theo luồng (streaming).

Token hóa video

• Video thêm một trục thứ ba — thời gian — vào các chiều không gian của ảnh. Một video là một chuỗi các khung hình, thường ở tốc độ 24-30 khung hình mỗi giây, và các khung hình liền kề có tính dư thừa (redundancy) rất cao vì thế giới thị giác không thay đổi chóng mặt trong 33 miligiây. Token hóa video khai thác sự dư thừa thời gian này để đạt mức nén cao hơn nhiều so với token hóa từng khung hình độc lập.

• Hãy nghĩ về nén video như một cuốn sách lật (flip-book). Nếu bạn vẽ mọi trang từ đầu, bạn sẽ cần hàng ngàn bức vẽ chi tiết. Nhưng hầu hết các trang gần giống hệt trang lân cận, nên bạn có thể vẽ một "khung hình khóa" (keyframe) đầy đủ mỗi 10 trang và chỉ ghi chú những thay đổi nhỏ trên các trang ở giữa. Các bộ token hóa video học thủ thuật này một cách tự động.

3D VQ-VAE

• Sự mở rộng trực tiếp nhất của VQ-VAE sang video là 3D VQ-VAE, thay thế các tích chập 2D trong bộ mã hóa và bộ giải mã bằng các tích chập 3D hoạt động đồng thời trên cả chiều không gian và chiều thời gian. Nếu bộ mã hóa giảm mẫu với hệ số \(f_s\) theo không gian và \(f_t\) theo thời gian, một đoạn video \(T \times H \times W\) trở thành một lưới token \((T/f_t) \times (H/f_s) \times (W/f_s)\).

• Ví dụ, với \(f_s = 16\) và \(f_t = 4\), một đoạn video 16 khung hình 256x256 trở thành một chuỗi token \(4 \times 16 \times 16 = 1024\). Điều này đủ gọn để một transformer mô hình hóa theo kiểu tự hồi quy, trong khi số lượng pixel thô sẽ là \(16 \times 256 \times 256 \times 3 \approx 3{,}1\) triệu giá trị.

• Các tích chập 3D cùng học các đặc trưng không gian và thời gian. Các lớp đầu tiên nắm bắt chuyển động cục bộ (các cạnh di chuyển giữa các khung hình) trong khi các lớp sâu hơn nắm bắt các động lực cấp cao (các vật thể xuất hiện, biến mất, hoặc thay đổi hình dạng). Đây là cùng nguyên lý trích xuất đặc trưng phân cấp từ mạng tích chập của Chương 8, được mở rộng dọc theo trục thời gian.

Các bộ token hóa video nhân quả (Causal Video Tokenisers)

• Một tích chập 3D chuẩn nhìn vào các khung hình quá khứ, hiện tại và tương lai, nghĩa là bạn cần toàn bộ đoạn video trước khi có thể token hóa bất kỳ phần nào của nó. Các bộ token hóa video nhân quả (causal video tokenisers) ràng buộc các tích chập thời gian sao cho mỗi đầu ra chỉ phụ thuộc vào khung hình hiện tại và trước đó, không bao giờ phụ thuộc khung hình tương lai. Điều này tương tự như mặt nạ nhân quả (causal masking) trong các transformer tự hồi quy (Chương 7): thông tin chảy về phía trước theo thời gian nhưng không bao giờ ngược lại.

• Token hóa nhân quả là thiết yếu cho hai trường hợp sử dụng. Thứ nhất, theo luồng (streaming): bạn có thể token hóa video theo thời gian thực khi các khung hình đến, mà không cần đệm (buffer) các khung hình tương lai. Thứ hai, sinh tự hồi quy: khi một transformer sinh video từng khung hình một, các token cho khung hình \(t\) phải tính được mà không cần biết khung hình \(t+1\), vì khung hình \(t+1\) chưa được sinh ra.

• Ràng buộc nhân quả được triển khai bằng cách đệm (pad) các tích chập thời gian bất đối xứng: một hạt nhân (kernel) có kích thước thời gian \(k\) được đệm bằng \(k-1\) số không ở phía quá khứ và không số không ở phía tương lai, đảm bảo đầu ra tại thời điểm \(t\) chỉ phụ thuộc vào các đầu vào tại thời điểm \(t-k+1, \ldots, t\).

• Một tính chất tinh tế của các bộ token hóa video nhân quả là chúng có thể token hóa một ảnh đơn lẻ (một "video" một khung hình) mà không cần xử lý đặc biệt. Khung hình đầu tiên không có ngữ cảnh quá khứ, nên các token của nó được tính từ chính khung hình đó. Sự thống nhất ảnh-video (image-video unification) này nghĩa là một bộ token hóa duy nhất phục vụ cả hai phương thức, đơn giản hóa kiến trúc và cho phép các mô hình sinh ảnh và video với cùng một bộ giải mã.

Các chiến lược nén thời gian (Temporal Compression Strategies)

• Các ứng dụng khác nhau đòi hỏi các tỷ lệ nén thời gian khác nhau. Với nhận dạng hành động (nơi các chuyển động tinh tế có ý nghĩa), nén nhẹ (\(f_t = 2\)) bảo toàn chi tiết thời gian. Với sinh video dài (nơi việc lưu hàng ngàn khung hình là quá tốn kém), nén mạnh (\(f_t = 8\) hoặc cao hơn) là cần thiết.

• Một số bộ token hóa dùng nén phân tích (factorised compression): nén không gian và nén thời gian được thực hiện ở các giai đoạn riêng biệt. Đầu tiên, một bộ mã hóa 2D nén từng khung hình độc lập, sinh ra một lưới tiềm ẩn mỗi khung hình. Sau đó, một bộ mã hóa thời gian 1D nén theo chiều thời gian. Sự phân tích này rẻ hơn về tính toán so với tích chập 3D đầy đủ và cho phép các tỷ lệ nén khác nhau cho không gian và thời gian. Đánh đổi là nó không thể nắm bắt các mẫu không-thời-gian (như một quả bóng di chuyển chéo) hiệu quả bằng mã hóa 3D chung.

• Token nội suy thời gian (temporal interpolation tokens) là một đổi mới gần đây, nơi bộ token hóa chỉ mã hóa đầy đủ các khung hình khóa (keyframe) và biểu diễn các khung hình trung gian như các mã nội suy nhẹ (lightweight interpolation codes) mô tả cách biến dạng (morph) giữa các khung hình khóa. Điều này phản chiếu nén video cổ điển (các I-frame và P-frame trong H.264/HEVC) nhưng trong một không gian tiềm ẩn học được.

Token liên tục so với token rời rạc

• Không phải mọi mô hình hạ nguồn (downstream) đều cần token rời rạc. Mô hình khuếch tán (diffusion models) (Chương 10, file 04) làm việc bản địa với các giá trị liên tục — chúng từ từ khử nhiễu (denoise) một mẫu Gauss, và các hàm mất mát của chúng (denoising score matching) được định nghĩa trên các không gian liên tục. Với các backend khuếch tán, bộ mã hóa token hóa xuất ra các vector tiềm ẩn liên tục không bao giờ được lượng tử hóa. Mô hình khuếch tán tiềm ẩn (latent diffusion models) (Stable Diffusion, DALL-E 3, Flux) dùng một bộ mã hóa-giải mã giống VQ-GAN nhưng bỏ qua hoàn toàn codebook, hoạt động trong không gian tiềm ẩn liên tục.

• Mô hình tự hồi quy (autoregressive models) (kiểu GPT), mặt khác, dự đoán token tiếp theo từ một từ vựng hữu hạn bằng một softmax trên \(K\) lớp. Về bản chất chúng đòi hỏi token rời rạc. Mọi hệ thống sinh ảnh dùng một transformer tự hồi quy (DALL-E, Parti, LlamaGen, Chameleon) đều phụ thuộc vào một bộ token hóa rời rạc.

• Do đó, sự lựa chọn giữa token liên tục và rời rạc được thúc đẩy bởi backend sinh:

• Dùng token rời rạc khi: mô hình là tự hồi quy (dự đoán token tiếp theo với mất mát cross-entropy), bạn muốn chia sẻ một từ vựng với token văn bản cho các mô hình đa phương thức thống nhất, hoặc bạn cần kiểm soát chính xác ở mức token (ví dụ, cho truy xuất hoặc chỉnh sửa bằng cách thay thế token).

• Dùng token liên tục khi: mô hình là mô hình khuếch tán hoặc flow-matching, nhiệm vụ đòi hỏi tái tạo độ trung thực rất cao (các tiềm ẩn liên tục tránh hoàn toàn sai số lượng tử hóa), hoặc bạn muốn dùng các mất mát hồi quy (regression loss) hoạt động trên các vector thực.

• Một số kiến trúc gần đây hỗ trợ cả hai chế độ. Bộ token hóa Cosmos, chẳng hạn, có thể xuất ra hoặc là các tiềm ẩn liên tục (cho chế độ khuếch tán của nó) hoặc các token đã lượng tử hóa FSQ (cho chế độ tự hồi quy của nó) từ cùng một bộ mã hóa, với một đầu lượng tử hóa nhẹ có thể bật hoặc tắt.

• Lượng tử hóa mềm (soft quantisation) là một điểm giữa: thay vì gán cứng \(\arg\min\), tính một trung bình có trọng số của \(k\) mục codebook gần nhất, với trọng số cho bởi một softmax trên các khoảng cách âm. Điều này bảo toàn nhiều thông tin hơn lượng tử hóa cứng trong khi vẫn xấp xỉ rời rạc. Một số hệ thống dùng lượng tử hóa mềm trong huấn luyện và lượng tử hóa cứng ở suy luận.

Ứng dụng

Sinh ảnh tự hồi quy

• Một khi ảnh là các chuỗi token rời rạc, bạn có thể huấn luyện một transformer tự hồi quy chuẩn để mô hình hóa chúng. Các token ảnh được làm phẳng (flatten) thành một chuỗi 1D (thường theo thứ tự quét raster: trái sang phải, trên xuống dưới) và transformer học \(p(\text{token}_i | \text{token}_1, \ldots, \text{token}_{i-1})\) với mất mát cross-entropy chuẩn. Tại thời điểm sinh, các token được lấy mẫu từng cái một và lưới hoàn chỉnh được đưa qua bộ giải mã của token hóa để tạo ra pixel.

• Điều kiện hóa (conditioning) trên văn bản rất đơn giản: thêm trước (prepend) các token văn bản vào chuỗi token ảnh, để mô hình học \(p(\text{token ảnh} | \text{token văn bản})\). Đây chính xác là cách DALL-E, Parti và LlamaGen thực hiện sinh văn-bản-thành-ảnh (text-to-image). Các token văn bản và ảnh chia sẻ cùng một transformer, cùng một cơ chế chú ý, và thường là cùng một bảng embedding (với token văn bản và ảnh chiếm các dải chỉ số khác nhau).

• Thứ tự quét raster tạo ra một bất đối xứng nhân tạo: góc trên-trái của ảnh được sinh ra trước, không có bất kỳ ngữ cảnh nào về góc dưới-phải. Một số công trình giải quyết điều này. Mô hình ảnh có mặt nạ (masked image modelling) (MaskGIT) huấn luyện một transformer hai chiều (bidirectional) sinh tất cả token đồng thời nhưng với độ tự tin thay đổi, dần gỡ mặt nạ (unmask) các token tự tin nhất. Sinh đa tỷ lệ (multi-scale generation) sinh các token thô trước (nắm bắt bố cục toàn cục) rồi tinh chỉnh bằng các token thặng dư. Các cách tiếp cận này đánh đổi sự đơn giản của sinh thuần trái-sang-phải lấy sự gắn kết toàn cục tốt hơn.

Các token thị giác-ngôn ngữ thống nhất

• Động lực sâu xa nhất cho token hóa ảnh là thống nhất (unification): đưa thị giác và ngôn ngữ vào cùng một định dạng biểu diễn để một kiến trúc mô hình duy nhất xử lý cả hai. Như ta đã thảo luận ở Chương 7, các mô hình ngôn ngữ là những cỗ máy chuỗi-sang-chuỗi (sequence-to-sequence) cực kỳ mạnh. Bằng cách biểu diễn ảnh như các chuỗi token, chúng ta kế thừa toàn bộ hạ tầng của mô hình ngôn ngữ — các công thức huấn luyện trước (pretraining), các quy luật mở rộng (scaling laws), RLHF, mở rộng độ dài ngữ cảnh — một cách miễn phí.

• Chameleon (Meta, 2024) là một ví dụ nổi bật: nó dùng một bộ token hóa VQ-GAN với 8192 mục codebook để chuyển đổi ảnh thành các token được xen kẽ với token văn bản trong một từ vựng duy nhất khoảng 65.000 mục (văn bản + ảnh). Một transformer chuẩn được huấn luyện trên các chuỗi văn bản-ảnh hỗn hợp, cho phép nó sinh văn bản từ ảnh, ảnh từ văn bản, hoặc nội dung văn bản-và-ảnh đan xen, tất cả chỉ với cùng một lượt truyền tiến.

• Gemini (Google, 2024) áp dụng cách tiếp cận tương tự ở quy mô khổng lồ, hiểu và sinh ảnh, âm thanh, văn bản một cách bản địa trong một transformer duy nhất, với các bộ token hóa đặc thù theo phương thức đưa vào một chuỗi chia sẻ.

• Thách thức kỹ thuật cốt lõi trong các mô hình thống nhất là cân bằng từ vựng (vocabulary balance): nếu 8192 trên 65.000 mục từ vựng là token ảnh, mô hình có thể phân bổ không đủ năng lực cho thị giác. Các giải pháp bao gồm các lớp embedding riêng biệt cho mỗi phương thức (chỉ chia sẻ ở mức độ chú ý), trọng số mất mát đặc thù theo phương thức, và tỷ lệ trộn dữ liệu cẩn thận trong lúc huấn luyện trước.

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

1. Cài đặt một lớp VQ tối giản trong JAX: cho một batch các vector đầu ra bộ mã hóa, thực hiện tra cứu codebook láng giềng gần nhất và tính mất mát VQ-VAE (tái tạo + codebook + cam kết). Trực quan hóa việc sử dụng codebook dưới dạng một biểu đồ cột (histogram).

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   # --- Lớp VQ tối giản ---
   key = jax.random.PRNGKey(42)
   d = 8          # số chiều embedding
   K = 64         # kích thước codebook
   n_vectors = 256  # batch các đầu ra bộ mã hóa
   
   # Các đầu ra bộ mã hóa và codebook ngẫu nhiên
   k1, k2 = jax.random.split(key)
   z_e = jax.random.normal(k1, (n_vectors, d))       # đầu ra bộ mã hóa
   codebook = jax.random.normal(k2, (K, d)) * 0.1     # codebook (khởi tạo nhỏ)
   
   # Tra cứu láng giềng gần nhất: tìm mục codebook gần nhất cho mỗi z_e
   # distances[i, k] = ||z_e[i] - codebook[k]||^2
   distances = (
       jnp.sum(z_e ** 2, axis=1, keepdims=True)
       - 2 * z_e @ codebook.T
       + jnp.sum(codebook ** 2, axis=1, keepdims=True).T
   )
   indices = jnp.argmin(distances, axis=1)       # chỉ số token
   z_q = codebook[indices]                        # các vector đã lượng tử hóa
   
   # Các số hạng mất mát VQ-VAE
   beta = 0.25
   loss_codebook = jnp.mean((jax.lax.stop_gradient(z_e) - z_q) ** 2)
   loss_commit   = jnp.mean((z_e - jax.lax.stop_gradient(z_q)) ** 2)
   loss_total    = loss_codebook + beta * loss_commit
   print(f"Codebook loss: {loss_codebook:.4f}, Commitment loss: {loss_commit:.4f}")
   
   # Việc sử dụng codebook
   unique, counts = jnp.unique(indices, return_counts=True, size=K, fill_value=-1)
   plt.figure(figsize=(10, 4))
   plt.bar(range(K), counts, color='#3498db', alpha=0.8)
   plt.xlabel('Chỉ số codebook'); plt.ylabel('Số lượng gán')
   plt.title(f'Việc sử dụng codebook ({jnp.sum(counts > 0)}/{K} mục được dùng)')
   plt.grid(True, alpha=0.3); plt.tight_layout(); plt.show()
   # Thử: tăng K lên 512 và quan sát sự sụp đổ. Sau đó thêm logic đặt lại codebook.
   

2. Xây dựng một bộ lượng tử hóa vector 2D đồ chơi học cách lát (tile) một phân bố 2D. Sinh các điểm 2D ngẫu nhiên, học một codebook qua các cập nhật EMA, và trực quan hóa các vùng Voronoi.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   # Sinh dữ liệu 2D từ một hỗn hợp các phân phối Gauss
   key = jax.random.PRNGKey(0)
   n_points = 2000
   K = 16  # các mục codebook
   gamma = 0.99  # suy giảm EMA
   
   # Bốn cụm
   keys = jax.random.split(key, 5)
   centres = jnp.array([[2, 2], [-2, 2], [-2, -2], [2, -2]], dtype=jnp.float32)
   data = jnp.concatenate([
       jax.random.normal(keys[i], (n_points // 4, 2)) * 0.5 + centres[i]
       for i in range(4)
   ])
   
   # Khởi tạo codebook từ các điểm dữ liệu ngẫu nhiên
   idx = jax.random.choice(keys[4], n_points, (K,), replace=False)
   codebook = data[idx]
   ema_count = jnp.ones(K)
   ema_sum = codebook.copy()
   
   # Chạy học codebook dựa trên EMA trong vài epoch
   for epoch in range(30):
       # Gán mỗi điểm cho mục codebook gần nhất
       dists = jnp.sum((data[:, None, :] - codebook[None, :, :]) ** 2, axis=2)
       assignments = jnp.argmin(dists, axis=1)
       # Cập nhật EMA
       for k in range(K):
           mask = (assignments == k)
           count_k = jnp.sum(mask)
           ema_count = ema_count.at[k].set(gamma * ema_count[k] + (1 - gamma) * count_k)
           if count_k > 0:
               sum_k = jnp.sum(data[mask], axis=0)
               ema_sum = ema_sum.at[k].set(gamma * ema_sum[k] + (1 - gamma) * sum_k)
       codebook = ema_sum / ema_count[:, None]
   
   # Trực quan hóa các gán và codebook
   fig, ax = plt.subplots(1, 1, figsize=(8, 8))
   colors = plt.cm.tab20(jnp.linspace(0, 1, K))
   for k in range(K):
       mask = assignments == k
       ax.scatter(data[mask, 0], data[mask, 1], c=[colors[k]], s=5, alpha=0.3)
   ax.scatter(codebook[:, 0], codebook[:, 1], c='black', s=120, marker='X',
              edgecolors='white', linewidths=1.5, zorder=10, label='Codebook')
   ax.set_title(f'Codebook VQ đã học ({K} mục) trên dữ liệu 2D')
   ax.legend(); ax.set_aspect('equal'); ax.grid(True, alpha=0.3)
   plt.tight_layout(); plt.show()
   # Thử: tăng K lên 64 và quan sát sự lát tinh hơn. Giảm gamma và xem sự bất ổn.
   

3. Minh họa lượng tử hóa thặng dư: mã hóa một batch các vector với \(T\) giai đoạn lượng tử hóa liên tiếp và đo xem sai số tái tạo giảm thế nào qua mỗi mức.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   key = jax.random.PRNGKey(7)
   d = 16         # số chiều embedding
   K = 32         # kích thước codebook mỗi mức
   T = 8          # số mức thặng dư
   n_vectors = 512
   
   # Dữ liệu ngẫu nhiên để lượng tử hóa
   k1, *cb_keys = jax.random.split(key, T + 1)
   z = jax.random.normal(k1, (n_vectors, d))
   
   # Các codebook ngẫu nhiên độc lập cho mỗi mức
   codebooks = [jax.random.normal(cb_keys[t], (K, d)) * (0.5 ** t)
                for t in range(T)]
   
   # Vòng lặp lượng tử hóa thặng dư
   residual = z.copy()
   z_hat = jnp.zeros_like(z)
   errors = []
   
   for t in range(T):
       cb = codebooks[t]
       dists = (jnp.sum(residual ** 2, axis=1, keepdims=True)
                - 2 * residual @ cb.T
                + jnp.sum(cb ** 2, axis=1, keepdims=True).T)
       indices = jnp.argmin(dists, axis=1)
       z_q_t = cb[indices]
       z_hat = z_hat + z_q_t
       residual = residual - z_q_t
       mse = jnp.mean(jnp.sum((z - z_hat) ** 2, axis=1))
       errors.append(float(mse))
       print(f"Mức {t+1}: MSE = {mse:.4f}")
   
   plt.figure(figsize=(8, 5))
   plt.plot(range(1, T + 1), errors, 'o-', color='#e74c3c', linewidth=2, markersize=8)
   plt.xlabel('Mức lượng tử hóa thặng dư')
   plt.ylabel('MSE tái tạo')
   plt.title('Giảm lỗi với lượng tử hóa thặng dư')
   plt.xticks(range(1, T + 1)); plt.grid(True, alpha=0.3)
   plt.tight_layout(); plt.show()
   # Thử: dùng một codebook duy nhất kích thước K*T và so sánh với RQ. Bên nào thắng?
   

4. Mô phỏng một "bộ token hóa video" 1D đơn giản: sinh một chuỗi các tín hiệu 1D (bắt chước các khung hình video), áp dụng nén thời gian nhân quả, và so sánh với nén không nhân quả về mặt chất lượng tái tạo.

   import jax
   import jax.numpy as jnp
   import matplotlib.pyplot as plt
   
   key = jax.random.PRNGKey(99)
   n_frames = 16
   frame_len = 64
   
   # Sinh một "video": một gauss di chuyển chậm qua các khung hình
   x_axis = jnp.linspace(-3, 3, frame_len)
   frames = jnp.stack([
       jnp.exp(-0.5 * (x_axis - (-2 + 4 * t / n_frames)) ** 2)
       for t in range(n_frames)
   ])  # shape: (n_frames, frame_len)
   
   # Nén thời gian nhân quả: code của mỗi khung hình chỉ phụ thuộc vào các khung hình quá khứ
   # Cách tiếp cận đơn giản: trung bình khung hình hiện tại với suy giảm hàm mũ của quá khứ
   alpha_causal = 0.6
   causal_codes = jnp.zeros_like(frames)
   causal_codes = causal_codes.at[0].set(frames[0])
   for t in range(1, n_frames):
       causal_codes = causal_codes.at[t].set(
           alpha_causal * frames[t] + (1 - alpha_causal) * causal_codes[t - 1]
       )
   
   # Không nhân quả: trung bình với cả quá khứ và tương lai (làm mượt hai phía)
   kernel = jnp.array([0.2, 0.6, 0.2])  # quá khứ, hiện tại, tương lai
   padded = jnp.concatenate([frames[:1], frames, frames[-1:]], axis=0)
   noncausal_codes = jnp.stack([
       kernel[0] * padded[t] + kernel[1] * padded[t+1] + kernel[2] * padded[t+2]
       for t in range(n_frames)
   ])
   
   # Sai số tái tạo
   mse_causal = jnp.mean((frames - causal_codes) ** 2)
   mse_noncausal = jnp.mean((frames - noncausal_codes) ** 2)
   print(f"Causal MSE: {mse_causal:.6f}, Non-causal MSE: {mse_noncausal:.6f}")
   
   fig, axes = plt.subplots(1, 3, figsize=(15, 5))
   for ax, data, title in zip(axes,
       [frames, causal_codes, noncausal_codes],
       ['Các khung hình gốc', f'Nhân quả (MSE={mse_causal:.5f})',
        f'Không nhân quả (MSE={mse_noncausal:.5f})']):
       ax.imshow(data, aspect='auto', cmap='viridis', origin='lower')
       ax.set_xlabel('Vị trí không gian'); ax.set_ylabel('Chỉ số khung hình')
       ax.set_title(title)
   plt.tight_layout(); plt.show()
   # Thử: thay đổi alpha_causal và trọng số hạt nhân. Điều gì xảy ra với alpha=1.0?