Vision Transformer

Patch là một đơn vị đọc ảnh

Patch là khối điểm ảnh không chồng lấn. Với ảnh $H\times W$ có $C$ kênh, patch $P\times P$, nếu hai chiều chia hết cho $P$ thì số patch là $L=(H/P)(W/P)$, mỗi patch có $CP^2$ số. Ta trải mỗi patch thành vector rồi nhân cùng một ma trận $E$ để có vector $D$ chiều. $D$ là chiều biểu diễn được chọn, không phải số pixel.

Với ảnh RGB $224\times224$, patch 16 cho $14\cdot14=196$ token ảnh, mỗi patch thô có $3\cdot16^2=768$ số. Bản ViT-Base chọn $D=768$; việc trùng với 768 số thô chỉ là tình cờ của cấu hình, vì ViT-Large cùng patch 16 dùng $D=1024$. Tên ViT-B/16 đọc là bản Base với patch $16\times16$.

Bài báo ViT (Dosovitskiy và cộng sự, 2020) đưa chuỗi patch embedding vào một Transformer encoder gần như giữ nguyên từ NLP; tên bài báo nói đúng ý đó: một ảnh đáng giá $16\times16$ từ. Các mục sau dựng từng bước trên một ảnh $4\times4$, đủ nhỏ để kiểm bằng tay.

Chia ảnh thành patch bằng tensor

Ta đánh số ảnh xám từ 0 đến 15 theo hàng. Patch trên trái gồm 0, 1, 4, 5; trên phải gồm 2, 3, 6, 7. Nếu flatten toàn ảnh rồi cứ cắt mỗi bốn số, ta nhận 0, 1, 2, 3 là một patch - đó là một hàng, không phải ô vuông.

unfold tạo cửa sổ theo chiều; permute đặt thứ tự batch, hàng patch, cột patch, kênh, hàng trong patch, cột trong patch. reshape chỉ dùng sau khi thứ tự trục đã đúng.

import torch
from torch import nn

torch.set_num_threads(1)
torch.manual_seed(54)
images = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4)
patch_size = 2
windows = images.unfold(2, patch_size, patch_size)
windows = windows.unfold(3, patch_size, patch_size)
patches = windows.permute(0, 2, 3, 1, 4, 5).reshape(1, 4, 4)
print(patches[0].to(torch.int64).numpy())
print(tuple(patches.shape))

Kết quả[[ 0  1  4  5]
 [ 2  3  6  7]
 [ 8  9 12 13]
 [10 11 14 15]]
(1, 4, 4)

Đọc dòng đầu để đối chiếu bốn ô tay. Shape (1,4,4) có hai chiều cùng bằng 4 nhưng ý nghĩa khác: bốn patch và bốn số trong mỗi patch. Khi tăng số kênh, chiều cuối tăng, số patch giữ nguyên.

Patch embedding là phép tuyến tính dùng chung

Với vector patch $\mathbf{p}_i$, ta tính $\mathbf{z}_i=\mathbf{p}_iE+\mathbf{b}$, trong đó $E$ có shape $(CP^2,D)$. Mọi patch dùng chung $E$. Cách khác là conv có kernel $P$, stride $P$, không padding; mỗi kênh đầu ra là một chiều embedding. Hai cách tương đương nếu xếp trọng số cùng thứ tự.

Ta đặt trọng số cụ thể để kết quả dễ kiểm, rồi đối chiếu:

projection = nn.Linear(4, 3)
with torch.no_grad():
    projection.weight.copy_(torch.tensor([[1., 0., 0., 0.],
                                          [0., 1., 0., 0.],
                                          [0.25, 0.25, 0.25, 0.25]]))
    projection.bias.zero_()
tokens = projection(patches)

patch_conv = nn.Conv2d(1, 3, kernel_size=2, stride=2)
with torch.no_grad():
    patch_conv.weight.copy_(projection.weight.reshape(3, 1, 2, 2))
    patch_conv.bias.copy_(projection.bias)
    conv_tokens = patch_conv(images).flatten(2).transpose(1, 2)
print(tokens[0].detach().tolist())
print(torch.allclose(tokens, conv_tokens))

Kết quả[[0.0, 1.0, 2.5], [2.0, 3.0, 4.5], [8.0, 9.0, 10.5], [10.0, 11.0, 12.5]]
True

Token đầu là $(0,1,2.5)$: hai chiều đầu đọc hai ô đầu, chiều thứ ba lấy trung bình $(0+1+4+5)/4$. True xác nhận hai cách xếp ô nhất quán. Trọng số ở đây đặt tay để dễ đọc; trong ViT thật, $E$ được học cùng phần còn lại của mạng.

Attention nối các mảnh bằng nội dung

Trong attention, từ token ta tạo query $Q$, key $K$ và value $V$. Điểm $QK^\top/\sqrt{d_k}$ đo mức phù hợp, softmax theo từng hàng tạo trọng số cộng thành 1, rồi lấy tổng có trọng số của $V$:

$d_k$ là số chiều query/key mỗi head, không phải số patch. Ô $(i,j)$ trong ma trận attention nói token $i$ đọc token $j$ bao nhiêu ở head này. Nó không tự là xác suất vật thể nằm ở patch $j$, cũng không chứng minh vùng ấy gây ra quyết định.

Điểm khác CNN nằm ở chỗ ngay từ lớp đầu, mọi patch đã có đường đọc trực tiếp mọi patch khác: quyển sách và cái đèn ở hai góc bàn trao đổi trong một bước. CNN phải xếp chồng nhiều lớp thì receptive field mới phủ được hai góc (bài 5.1).

Mỗi block của ViT ghép attention và MLP; trước mỗi phần có LayerNorm, sau mỗi phần có residual connection.

Class token (CLS) là một vector học được, ghép vào đầu chuỗi nên chuỗi dài $L+1$. Nó không ứng với vùng ảnh nào, chỉ đọc các patch qua attention, và output của nó ở lớp cuối đi vào head phân loại. Với ViT-B/16 trên ảnh 224, chuỗi vào encoder có shape $(B,197,768)$; mỗi lớp có 12 head, mỗi head một ma trận $197\times197$.

Vì sao cần thông tin vị trí?

Attention dựa vào nội dung không biết token này từ góc trái hay góc phải. Nếu đảo thứ tự token và đảo đầu ra tương ứng, self-attention thông thường không có position vẫn cho cùng quan hệ nội dung. Mean pooling trên đầu ra vì vậy không phân biệt cùng tập mảnh bị đổi chỗ.

Position embedding thêm thông tin vị trí vào token trước encoder. Nếu đổi ảnh nhưng giữ các mã tại vị trí của chúng, model có cơ sở phân biệt bố cục. Đổi cả token và mã đi cùng nó chỉ là đổi cách sắp danh sách, chưa thay quan hệ vị trí được biểu diễn.

ViT gốc dùng bảng vị trí 1D học được: patch đánh số theo hàng từ 0 đến $L-1$, mỗi chỉ số và CLS có một vector riêng, cả bảng có shape $(1,L+1,D)$ và được cộng vào token.

Bảng không ghi sẵn rằng trên lưới $14\times14$, patch 0 và patch 14 kề nhau theo chiều dọc; model tự học điều đó từ dữ liệu. Bài báo cho thấy sau huấn luyện, vector vị trí của các patch cùng hàng hoặc cùng cột giống nhau hơn, và các mã vị trí 2D thiết kế sẵn không cho lợi ích đáng kể.

Ta kiểm tính đối xứng bằng một block nhỏ khởi tạo ngẫu nhiên, vì câu hỏi ở đây là về cấu trúc chứ chưa phải về điều model học được. Dropout tắt để khác biệt không do ngẫu nhiên.

encoder = nn.TransformerEncoderLayer(
    d_model=3, nhead=1, dim_feedforward=6,
    dropout=0.0, batch_first=True,
)
encoder.eval()
permutation = torch.tensor([2, 0, 3, 1])
positions = torch.tensor([[[0.4, 0.0, 0.0],
                           [0.0, 0.4, 0.0],
                           [0.0, 0.0, 0.4],
                           [0.2, -0.3, 0.1]]])
with torch.no_grad():
    pooled = encoder(tokens).mean(dim=1)
    shuffled = encoder(tokens[:, permutation]).mean(dim=1)
    with_position = encoder(tokens + positions).mean(dim=1)
    shuffled_content = encoder(tokens[:, permutation] + positions).mean(dim=1)
print("without position:", torch.allclose(pooled, shuffled, atol=1e-5))
print("with fixed positions:", torch.allclose(with_position, shuffled_content, atol=1e-5))

Kết quảwithout position: True
with fixed positions: False

Dòng True xác nhận đối xứng khi chưa có vị trí: đảo các mảnh thì kết quả sau pooling không đổi, nên sách ở góc trái hay góc phải là như nhau với block này. Dòng False cho thấy cộng mã vị trí cố định đã phá đối xứng đó: cùng tập mảnh ở bố cục khác cho output khác. Dùng thông tin vị trí ấy thế nào là phần model tự học từ dữ liệu.

CNN và ViT: so giả định trước khi so điểm

CNN có thiên hướng cục bộ và chia sẻ trọng số theo vị trí. ViT với attention toàn cục có đường trao đổi trực tiếp giữa patch xa, nhưng mất thông tin bên trong patch nếu projection nén quá mạnh, và phải học quan hệ không gian qua biểu diễn vị trí. “Toàn cục” không đồng nghĩa tự hiểu mọi quan hệ.

Ít giả định có sẵn thì phải học nhiều hơn từ dữ liệu. Bài báo ViT báo cáo rằng khi chỉ huấn luyện trên ImageNet (khoảng 1.3 triệu ảnh) mà không dùng regularization mạnh, ViT kém ResNet cùng cỡ vài điểm phần trăm; khi pretrain trên JFT-300M (303 triệu ảnh), ViT vượt ResNet và tốn ít tài nguyên pretrain hơn, rồi fine-tune tốt sang các tập nhỏ.

Với attention đầy đủ, ma trận điểm có $L^2$ ô cho mỗi head, hoặc $(L+1)^2$ khi có CLS. Giữ patch size, tăng gấp đôi cả hai chiều ảnh làm số patch gấp 4 và số ô attention gấp 16. Convolution có kernel cố định thường tăng công việc theo diện tích, khoảng gấp 4.

Phần MLP và các phép chiếu của ViT chỉ tăng theo $L$, tức gấp 4, nên tổng thời gian tăng ít hơn 16 lần. Thời gian thật còn phụ thuộc cách cài và bộ nhớ activation, vì vậy ước lượng trước rồi đo.

ViT pretrained có thể là baseline mạnh khi có checkpoint và preprocessing phù hợp. Train ViT lớn từ đầu trên 60 ảnh bàn học khó kiểm soát hơn một baseline nhỏ. Patch quá lớn có thể làm chữ rất nhỏ bị gộp; patch quá nhỏ tăng chi phí. Chọn theo đối tượng, dữ liệu và ngân sách, rồi so trên cùng split.

Đổi độ phân giải khi fine-tune

Bài báo ViT fine-tune ở độ phân giải cao hơn lúc pretrain, giữ nguyên patch size, và nội suy 2D bảng vị trí theo vị trí của patch trong ảnh gốc. Các bước:

1. Tách vector CLS.
2. Đưa $L$ vector còn lại về lưới $h\times w\times D$.
3. Nội suy lưới sang kích thước mới.
4. Trải lại rồi ghép CLS vào đầu.

Để thấy hình học được giữ, ta dùng một bảng vị trí tự đặt, trong đó vector của mỗi patch chính là (hàng, cột) của nó, còn CLS là $(-1,-1)$:

import torch.nn.functional as F

def resize_position_table(table, old_grid, new_grid):
    cls_part, grid_part = table[:, :1], table[:, 1:]
    dim = table.shape[-1]
    grid = grid_part.reshape(1, old_grid, old_grid, dim).permute(0, 3, 1, 2)
    grid = F.interpolate(grid, size=(new_grid, new_grid),
                         mode="bilinear", align_corners=True)
    grid_part = grid.permute(0, 2, 3, 1).reshape(1, new_grid * new_grid, dim)
    return torch.cat([cls_part, grid_part], dim=1)

rows, cols = torch.meshgrid(torch.arange(4.0), torch.arange(4.0), indexing="ij")
coords = torch.stack([rows, cols], dim=-1).reshape(1, 16, 2)
table = torch.cat([torch.full((1, 1, 2), -1.0), coords], dim=1)
resized = resize_position_table(table, old_grid=4, new_grid=7)
print(tuple(table.shape), "->", tuple(resized.shape))
print("CLS:", resized[0, 0].tolist())
print("row of first column:", resized[0, 1:, 0].reshape(7, 7)[:, 0].tolist())

Kết quả(1, 17, 2) -> (1, 50, 2)
CLS: [-1.0, -1.0]
row of first column: [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0]

Bảng đi từ 17 lên 50 vector, tức CLS cộng lưới $7\times7$. CLS giữ nguyên $(-1,-1)$, còn tọa độ hàng ở cột đầu chạy đều từ 0 đến 3: mỗi vị trí mới nhận vector của đúng chỗ tương ứng trên lưới cũ. Ví dụ dùng bilinear để số dễ đọc; thư viện thường dùng bicubic. Sau khi nội suy vẫn cần fine-tune ở độ phân giải mới, vì bảng chỉ là điểm xuất phát hợp lý.