CoAtNet - Kết hợp Convolution và Attention cho tất cả các kích thước dữ liệu

1. Giới thiệu

Convolutional Neural Network (CNN) là một trong những mô hình học sâu phổ biến nhất hiện nay. CNN thường được dùng trong các bài toán thuộc lĩnh vực thị giác máy tính như: phân loại hình ảnh, phát hiện đối tượng,...

Transformer đạt được nhiều thành công lớn trong lĩnh vực xử lý ngôn ngữ tự nhiên (NLP) với các tác vụ: phân loại văn bản, dịch máy,... Gần đây, nhiều nghiên cứu cố gắng đưa kiến trúc này vào lĩnh vực thị giác máy tính, tiêu biểu nhất là mô hình Vision Transformer (ViT). Tuy nhiên, các mô hình dựa trên Transformer này vẫn chưa hoàn toàn vượt qua được các kiến trúc CNN tốt nhất. Một trong những hạn chế chính của mô hình Vision Transformer so với CNN là thiếu tính chất translation equivariance.

Do đó, khi một hình ảnh bị dịch chuyển một ít, mô hình sẽ gặp khó khăn khi dự đoán. Vì vậy, để vượt qua CNN, ViT phải được huấn luyện trên bộ dữ liệu khổng lồ (JFT-300M).

Mặc dù thế, Self Attention trong Transformer vẫn có một số ưu điểm. Dưới đây là một vài tính chất quan trọng của CNN và Self Attention:

• Translation Equivariance: Tính chất của CNN. Sự dịch chuyển của một vật thể trên ảnh đầu vào so với đầu ra là bất biến trên một mạng CNN
• Input-adaptive Weighting: Tính chất của Self-Attention, trọng số chú ý phụ thuộc vào đầu vào. Điều này giúp mô hình hiểu hơn về các mối quan hệ phức tạp của đầu vào
• Global Receptive Field: Self-Attention có trường tiếp nhận toàn cục

CoAtNet ra đời nhằm kết hợp ưu điểm của cả hai kiến trúc trên. Trong bài viết này, mình sẽ giới thiệu về kiến trúc, các thử nghiệm, kết quả của mô hình CoAtNet.

2. Mô hình

- Kết hợp Convolution và Attention

Trong mô hình CoAtNet, Depth wise Convolution được các tác giả lựa chọn để kết hợp với Self-Attention

Công thức của Convolution và Self-Attention:

• Convolution:

\[y_{i}=\sum_{j \in \mathcal{L}(i)}w_{i-j} \odot x_{j}\]

• Self-Attention:

\[y_{i}=\sum_{j \in \mathcal{G}} \underbrace{\frac{\exp \left(x_{i}^{T} x_{j}\right)}{\sum_{j \in \mathcal{G}} \exp \left(x_{i}^{T} x_{k}\right)}}_{A_{i, j}} x_{j}\]

Kết hợp Convolution và Self-Attention. Công thức:

\[y_{i}^{p r e}=\sum_{j \in G}\left(\frac{\exp \left(x_{i}^{T} x_{j}+w_{i} x_{j}\right)}{\sum_{j \in \mathcal{G}} \exp \left(x_{i}^{T} x_{k}+w_{i-k}\right)}\right) x_{j}\]

Trong đó: \(w\) là global static convolution kernel (hoặc relative position bias)

- Biểu diễn thông tin vị trí (\(w_{i-j}\))

+ Relative position representations 1D:

Khác với position encoding trong Transformer gốc. Trong relative position representations, khi muốn tính trọng số chú ý cho một phần tử so với một phần tử khác, ta phải tính khoảng cách (số phần tử) giữa hai phần tử này. Sau đó, chúng ta phải tra cứu trong một bảng (được gọi là relative position table). Bảng này chứa trọng số cho mỗi khoảng cách. Nếu chúng ta có N phần tử, bảng này có (2 * N) - 1 trọng số. Ngoài ra, giữa các phần tử có khoảng cách giống nhau sử dụng cùng một trọng số trong bảng

+ Relative position representations 2D:

Chúng ta xem xét hình ảnh có kích thước 2 x 2 này. Tương tự relative position representations 1D, ta có 9 trọng số tương ứng với khoảng cách như hình bên dưới.

Cách tính toán relative position bias trong Self-Attention:

Phép tính này tương đương với phép tích chập:

Trong mô hình CoAtNet, tác giả sử dụng các khối MBConv làm ConvNet trích xuất đặc trưng ảnh đầu vào giúp giảm khối lượng tính toán cho Attention layer

Các kiến trúc thử nghiệm này được đánh giá trên hai tiêu chí:

- Tính tổng quát hóa: Độ chênh lệch độ chính xác trên tập huấn luyện và tập kiểm thử trên tập dữ liệu nhỏ. Kết quả cho thấy, mô hình có càng nhiều lớp Convolution càng tốt, mô hình chỉ có các lớp Attention (ViT) có kết quả kém nhất

\[\text{C-C-C-C}\approx\text{C-C-C-T}\geq\text{C-C-T-T}>\text{C-T-T-T}\gg\text{ViT}_{\text{REL}} .\]

- Hiệu suất: Độ chính xác trên tập kiểm thử tốt hơn trên bộ dữ liệu lớn với cùng số epoch.

\[\text{C-C-T-T}\approx\text{C-T-T-T}>\text{ViT}_{\text{REL}} .>\text{C-C-C-T}>\text{C-C-C-C}\]

Thử nghiệm bổ sung: CCTT và CTTT được huấn luyện trước trên JFT và được fine tune trên ImageNet-1K trong 30 epoch. Kiến trúc CCTT đạt kết quả tốt hơn. Do đó, tác giả chọn kiến trúc này làm kiến trúc tổng thể cho mô hình CoAtNet.

3. Kết quả

Dưới đây là kết quả các kiến trúc của mô hình CoAtNet:

Nhìn vào kết quả trên, ta có một số đánh giá sau:

- Khi CoAtNet chỉ được đào tạo trước trên ImageNet-1K, CoAtNet (CoAtNet-3) có kết quả tương đương với EfficientNet-V2 và NFNets và vượt trội hơn so với ViT

- Khi sử dụng ImageNet-21K để huấn luyện trước, CoAtNet (CoAtNet-4) đạt độ chính xác 88,56% trên ImageNet-1K vượt qua độ chính xác ViT với kích thước lữ liệu ít hơn 23 lần

- Khi sử dụng JFT-300M, CoAtNet (CoAtNet-5) đạt kết quả 89,77% vượt qua các mô hình khác khi sử dụng bộ dữ liệu này

- Khi sử dụng JFT-3B, CoAtNet-6 và CoAtNet-7 đạt độ chính xác lần lượt là 90,88% và 90,45% trên ImageNet-1K đạt top-3 và top-1 trên bảng xếp bạng ImageNet trên Paper With Code

4. Thử nghiệm

Trong phần này, mình sẽ trình bày kết quả của mô hình CoAtNet trong bài phân loại hình ảnh.

Dưới đây là thông tin thử nghiệm mô hình:

- Bộ dữ liệu: CIFAR10 có 10 lớp

- Tập huấn luyện: 50000 ảnh

- Tập kiểm thử: 10000 ảnh

- Kích thước hình ảnh đầu vào: (32, 32, 3); (160, 160, 1)

- Kiến trúc mô hình thử nghiệm: CoAtNet-0

Kết quả:

5. Kết luận

Trong bài viết này, mình đã giới thiệu về mô hình CoAtNet. CoAtNet đã phá vỡ các quan niệm về việc các mô hình dựa trên Transformer cần lượng dữ liệu khổng lồ để đạt được SOTA. Điều này có ý nghĩa quan trọng trong các dự án thực tế do lượng dữ liệu có nhãn thu thập được để huấn luyện không phải lúc nào cũng dễ dàng. Ngoài ra, mô hình này đã chứng minh được tính tổng quát hóa của nó tốt hơn so với các mô hình dựa trên Transformer khác và hiệu suất tốt hơn so với các mô hình CNN thông qua các thử nghiệm. Trong thời gian tới, team mình sẽ tiếp tục nghiên cứu và thử nghiệm các kiến trúc lớn hơn của mô hình này để so sánh, cải thiện kết quả các service phân loại hình ảnh.

6. Tài liệu tham khảo

- Peter Shaw, Jakob Uszkoreit, Ashish Vaswani: Self-Attention with Relative Position Representations.

- Zihang Dai, Hanxiao Liu, Quoc V. Le, Mingxing Tan: CoAtNet: Marrying Convolution and Attention for All Data Sizes.

- https://paperswithcode.com/