Masked Autoencoders (MAE)

Autoencoder có mặt nạ (MAE) là một phương pháp học tự giám sát cho hình ảnh: model che ngẫu nhiên một phần lớn các patch của ảnh, thường là 75%, rồi học cách khôi phục các pixel bị thiếu từ phần còn lại nhìn thấy được. Kaiming He và các đồng nghiệp tại Meta AI giới thiệu phương pháp này trong bài báo năm 2021 Masked Autoencoders Are Scalable Vision Learners, điều chỉnh phương pháp dự đoán từ bị che của BERT cho vision transformer. Vì encoder chỉ xử lý các patch nhìn thấy được, quá trình tiền huấn luyện MAE nhanh hơn ít nhất 3 lần so với xử lý mọi patch; model ViT-Huge được tiền huấn luyện theo cách này đạt độ chính xác 87.8% trên ImageNet, kết quả tốt nhất vào thời điểm đó trong số các phương pháp chỉ sử dụng dữ liệu ImageNet-1K. Sau đó, encoder đã tiền huấn luyện được dùng làm backbone cho phân loại, phát hiện và phân đoạn.

Cách hoạt động của Autoencoder có che giấu#

MAE chia ảnh thành một lưới các patch không chồng lấn — 196 patch có kích thước 16x16 pixel với đầu vào 224x224 — rồi che ngẫu nhiên phần lớn trong số đó. Encoder và decoder được thiết kế không cân xứng có chủ đích. Encoder, một ViT tiêu chuẩn, chỉ chạy trên các patch nhìn thấy được, giúp giảm lượng tính toán và bộ nhớ xuống còn khoảng một phần tư. Decoder gọn nhẹ, mặc định gồm 8 khối transformer với lượng tính toán trên mỗi token dưới 10% so với encoder, nhận các patch nhìn thấy được đã mã hóa cùng một token mặt nạ có thể học dùng chung cho mỗi vị trí bị che, rồi dự đoán giá trị pixel của từng patch bị che. Hàm mất mát là sai số bình phương trung bình chỉ tính trên các patch bị che; việc dùng pixel đã chuẩn hóa theo từng patch làm mục tiêu còn cải thiện thêm các đặc trưng đã học. Sau khi tiền huấn luyện, decoder bị loại bỏ:

Tỷ lệ che cao là yếu tố khiến nhiệm vụ này hữu ích. Các patch ảnh lân cận có tính dư thừa cao, vì vậy nếu chỉ che một vài patch, model có thể lấp từng khoảng trống bằng cách nội suy từ các patch bên cạnh. Khi 75% bị che, model phải suy luận toàn bộ đối tượng và cảnh vật từ các mảnh rời rạc, buộc model học hình dạng, bộ phận và ngữ cảnh. Ngược lại, BERT chỉ che 15% số từ vì mỗi từ đã mang ý nghĩa cô đọng.

Trọng tâm của MAE là bước tạo mặt nạ, được minh họa ở đây bằng PyTorch trên một batch gồm bốn ảnh:

import torch

torch.manual_seed(0)
images = torch.rand(4, 3, 224, 224)  # a batch of 4 RGB images

# Patchify: each 224x224 image becomes 196 patches of 16x16 pixels, flattened to 768 values
patches = images.unfold(2, 16, 16).unfold(3, 16, 16).permute(0, 2, 3, 1, 4, 5).reshape(4, 196, 768)

# Random masking as in the MAE reference code: rank patches by random noise and keep the first 25%
keep = int(196 * (1 - 0.75))
ids_keep = torch.argsort(torch.rand(4, 196), dim=1)[:, :keep]
visible = torch.gather(patches, 1, ids_keep.unsqueeze(-1).expand(-1, -1, 768))

# Mask of hidden patches (1 = hidden): the decoder predicts these, and the loss is averaged over them only
mask = torch.ones(4, 196).scatter(1, ids_keep, 0)
print(f"Encoder input: {visible.shape[1]} of {patches.shape[1]} patches per image")
print(f"Reconstruction targets: {int(mask[0].sum())} hidden patches per image")

Kết quả in ra là Encoder input: 49 of 196 patches per image và Reconstruction targets: 147 hidden patches per image. Do đó, các lớp transformer của encoder xử lý một chuỗi có độ dài bằng một phần tư ảnh đầy đủ; và vì attention so sánh từng token với mọi token khác, mỗi ma trận điểm attention có kích thước bằng một phần mười sáu so với khi xử lý toàn bộ 196 patch.

Kết quả và các biến thể#

Đặc trưng MAE phát huy tốt nhất sau khi fine-tune. Với linear probing, encoder ViT-L được tiền huấn luyện bằng MAE đạt độ chính xác 73.5% trên ImageNet, thấp hơn các phương pháp tương phản; nhưng khi fine-tune và chuyển giao sang phát hiện đối tượng và phân đoạn trên COCO, tiền huấn luyện MAE cho kết quả tốt hơn tiền huấn luyện có giám sát. Quy trình này đã được điều chỉnh vượt xa phạm vi ảnh tĩnh: VideoMAE che ít nhất 90% các patch không-thời gian của một đoạn video, SatMAE tiền huấn luyện trên ảnh vệ tinh, còn ConvNeXt V2 đưa phương pháp mã hóa tự động có mặt nạ vào mạng tích chập bằng các phép tích chập thưa.

Tiền huấn luyện có mặt nạ cũng là nền tảng của nhiều foundation model được sử dụng rộng rãi. Segment Anything Model (SAM) của Meta, có trong gói Ultralytics, xây dựng image encoder dựa trên ViT được tiền huấn luyện bằng MAE; phiên bản ViT-B của model này, sam_b.pt, là model mặt nạ mặc định trong Ultralytics auto_annotate, công cụ biến các bounding box của bộ phát hiện YOLO26 thành nhãn phân đoạn.

Autoencoder có mặt nạ so với các thuật ngữ liên quan#

• Autoencoder: Nén toàn bộ đầu vào thành một mã nhỏ rồi tái tạo đầu vào. MAE nhận đầu vào không đầy đủ và phải dự đoán những phần chưa từng thấy; đầu ra encoder của MAE giữ lại một vector cho mỗi patch nhìn thấy được.
• Mô hình hóa ảnh có mặt nạ: Họ phương pháp rộng hơn mà MAE thuộc về. BEiT dự đoán các token thị giác rời rạc cho patch bị che, còn SimMIM đưa token mặt nạ qua encoder; trong khi đó, MAE dự đoán pixel và không đưa token mặt nạ vào encoder.
• Học tương phản: Học bằng cách khớp các góc nhìn đã tăng cường của cùng một ảnh và không cần decoder; phương pháp này tạo ra các đặc trưng hoạt động tốt hơn khi được giữ cố định, trong khi đặc trưng MAE cải thiện nhiều hơn sau khi fine-tune.
• Kiến trúc dự đoán nhúng chung (JEPA): Cũng che các vùng ảnh, nhưng dự đoán embedding của chúng thay vì pixel, giúp model không tiêu tốn năng lực vào chi tiết kết cấu.