Phân loại ảnh với Transformer: Attention, ViT

Mục tiêu

Sau bài này, bạn sẽ:

Hiểu vì sao self-attention trên các patch cho mỗi vùng ảnh nhìn được cả ảnh ngay từ lớp đầu, khác với trường tiếp nhận (receptive field) lớn dần qua từng lớp của CNN.
Mô tả được đường đi của một ảnh qua Vision Transformer (ViT): patch embedding, token lớp (class token, [CLS]), embedding vị trí, bộ mã hóa Transformer, đầu phân loại. Tính được số token và chi phí $O(N^2)$ theo số patch.
Giải thích được thiên kiến quy nạp (inductive bias) và vì sao ViT cần nhiều dữ liệu (JFT-300M), còn DeiT học được chỉ trên ImageNet nhờ augmentation mạnh và chắt lọc tri thức (knowledge distillation). Nêu được ý chính của Swin: attention trong cửa sổ cục bộ được dịch chuyển.
Cài được phép cắt patch và scaled dot-product attention bằng NumPy, rồi huấn luyện một ViT tí hon bằng PyTorch và so với một CNN nhỏ.

Cần biết trước

Phân loại ảnh với CNN: AlexNet → VGG → ResNet → MobileNet → EfficientNet, thuộc phần Thị giác máy tính (Đã học)
Transformers, thuộc phần Học sâu (Đã học)

Bài Transformer của phần Học sâu chưa xuất bản, nên nhắc ở đây phần bài này dùng; bản dài hơn, có mặt nạ nhân quả, ở bài Xây dựng LLM đầu tiên. Self-attention cho mỗi vị trí trong một dãy vectơ nhìn mọi vị trí của chính dãy đó: nó chấm điểm mức liên quan giữa từng cặp, đổi điểm thành trọng số bằng softmax, rồi lấy trung bình có trọng số. Mỗi vectơ được nhân với ba ma trận học được để thành truy vấn, khóa, giá trị (query, key, value) $\mathbf{Q}$ , $\mathbf{K}$ , $\mathbf{V}$ : truy vấn là “tôi đang tìm gì”, khóa là “tôi có gì”, giá trị là “thông tin tôi đưa ra”. Attention nhiều đầu (multi-head attention) chạy $h$ phép như vậy song song, mỗi đầu $d = D/h$ chiều, rồi ghép lại, để mỗi đầu theo dõi một kiểu quan hệ. Một khối Transformer gồm attention nhiều đầu và một mạng truyền thẳng nhỏ (MLP), mỗi phần có chuẩn hóa lớp (layer normalization) ở đầu vào và kết nối tắt (residual connection) cộng đầu vào vào đầu ra. Cắt ảnh thành patch và chiếu mỗi patch thành một vectơ đã có ở bài Multimodal; bài này đi tiếp phần ViT dùng các vectơ đó để phân loại.

Trực giác

Tai mèo ở góc trên, đuôi ở góc dưới

Bạn nhìn một bức ảnh con mèo đang vươn vai: tai ở góc trên bên trái, đuôi cong ở góc dưới bên phải. Muốn chắc đây là một con mèo chứ không phải hai con vật, mắt bạn phải nối hai chỗ cách xa nhau đó lại.

Ở bài Phân loại ảnh với CNN, mỗi lớp tích chập chỉ nhìn một ô nhỏ, ví dụ $3 \times 3$ điểm ảnh. Thông tin ở tai muốn gặp thông tin ở đuôi thì phải đi qua nhiều lớp, mỗi lớp lan ra thêm một chút, như tin đồn truyền qua hàng xóm. Pooling và stride giúp lan nhanh hơn, nhưng hai vùng xa nhau vẫn chỉ “gặp” nhau ở các lớp sâu.

Giờ hãy hình dung một cuộc họp. Cắt ảnh thành lưới ô vuông (patch), mỗi ô là một người dự họp. Ngay phút đầu, mỗi người được hỏi thẳng mọi người khác: “bạn có gì liên quan tới tôi không?”. Ô chứa tai hỏi được ô chứa đuôi ngay lập tức. Đó là self-attention trên các patch, và là ý chính của ViT.

Cái giá của cuộc họp

Cuộc họp mà ai cũng nói với ai có hai cái giá.

Thứ nhất, số cặp tăng theo bình phương. 4 người thì 16 cặp (kể cả tự hỏi mình), 200 người thì 40 nghìn cặp. Chia ảnh thành ô nhỏ hơn để thấy chi tiết thì số người tăng, và số cặp tăng nhanh hơn nhiều.

Thứ hai, không ai được dặn trước “hãy để ý người ngồi cạnh”. CNN có sẵn giả định đó trong kiến trúc: điểm ảnh gần nhau thì liên quan, và một vật dịch sang trái vẫn là vật đó. ViT phải tự học điều này từ dữ liệu. Có đủ ảnh thì tự học lại là lợi thế, vì mô hình không bị bó buộc; ít ảnh thì nó học kém hơn CNN.

Bài toán của bài này là: làm sao dùng một Transformer, vốn làm cho chữ, để phân loại ảnh; nó tốn bao nhiêu khi ảnh lớn; và khi nào nó thắng, khi nào nó thua một CNN?

Lý thuyết

Ký hiệu

Ảnh $\mathbf{X} \in \mathbb{R}^{H \times W \times C}$ có chiều cao $H$ , chiều rộng $W$ , $C$ kênh màu. Cỡ patch là $P$ , số patch là $N = (H/P)(W/P)$ , số chiều của mỗi token là $D$ , như bài Multimodal. Bộ mã hóa có $J$ khối, đánh số $\ell = 1, \dots, J$ (bài báo ViT viết $L$ ; bài này để dành $L$ cho độ dài chuỗi như phần LLM). Mỗi khối có $h$ đầu attention, mỗi đầu $d = D/h$ chiều.

Bước 1: từ ảnh tới dãy token

Phần này bài Multimodal đã dẫn từng bước, nên ở đây chỉ viết lại gọn. Ảnh được cắt thành $N$ patch không chồng nhau, mỗi patch trải phẳng thành vectơ hàng $\mathbf{x}_i \in \mathbb{R}^{P^2 C}$ , rồi nhân với một ma trận chung $\mathbf{E} \in \mathbb{R}^{(P^2 C) \times D}$ . ViT đặt thêm token [CLS] lên đầu dãy và cộng embedding vị trí (Dosovitskiy và cộng sự, 2021, phương trình 1):

\mathbf{Z}_0 = \begin{bmatrix} \mathbf{x}_{\text{cls}} \\ \mathbf{x}_1 \mathbf{E} \\ \vdots \\ \mathbf{x}_N \mathbf{E} \end{bmatrix} + \mathbf{E}_{\text{pos}}, \qquad \mathbf{Z}_0, \mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times D}

Đọc từng phần:

$\mathbf{x}_{\text{cls}}$ là một vectơ $D$ chiều học được, giống nhau cho mọi ảnh. Nó không mang nội dung ảnh nào. Chính vì “trống” như vậy, nó là chỗ tốt để gom thông tin từ mọi patch qua các lớp attention.
$\mathbf{E}_{\text{pos}}$ có $N + 1$ hàng, mỗi hàng là embedding vị trí $\mathbf{p}_i$ của bài Multimodal, cộng cho token ở vị trí đó. Dãy có $N + 1$ hàng nên ma trận vị trí cũng có $N + 1$ hàng.
Với chữ số $8 \times 8$ của load_digits ( $C = 1$ ) và $P = 4$ : $N = 2 \cdot 2 = 4$ patch, mỗi patch $16$ số, nên dãy có 5 token. Ô NumPy trong bài làm đúng phép này.

Bước 2: self-attention trên các patch

Mỗi khối cho dãy token đi qua attention. Với một đầu, ba phép chiếu học được cho truy vấn, khóa, giá trị:

\mathbf{Q} = \mathbf{Z}\mathbf{W}_Q, \quad \mathbf{K} = \mathbf{Z}\mathbf{W}_K, \quad \mathbf{V} = \mathbf{Z}\mathbf{W}_V, \qquad \mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V \in \mathbb{R}^{D \times d}

Ở đây $\mathbf{K}$ đậm là ma trận khóa, khác kernel tích chập của bài CNN. Rồi scaled dot-product attention:

\mathbf{A} = \mathrm{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d}}\right), \qquad \mathrm{Attn}(\mathbf{Z}) = \mathbf{A}\mathbf{V}

Đọc từng phần:

$\mathbf{Q}\mathbf{K}^\top$ có cỡ $(N+1) \times (N+1)$ . Ô $(i, j)$ là tích vô hướng của truy vấn ở token $i$ với khóa ở token $j$ : patch $i$ thấy patch $j$ liên quan bao nhiêu.
Chia $\sqrt{d}$ giữ cho điểm không quá lớn khi $d$ lớn, để softmax không dồn hết về một ô (lý do chi tiết ở bài Xây dựng LLM đầu tiên).
Softmax theo từng hàng đổi điểm thành trọng số dương, cộng lại bằng 1. Hàng $i$ của $\mathbf{A}$ là cách patch $i$ chia sự chú ý của nó cho mọi token.
$\mathbf{A}\mathbf{V}$ : hàng $i$ của kết quả là trung bình có trọng số của mọi hàng giá trị. Patch $i$ “đọc” thông tin của cả ảnh.

Không có mặt nạ. GPT che các vị trí ở sau, vì khi sinh chữ thì tương lai chưa có. Ảnh thì có sẵn cả tấm, không có “tương lai”, nên ViT dùng encoder: patch ở góc dưới đọc được patch ở góc trên và ngược lại.

Trường tiếp nhận toàn cục từ lớp đầu. Vì $e^{s} > 0$ với mọi số $s$ , mọi trọng số $A_{ij}$ đều dương: mỗi token nhận thông tin từ mọi token, ngay ở khối thứ nhất. So với CNN: một lớp $3 \times 3$ , stride 1, cho mỗi điểm nhìn $3 \times 3$ điểm ảnh. Lớp thứ hai nhìn $3 \times 3$ điểm của lớp trước, mỗi điểm đó lại nhìn $3 \times 3$ , nên vùng nhìn nới thêm 1 điểm ảnh mỗi phía. Sau $l$ lớp, trường tiếp nhận là $(2l + 1) \times (2l + 1)$ . Muốn phủ cả chữ số $8 \times 8$ phải cần 4 lớp (tự tính: $2 \cdot 4 + 1 = 9 \ge 8$ ), còn ảnh $224 \times 224$ thì cần tới 112 lớp nếu không có pooling hay stride (tự tính).

“Nhìn được” chưa có nghĩa là “nhìn”. Bài báo ViT đo khoảng cách attention (attention distance) trung bình của từng đầu, một đại lượng giống trường tiếp nhận của CNN. Ở các lớp thấp, có đầu nhìn gần như cả ảnh, có đầu chỉ nhìn rất gần, giống các lớp tích chập đầu; càng lên sâu, khoảng cách càng lớn (mục 4.5). Mô hình tự học khi nào nhìn xa, khi nào nhìn gần.

Nhiều đầu. $h$ đầu chạy song song với bộ $\mathbf{W}_Q, \mathbf{W}_K, \mathbf{W}_V$ riêng, mỗi đầu cho $N + 1$ hàng $d$ chiều. Ghép $h$ kết quả được lại $D = hd$ chiều, rồi nhân một ma trận $D \times D$ để trộn các đầu. ViT-Base dùng $D = 768$ , $h = 12$ , nên $d = 64$ .

Bước 3: khối encoder và đầu phân loại

Mỗi khối $\ell$ làm hai bước, mỗi bước có chuẩn hóa lớp ( $\mathrm{LN}$ ) ở đầu vào và kết nối tắt ở đầu ra (phương trình 2 và 3 của bài báo):

\begin{aligned} \mathbf{Z}'_\ell &= \mathrm{MSA}\big(\mathrm{LN}(\mathbf{Z}_{\ell-1})\big) + \mathbf{Z}_{\ell-1} \\ \mathbf{Z}_\ell &= \mathrm{MLP}\big(\mathrm{LN}(\mathbf{Z}'_\ell)\big) + \mathbf{Z}'_\ell \end{aligned}

$\mathrm{MSA}$ là attention nhiều đầu (multi-head self-attention) của Bước 2. $\mathrm{MLP}$ gồm hai lớp tuyến tính với hàm kích hoạt GELU ở giữa, áp riêng cho từng token; lớp ẩn của ViT-Base rộng $3\,072 = 4D$ . Cách đặt $\mathrm{LN}$ trước mỗi phần (Pre-LN) giống GPT-2 ở bài Xây dựng LLM đầu tiên. Dấu cộng giữ cho mỗi khối chỉ thêm một phần hiệu chỉnh, nên chồng nhiều khối vẫn huấn luyện được.

Sau $J$ khối, ViT chỉ lấy hàng đầu, tức trạng thái của [CLS], chuẩn hóa thêm một lần rồi đưa vào đầu phân loại (phương trình 4):

\mathbf{y} = \mathrm{LN}\big(\mathbf{Z}_J[0]\big), \qquad \text{logit} = \mathbf{y}\,\mathbf{W}_{\text{head}} + \mathbf{b}_{\text{head}}

$\mathbf{Z}_J[0]$ là hàng 0 (token [CLS]) của $\mathbf{Z}_J$ . $\mathbf{W}_{\text{head}}$ có $D$ hàng và số cột bằng số lớp; 10 lớp với chữ số. Softmax của logit cho xác suất từng lớp, và hàm mất mát là cross-entropy như mọi bộ phân loại. Bài báo dùng một MLP một lớp ẩn làm đầu phân loại khi pre-train, và thay bằng một lớp tuyến tính khởi tạo bằng 0 khi fine-tune sang bài toán mới. Thay vì [CLS], có thể lấy trung bình mọi token patch (global average pooling) như ResNet; phụ lục của bài báo cho thấy hai cách cho kết quả tương tự khi chỉnh tốc độ học cho hợp.

Ba cỡ trong bài báo (bảng 1). Tên như ViT-B/16 nghĩa là cỡ Base, patch $16 \times 16$ .

Cỡ	Số khối $J$	$D$	Lớp ẩn MLP	Số đầu $h$	Tham số
ViT-Base	12	768	3 072	12	86 triệu
ViT-Large	24	1 024	4 096	16	307 triệu
ViT-Huge	32	1 280	5 120	16	632 triệu

Chi phí: $O(N^2)$ theo số patch

Bảng điểm $\mathbf{Q}\mathbf{K}^\top$ có $(N+1)^2$ ô cho mỗi đầu ở mỗi khối. Mỗi ô là một tích vô hướng $d$ chiều, nên một đầu tốn khoảng $(N+1)^2 d$ phép nhân, và $h$ đầu tốn $(N+1)^2 hd = (N+1)^2 D$ . Nhân $\mathbf{A}$ với $\mathbf{V}$ tốn thêm chừng đó. Các phép chiếu và MLP chỉ tốn tỉ lệ với $N$ , vì chúng làm riêng trên từng token. Khi $N$ lớn, phần bình phương thắng.

Thay $N = HW/P^2$ vào: giảm nửa cỡ patch hay gấp đôi mỗi cạnh ảnh thì $N$ gấp 4, bảng điểm gấp khoảng 16. Ô NumPy bên dưới in ra: ảnh $224 \times 224$ với $P = 16$ có 197 token và $38\,809$ cặp; với $P = 8$ có 785 token và $616\,225$ cặp, gấp khoảng 15.9 lần. Vì vậy ViT dùng patch tương đối lớn (14 hay 16 điểm ảnh), và tên bài báo là “một tấm ảnh đáng giá $16 \times 16$ từ”.

Bài báo hay fine-tune ở ảnh lớn hơn lúc pre-train. Khi đó $P$ giữ nguyên, $N$ tăng, và embedding vị trí học được cho lưới cũ không còn khớp: tác giả nội suy hai chiều (2D interpolation) các embedding đó theo vị trí của chúng trên ảnh (mục 3.2).

Thiên kiến quy nạp: vì sao ViT cần nhiều dữ liệu

Thiên kiến quy nạp là những giả định có sẵn trong mô hình trước khi thấy dữ liệu. Bài báo ViT chỉ ra ba giả định của CNN, có ở mọi lớp: tính cục bộ (locality, điểm gần nhau liên quan nhau), cấu trúc lân cận hai chiều, và đẳng biến tịnh tiến (translation equivariance: dịch ảnh thì bản đồ đặc trưng dịch theo, vì cùng một kernel trượt khắp ảnh). Ở ViT, chỉ MLP là cục bộ và đẳng biến tịnh tiến, còn attention là toàn cục. Cấu trúc hai chiều chỉ được dùng ở hai chỗ: lúc cắt patch, và lúc nội suy embedding vị trí. Lúc khởi tạo, embedding vị trí không chứa thông tin gì về vị trí hai chiều; quan hệ không gian giữa các patch phải học từ đầu.

Ít giả định thì mô hình linh hoạt hơn nhưng cần nhiều ví dụ hơn để tìm ra quy luật, đúng cái cân bằng của bài Overfitting và underfitting. Bài báo kiểm chứng bằng ba bộ dữ liệu pre-train (mục 4.1 và 4.3):

Bộ pre-train	Số ảnh	Số lớp	Kết quả của ViT (Dosovitskiy và cộng sự, 2021)
ImageNet	1.3 triệu	1 000	Kém ResNet cùng cỡ vài điểm phần trăm; ViT-Large còn kém ViT-Base
ImageNet-21k	14 triệu	21 000	ViT-Large và ViT-Base ngang nhau
JFT-300M (nội bộ Google)	303 triệu	18 000	Mô hình lớn mới thể hiện hết; vượt các mô hình ResNet làm mốc

Thí nghiệm thứ hai rõ hơn: trên tập con 9 triệu ảnh của JFT, ViT-B/32 kém hẳn ResNet-50 có chi phí tính tương đương; từ 90 triệu ảnh trở lên thì ViT-B/32 tốt hơn. Kết luận của tác giả: thiên kiến của tích chập có ích khi dữ liệu nhỏ, còn với dữ liệu lớn thì học thẳng từ dữ liệu là đủ, thậm chí tốt hơn.

DeiT: ViT chỉ với ImageNet

JFT-300M là dữ liệu nội bộ, không phải ai cũng có. DeiT (Touvron và cộng sự, 2021) huấn luyện ViT chỉ trên ImageNet, trên một máy 8 GPU trong khoảng 53 giờ pre-train. Hai thành phần chính:

Augmentation và regularization mạnh. Tác giả ghi rằng Transformer cần nhiều dữ liệu hơn mô hình có sẵn giả định như tích chập, nên họ “tạo thêm dữ liệu” bằng RandAugment, Mixup, CutMix, random erasing, repeated augmentation, cộng stochastic depth. Gần như mọi kiểu augmentation họ thử đều có ích. Các phép này có ở bài Data augmentation trong CV.
Token chắt lọc (distillation token). Thêm một token học được nữa, đứng cạnh [CLS], cũng đi qua mọi lớp attention. Đầu ra của [CLS] học nhãn thật $y$ ; đầu ra của token chắt lọc học nhãn mà một mô hình thầy (teacher) đã huấn luyện xong dự đoán, $y_t = \arg\max_c Z_t(c)$ với $Z_t$ là logit của thầy. Bản “nhãn cứng” (hard distillation) lấy trung bình hai cross-entropy:

\mathcal{L} = \tfrac{1}{2}\,\mathrm{CE}\big(\mathrm{softmax}(\mathbf{z}_{\text{cls}}), y\big) + \tfrac{1}{2}\,\mathrm{CE}\big(\mathrm{softmax}(\mathbf{z}_{\text{dist}}), y_t\big)

$\mathbf{z}_{\text{cls}}$ và $\mathbf{z}_{\text{dist}}$ là logit tính từ hai token đó. Điều thú vị: thầy là một CNN (RegNetY-16GF) thì trò học tốt hơn so với khi thầy là một Transformer có độ chính xác tương đương. Một cách hiểu là CNN truyền một phần thiên kiến quy nạp của nó sang ViT qua các nhãn nó đoán. Kết quả trong tóm tắt bài báo: mô hình tham chiếu 86 triệu tham số đạt 83.1% top-1 trên ImageNet không dùng dữ liệu ngoài, và tới 85.2% khi có chắt lọc.

Swin: attention trong cửa sổ dịch chuyển

ViT gốc có hai điểm chưa hợp với các bài toán dày đặc như phân đoạn ảnh: chi phí bình phương theo số patch, và chỉ có một độ phân giải đặc trưng. Swin Transformer (Liu và cộng sự, 2021) sửa cả hai:

Cửa sổ cục bộ. Patch nhỏ $4 \times 4$ ; attention chỉ tính trong từng cửa sổ $M \times M$ patch không chồng nhau ( $M = 7$ mặc định). Với lưới $h \times w$ patch và $C$ kênh (ký hiệu của bài báo, ứng với $D$ ở đây), bài báo đếm chi phí là $4hwC^2 + 2(hw)^2C$ cho attention toàn cục, và $4hwC^2 + 2M^2hwC$ cho attention trong cửa sổ. Số hạng thứ hai đổi từ bình phương sang tuyến tính theo số patch $hw$ khi $M$ cố định.
Cửa sổ dịch chuyển (shifted window). Cửa sổ không chồng nhau thì patch ở hai cửa sổ kề nhau không bao giờ nói chuyện. Swin xen kẽ hai cách chia: khối này chia lưới từ góc trên trái, khối sau dời lưới đi $\lfloor M/2 \rfloor$ patch mỗi chiều, nên mỗi cửa sổ mới vắt ngang các cửa sổ cũ.
Phân tầng. Sau mỗi giai đoạn, lớp gộp patch (patch merging) ghép từng nhóm $2 \times 2$ patch kề nhau thành một, nên số token giảm 4 lần còn số kênh tăng gấp đôi, giống cách CNN giảm độ phân giải và tăng số kênh.

Để ý rằng Swin đưa tính cục bộ và cấu trúc phân tầng của CNN quay lại vào Transformer. Bài báo báo cáo 87.3% top-1 trên ImageNet-1K với bản lớn nhất Swin-L (pre-train trên ImageNet-22K), và dùng Swin làm xương sống cho phát hiện vật thể và phân đoạn ảnh.

Giới hạn

Bài chỉ nói về phân loại. Phân đoạn ảnh dùng xương sống CNN hay Transformer có ở bài Semantic Segmentation.
ViT học không cần nhãn (như MAE, DINO) thuộc bài Self-supervised learning với ảnh; ViT học cùng chữ thuộc bài CLIP, BLIP, GLIP.
Các mẹo tiết kiệm attention khác (FlashAttention, attention thưa) không có trong bài; ý chung là cùng phép tính nhưng ít đọc ghi bộ nhớ hơn, hoặc bỏ bớt cặp.

Minh họa

Hình dưới vẽ đường đi của một ảnh qua ViT, từ trái sang phải. Hãy nhìn vào token [CLS]: nó vào bộ mã hóa không mang nội dung ảnh nào, và đi ra là thứ duy nhất được đưa vào lớp phân loại.

Bảng dưới đặt CNN và các họ ViT của bài cạnh nhau. Các dòng về dữ liệu là kết luận của bài báo ViT và DeiT, không phải luật cứng.

Tiêu chí	CNN (ResNet)	ViT	DeiT	Swin
Phép chính	Tích chập kernel nhỏ	Self-attention toàn cục	Như ViT, thêm token chắt lọc	Self-attention trong cửa sổ dịch chuyển
Vùng nhìn ở lớp đầu	Cục bộ ( $3 \times 3$ , $7 \times 7$ )	Cả ảnh	Cả ảnh	Một cửa sổ $M \times M$ patch
Thiên kiến quy nạp	Mạnh: cục bộ, đẳng biến tịnh tiến	Yếu: chỉ ở bước cắt patch và vị trí	Yếu, bù bằng augmentation và thầy CNN	Vừa: cục bộ và phân tầng quay lại
Chi phí theo số patch $N$	Tuyến tính theo số điểm ảnh	$O(N^2)$	$O(N^2)$	Tuyến tính khi $M$ cố định
Dữ liệu cần để học tốt	Ít tới vừa	Rất nhiều (14 triệu tới 300 triệu ảnh)	ImageNet 1.3 triệu ảnh, kèm augmentation mạnh	ImageNet-1K hoặc 22K
Độ phân giải đặc trưng	Nhiều tầng	Một tầng	Một tầng	Nhiều tầng

Ví dụ code

Code trong bài chạy thử lần cuối tháng 11/2026: ô NumPy với Python 3.14, NumPy 2.4 và scikit-learn 1.8 (bản có trong Pyodide), ô PyTorch với Python 3.14, PyTorch 2.15 và scikit-learn 1.9 trên CPU. Ô NumPy chạy ngay trên trang, ô PyTorch chạy trên máy bạn. Hai ô độc lập, không dùng biến của nhau. Ô NumPy lần đầu chạy lâu hơn, vì trình duyệt phải tải thêm scikit-learn để lấy load_digits.

Cắt patch và attention bằng NumPy

Ô dưới đây lấy chữ số đầu tiên của load_digits (một chữ số 0, ảnh $8 \times 8$ ), cắt thành 4 patch $4 \times 4$ , chiếu mỗi patch thành vectơ $D = 8$ chiều bằng một ma trận ngẫu nhiên (hạt giống 0), thêm [CLS] và embedding vị trí, rồi tính một đầu attention. Các ma trận ở đây ngẫu nhiên, chưa học gì: ô chỉ kiểm hình dạng và tính chất của phép tính, không kiểm ý nghĩa.

Cắt chữ số thành patch, tính self-attention

# Cắt một chữ số 8x8 thành patch rồi tính self-attention bằng NumPy.
import numpy as np
from sklearn.datasets import load_digits

rng = np.random.default_rng(0)
digits = load_digits()
img = digits.images[0] / 16.0         # ảnh 8x8, giá trị từ 0 tới 1
print("label:", digits.target[0], "image:", img.shape)

P, D = 4, 8                            # cỡ patch, số chiều embedding
H, W = img.shape
# Cắt thành lưới (H/P) x (W/P) ô, đọc theo hàng, trải phẳng mỗi ô
patches = (img.reshape(H // P, P, W // P, P)
           .transpose(0, 2, 1, 3).reshape(-1, P * P))
N = patches.shape[0]
print("patches:", patches.shape, "N =", N)
print("patch 1 = top-right block:",
      np.array_equal(patches[1], img[:P, P:].ravel()))

E = rng.normal(0, 0.3, (P * P, D))     # phép chiếu dùng chung cho mọi patch
cls = rng.normal(0, 0.3, (1, D))       # token [CLS] học được
pos = rng.normal(0, 0.3, (N + 1, D))   # embedding vị trí học được
Z = np.vstack([cls, patches @ E]) + pos
print("tokens Z:", Z.shape)

# Một đầu attention: truy vấn, khóa, giá trị
Wq, Wk, Wv = (rng.normal(0, 0.5, (D, D)) for _ in range(3))
Q, K, V = Z @ Wq, Z @ Wk, Z @ Wv
S = Q @ K.T / np.sqrt(D)               # điểm của mọi cặp token
S = S - S.max(axis=1, keepdims=True)   # trừ max cho ổn định số
A = np.exp(S) / np.exp(S).sum(axis=1, keepdims=True)
out = A @ V

np.set_printoptions(precision=3, suppress=True)
print("attention A (row = query, col = key):")
print(A)
print("row sums:", A.sum(axis=1))
print("all weights > 0:", bool((A > 0).all()))
print("output:", out.shape)
# Số cặp điểm attention tăng theo bình phương số token
for side, p in [(8, 4), (8, 2), (224, 16), (224, 8)]:
    n = (side // p) ** 2 + 1           # thêm 1 cho [CLS]
    print(f"{side}x{side}, P={p}: {n} tokens, {n * n} pairs")

Lần chạy đầu cần tải Python, khoảng 10 MB.

Đọc kết quả:

patches: (4, 16) N = 4: $N = (8/4)^2 = 4$ patch, mỗi patch $4 \cdot 4 = 16$ số. Dòng True tiếp theo xác nhận patch 1 là khối trên bên phải: đọc hết hàng rồi mới xuống.
tokens Z: (5, 8): 4 token patch cộng [CLS] ở hàng 0, mỗi token 8 chiều, đúng cỡ $(N+1) \times D$ của $\mathbf{Z}_0$ .
Ma trận $\mathbf{A}$ cỡ $5 \times 5$ : hàng là token hỏi (truy vấn), cột là token được hỏi (khóa). Hàng 0 là [CLS]: nó dồn 0.457 cho token 2 (patch 1) và chia phần còn lại cho các token khác. Với trọng số ngẫu nhiên, con số này không mang nghĩa gì; sau khi học, đây là chỗ [CLS] quyết định đọc patch nào.
row sums: [1. 1. 1. 1. 1.] và all weights > 0: True: softmax làm mỗi hàng là một phân phối, và không ô nào bằng 0. Mỗi patch nhận thông tin từ cả 5 token ngay ở lớp này: trường tiếp nhận toàn cục.
Bảng cuối: số cặp là $(N+1)^2$ . Chữ số $8 \times 8$ với $P = 2$ đã có 17 token và 289 cặp; ảnh $224 \times 224$ với $P = 8$ có $616\,225$ cặp, gấp khoảng 15.9 lần $38\,809$ của $P = 16$ .

ViT tí hon so với CNN nhỏ bằng PyTorch

File dưới đây dựng một ViT tí hon: patch $2 \times 2$ trên chữ số $8 \times 8$ ( $N = 16$ , cộng [CLS] là 17 token), $D = 32$ , 2 khối, 4 đầu, lớp ẩn MLP $2D$ . Patch embedding viết bằng Conv2d có kernel bằng stride bằng $P$ , như bài Multimodal đã kiểm. Bộ mã hóa dùng nn.TransformerEncoderLayer với norm_first=True (Pre-LN, đúng phương trình ở Bước 3). Để so sánh, file dựng thêm một CNN nhỏ hai lớp tích chập. Cả hai học 60 epoch trên 1 437 ảnh, kiểm trên 360 ảnh còn lại, cùng bộ tối ưu AdamW. Cách cài PyTorch bản CPU có ở bài Xây dựng LLM đầu tiên.

Chạy trên máy bạn

ViT tí hon và CNN nhỏ trên load_digits

Cần cài: pip install torch scikit-learn
Phần cứng: CPU là đủ, không cần GPU. Khoảng 2 phút trên máy phát triển, có thể lâu hơn trên máy chậm hoặc đang bận

# ViT tí hon và CNN nhỏ học phân loại chữ số 8x8 của load_digits.
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.manual_seed(0)
X, y = load_digits(return_X_y=True)
X = torch.tensor(X / 16.0, dtype=torch.float32).view(-1, 1, 8, 8)
y = torch.tensor(y)
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.2, random_state=0)


class TinyViT(nn.Module):
    def __init__(self, P=2, D=32, depth=2, heads=4):
        super().__init__()
        N = (8 // P) ** 2                       # số patch
        # Conv2d kernel = stride = P: cắt patch rồi chiếu tuyến tính
        self.embed = nn.Conv2d(1, D, kernel_size=P, stride=P)
        self.cls = nn.Parameter(torch.zeros(1, 1, D))
        self.pos = nn.Parameter(torch.randn(1, N + 1, D) * 0.02)
        layer = nn.TransformerEncoderLayer(
            D, heads, dim_feedforward=2 * D, dropout=0.0,
            activation="gelu", batch_first=True, norm_first=True)
        self.encoder = nn.TransformerEncoder(
            layer, depth, enable_nested_tensor=False)
        self.head = nn.Sequential(nn.LayerNorm(D), nn.Linear(D, 10))

    def forward(self, x):
        z = self.embed(x).flatten(2).transpose(1, 2)    # (B, N, D)
        z = torch.cat([self.cls.expand(len(z), -1, -1), z], dim=1)
        z = self.encoder(z + self.pos)
        return self.head(z[:, 0])               # chỉ đọc token [CLS]


cnn = nn.Sequential(                            # CNN nhỏ để so sánh
    nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Conv2d(32, 48, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Flatten(), nn.Linear(48 * 2 * 2, 10))


def train_eval(model, epochs=60):
    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
    for _ in range(epochs):
        model.train()
        for idx in torch.randperm(len(Xtr)).split(64):
            loss = nn.functional.cross_entropy(model(Xtr[idx]), ytr[idx])
            opt.zero_grad()
            loss.backward()
            opt.step()
    model.eval()
    with torch.no_grad():
        return (model(Xte).argmax(1) == yte).float().mean().item()


for name, model in [("ViT", TinyViT()), ("CNN", cnn)]:
    n = sum(p.numel() for p in model.parameters())
    print(f"{name}: {n} params, test acc {train_eval(model):.3f}")

Kết quả mẫu

Lấy từ một lần chạy thật. Câu trả lời của mô hình mỗi lần một khác.

ViT: 18218 params, test acc 0.958
CNN: 16122 params, test acc 0.986