Swin Transformer

Swin Transformer, viết tắt của Shifted Window Transformer, là một vision transformer phân cấp, tính self-attention bên trong các cửa sổ cục bộ nhỏ và dịch chuyển những cửa sổ này giữa các layer liên tiếp. Ze Liu và cộng sự tại Microsoft Research Asia giới thiệu kiến trúc này trong bài báo năm 2021 Swin Transformer: Vision Transformer phân cấp sử dụng cửa sổ dịch chuyển, bài báo giành giải Marr cho bài xuất sắc nhất tại ICCV 2021. Các cửa sổ cục bộ khiến chi phí attention tăng tuyến tính theo số pixel thay vì bậc hai, còn việc gộp patch theo từng giai đoạn tạo ra feature map đa tỉ lệ như mạng neural tích chập. Những đặc điểm này giúp transformer trở thành backbone đa dụng thiết thực cho phát hiện đối tượng và phân đoạn trên ảnh độ phân giải cao, vượt ra ngoài tác vụ phân loại.

Cách hoạt động của cửa sổ dịch chuyển#

Đầu tiên, Swin chia ảnh thành các patch 4×4 pixel và chiếu từng patch thành embedding token, vì vậy ảnh 224×224 trở thành lưới token 56×56 với 96 channel trong model Tiny. Attention toàn cục trên cả 3,136 token sẽ so sánh mọi token với tất cả token còn lại; Swin chia lưới thành các cửa sổ không chồng lấp gồm 7×7 token và chỉ tính attention trong từng cửa sổ. Bài báo cho biết chi phí của attention toàn cục là 4hwC² + 2(hw)²C trên lưới h×w với C channel, trong khi các cửa sổ M×M token có chi phí 4hwC² + 2M²hwC, tuyến tính theo hw.

Nếu chỉ dùng cửa sổ, thông tin sẽ không bao giờ được trao đổi giữa chúng, nên cứ mỗi block thứ hai, lưới cửa sổ được dịch chuyển nửa chiều rộng cửa sổ, tức 3 token, khiến các cửa sổ mới băng qua ranh giới cũ. Để giữ số cửa sổ không đổi, quá trình triển khai cuộn feature map theo chu kỳ và áp dụng mask attention giữa các token vốn không nằm cạnh nhau trong ảnh gốc. Mỗi block cũng cộng thêm bias vị trí tương đối được học vào các điểm attention.

Giữa các giai đoạn, một layer gộp patch nối từng nhóm 2×2 token lân cận rồi chiếu tuyến tính chúng, giảm một nửa độ phân giải và tăng gấp đôi số channel. Bốn giai đoạn xuất feature map ở mức 1/4, 1/8, 1/16 và 1/32 độ phân giải đầu vào, cùng các stride như ResNet, nên Swin có thể tích hợp vào các head phát hiện và phân đoạn hiện có như mạng kim tự tháp đặc trưng. Bản triển khai trong torchvision minh họa cả hai cơ chế:

import torch
from torchvision.models import swin_t

# Swin-T architecture; random weights are enough to inspect its structure
model = swin_t().eval()

# Each pair of blocks uses 7x7 windows, the second shifted by 3 patches
regular, shifted = model.features[1]
print(f"Window {regular.attn.window_size}, shift {regular.attn.shift_size} then {shifted.attn.shift_size}")

# features alternates stages with patch merging; stage outputs are (batch, height, width, channels)
x = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    for i, layer in enumerate(model.features):
        x = layer(x)
        if i % 2 == 1:
            print(f"Stage {i // 2 + 1}: {tuple(x.shape)}")

Đoạn mã in ra Window [7, 7], shift [0, 0] then [3, 3], tiếp theo là đầu ra các giai đoạn có kích thước (1, 56, 56, 96), (1, 28, 28, 192), (1, 14, 14, 384) và (1, 7, 7, 768): mỗi giai đoạn giảm một nửa kích thước lưới và tăng gấp đôi số channel, tạo ra kim tự tháp mà các head dự đoán dày đặc cần đến.

Swin Transformer so với ViT và CNN#

ViT nguyên bản chia ảnh thành các patch 16×16, duy trì một độ phân giải duy nhất qua mọi layer và dùng attention toàn cục; cách này phù hợp với phân loại ảnh nhưng tốn kém với đầu vào độ phân giải cao trong phát hiện và phân đoạn. Swin đưa trở lại hai prior của mạng tích chập là tính cục bộ và hệ phân cấp đặc trưng, đồng thời giữ trọng số phụ thuộc đầu vào của cơ chế attention. Cửa sổ cục bộ của Swin khác với sliding window attention, trong đó mỗi token attention đến vùng lân cận riêng, vì cửa sổ của Swin là các phân vùng cố định được nhiều token dùng chung, cho phép ánh xạ hiệu quả lên phép nhân ma trận của GPU.

Khi ra mắt, model lớn nhất, Swin-L được tiền huấn luyện trên ImageNet-22K, đạt độ chính xác top-1 87.3% trên ImageNet-1K, box AP 58.7 và mask AP 51.1 trên COCO test-dev, cùng mIoU 53.5 trên ADE20K, những kết quả được công bố tốt nhất trên các benchmark phát hiện và phân đoạn ngữ nghĩa đó. Thành công của model cũng thúc đẩy nghiên cứu mạng tích chập: ConvNeXt của Meta (2022) hiện đại hóa ResNet bằng các lựa chọn thiết kế lấy từ Swin, đạt độ chính xác ngang bằng hoặc cao hơn.

Swin có thể làm backbone cho model Ultralytics tùy chỉnh: module TorchVision trong file YAML của model tải swin_t của torchvision cùng trọng số ImageNet đã tiền huấn luyện, còn head phát hiện YOLO được huấn luyện trên feature map stride-32 duy nhất mà module trả về.

Các biến thể và ứng dụng#

• Swin Transformer V2: Mở rộng thiết kế lên 3 tỷ tham số và ảnh có kích thước đến 1,536×1,536 pixel vào năm 2022, sử dụng cosine attention có scale và bias vị trí liên tục theo thang log để duy trì độ ổn định khi huấn luyện.
• Nhận diện hành động trong video: Video Swin Transformer mở rộng cửa sổ sang ba chiều, bao gồm không gian và thời gian, để phân loại hành động trong các đoạn video.
• Siêu phân giải và khôi phục ảnh: SwinIR sử dụng các block Swin để phóng to ảnh, khử nhiễu và loại bỏ artifact JPEG.
• Phân tích ảnh y khoa: Swin-Unet và Swin UNETR kết hợp encoder Swin với decoder theo kiểu U-Net; Swin UNETR phân đoạn khối u não trong thể tích MRI 3D.