A Simple Framework for Contrastive Learning of Visual Representations

Hướng dẫn đọc

Vấn đề và động lực

Việc học biểu diễn thị giác mà không có sự giám sát của con người thường dựa vào mô hình sinh (generative modeling) hoặc các tác vụ phụ giả định theo kinh nghiệm (ví dụ: giải đố ghép hình, dự đoán góc xoay, tái tạo màu sắc). Các phương pháp học tương phản đạt hiệu năng tốt hơn nhưng thường đưa vào các kiến trúc tổng hợp ngữ cảnh chuyên biệt (như PixelCNNs trong CPC), trường tiếp nhận (receptive field) hạn chế, hoặc các bộ nhớ đệm tường minh (như MoCo hoặc InstDisc) để cung cấp các cặp âm. SimCLR giải quyết những phức tạp này bằng cách chứng minh rằng một khuôn khổ đơn giản sử dụng các bộ mã hóa tiêu chuẩn, các cặp âm ngay trong batch (in-batch negatives) và thiết kế tăng cường dữ liệu có hệ thống có thể vượt qua các phương pháp trước đó mà không cần giải pháp can thiệp kiến trúc nào.

Ý tưởng cốt lõi

• Sự kết hợp tăng cường dữ liệu xác định tác vụ dự đoán: Thay vì thay đổi kiến trúc mạng để ép buộc các dự đoán ngữ cảnh không gian, việc kết hợp ngẫu nhiên các phép tăng cường dữ liệu—cụ thể là cắt ngẫu nhiên (random cropping) kết hợp với biến dạng màu sắc (color distortion)—sẽ ngăn ngừa các lối tắt (shortcuts, chẳng hạn như khớp lược đồ màu) và buộc mạng phải trích xuất các đặc trưng ngữ nghĩa.
• Tăng cường dữ liệu mạnh hơn có lợi cho học tương phản: Các phép biến đổi như làm nhiễu màu mạnh (color jittering) vốn làm giảm hiệu năng huấn luyện có giám sát lại cải thiện đáng kể cho học tương phản tự giám sát.
• Projection head phi tuyến: Việc bổ sung một mạng MLP nhỏ $g(\cdot )$ nằm giữa biểu diễn của bộ mã hóa $\mathbf{h}$ và hàm mất mát tương phản giúp cải thiện chất lượng biểu diễn. Vì $\mathbf{z}=g(\mathbf{h})$ được huấn luyện để bất biến với các phép biến đổi, việc loại bỏ $g(\cdot )$ sau khi huấn luyện sẽ giữ lại thông tin nhạy cảm với phép biến đổi trong $\mathbf{h}$, mang lại lợi ích cho các tác vụ xuôi dòng (downstream tasks).
• Mở rộng cặp âm thông qua batch lớn và thời gian huấn luyện dài: Thay vì sử dụng bộ nhớ đệm tường minh, SimCLR dựa vào kích thước batch lớn (lên tới 8192) được tối ưu hóa bằng LARS, cung cấp hàng nghìn cặp âm trong batch cho mỗi cặp dương.
• Chuẩn hóa theo Batch Toàn cục (Global Batch Normalization): Để tránh lối tắt mà mạng có thể khai thác thống kê chuẩn hóa theo batch nội bộ trên từng thiết bị nhằm nhận diện các cặp dương, giá trị trung bình và phương sai của batch normalization được đồng bộ hóa trên toàn bộ các bộ tăng tốc phân tán.

Phương pháp

Với một minibatch gồm $N$ ảnh được lấy mẫu, mỗi mẫu ${\mathbf{x}}_k$ được biến đổi bởi hai phép tăng cường ngẫu nhiên khác biệt $t,t^{\prime }\sim \mathcal{T}$, tạo ra một cặp dương $({\tilde{\mathbf{x}}}_i,{\tilde{\mathbf{x}}}_j)$ trong số $2N$ mẫu đã tăng cường.

1. Bộ mã hóa cơ sở (Base Encoder): ${\mathbf{h}}_i=f({\tilde{\mathbf{x}}}_i)=\mathrm{ResNet}({\tilde{\mathbf{x}}}_i)$, trích xuất các biểu diễn từ tầng average pooling.
2. Projection Head: ${\mathbf{z}}_i=g({\mathbf{h}}_i)=W^{(2)}\sigma (W^{(1)}{\mathbf{h}}_i)$, ánh xạ các biểu diễn vào một không gian ẩn 128 chiều với một tầng ẩn duy nhất cùng hàm kích hoạt ReLU $\sigma$.
3. Hàm mất mát NT-Xent: Độ tương đồng cosine giữa các vector chuẩn hóa được định nghĩa là $\mathrm{sim}(\mathbf{u},\mathbf{v})={\mathbf{u}}^{\top }\mathbf{v}/(\Vert \mathbf{u}\Vert \Vert \mathbf{v}\Vert )$. Đối với một cặp dương $(i,j)$, hàm mất mát là: $${\ell }_{i,j}=-\log \frac{\exp (\mathrm{sim}({\mathbf{z}}_i,{\mathbf{z}}_j)/\tau )}{{\sum }_{k=1}^{2N}\text{mathbbm}{1}_{[k\ne i]}\exp (\mathrm{sim}({\mathbf{z}}_i,{\mathbf{z}}_k)/\tau )}$$ trong đó $\tau$ là siêu tham số nhiệt độ (temperature) và $\text{mathbbm}{1}_{[k\ne i]}$ có giá trị bằng $1$ khi và chỉ khi $k\ne i$. Hàm mất mát tổng $\mathcal{L}$ tính trung bình trên toàn bộ $2N$ cặp trong batch.

Sau khi huấn luyện, $g(\cdot )$ bị loại bỏ và $\mathbf{h}$ được sử dụng cho bài toán phân loại tuyến tính hoặc tinh chỉnh xuôi dòng.

Kết quả

• Đánh giá Tuyến tính trên ImageNet: SimCLR với ResNet-50 tiêu chuẩn đạt độ chính xác top-1 69,3% (ngang bằng hoặc vượt qua các mô hình tự giám sát trước đó). Mở rộng lên ResNet-50 ($4\times$) đạt độ chính xác top-1 76,5%, ngang ngửa với mô hình chuẩn ResNet-50 được giám sát hoàn toàn.
• Tinh chỉnh Bán giám sát: Khi tinh chỉnh chỉ trên 1% nhãn của ImageNet, ResNet-50 ($4\times$) đạt độ chính xác top-5 85,8% (và 63,0% top-1), vượt qua AlexNet được huấn luyện trên 100% nhãn.
• Học chuyển giao (Transfer Learning): Trên 12 bộ dữ liệu phân loại ảnh tự nhiên, SimCLR ResNet-50 ($4\times$) sau khi tinh chỉnh đạt hiệu năng tương đương hoặc vượt qua baseline có giám sát trên 10/12 bộ dữ liệu.
• Phân tích thành phần (Ablation) trên Projection Head: Sử dụng projection head phi tuyến vượt trội hơn projection tuyến tính +3% và vượt hơn việc không dùng projection >10% trong đánh giá tuyến tính trên ImageNet.

Hạn chế và các câu hỏi mở

• Đòi hỏi năng lực tính toán lớn: Quá trình huấn luyện đòi hỏi kích thước mini-batch rất lớn (ví dụ: 4096 hoặc 8192) cùng lịch trình huấn luyện kéo dài (lên tới 1000 epoch) để đạt kết quả tốt nhất, cần đến các cụm tăng tốc tính toán lớn (32–128 lõi TPU).
• Chênh lệch kiến trúc khi chuyển giao: Mặc dù ResNet-50 ($4\times$) bắt kịp hoặc vượt qua các baseline có giám sát trên các benchmark chuyển giao, ResNet-50 tiêu chuẩn vẫn xếp sau các baseline có giám sát trên mọi tác vụ đánh giá tuyến tính.
• Sự đánh đổi thông tin: Cơ chế mà projection head dùng để giữ lại thông tin đặc thù theo phép biến đổi trong $\mathbf{h}$ trong khi vẫn ép buộc tính bất biến ở $\mathbf{z}$ chủ yếu dựa trên các quan sát thực nghiệm thay vì một khuôn khổ phân tích toán học rõ ràng.