개요
Vision Transformer의 문제점 중 하나인 이미지 크기에 따른 메모리 및 연산량 문제를 기존의 채널 단위의 Self-Attention을 수행함으로써 문제를 해결하는 XCiT 모듈을 제안하는 논문이다.
Motivation
Transformer는 고질적으로 Self-Attention을 하는 과정에서 메모리와 연산시간이 많이 요구되는 문제를 가지고 있다. Vision Transformer(ViT)의 측면에서 본다면, 이미지 크기가 커질수록 토큰의 수가 많아지게 되면서 메모리와 연산시간이 quadratic 하게 증가하기 때문에 고해상도 이미지에는 적용이 어렵다는 문제를 가지고 있다.
저자들은 이미지 크기에 따른 메모리 및 연산시간의 증가를 다루기 위해 Self-Attention을 토큰이 아닌 채널에 대해 적용하는 방법을 제안한다. 이 방법은 결과적으로 메모리와 연산시간이 이미지의 크기에 선형적인 영향만을 받도록 만들어준다.
Contribution
- 기존 토큰을 대상으로하는 Self-Attention 대신 채널 사이의 Attention을 계산하는 Cross-Covariance Attention(XCA)를 적용. 이를 통해 고해상도 이미지에 대해서도 효율적인 처리가 가능해짐
- 고정된 길이의 채널에 대해 Attention을 수행하게 되면서 이미지 크기와의 의존성을 줄임. 이미지 크기에 따른 성능 변화에 대해 더 robust 한 결과를 가질 수 있게됨
- image classification 에서 state-of-the-art 모델과 비슷한 성능을 보여줌
- object detection 에서 비슷한 크기의 모델과 비교했을 때 더 좋은 성능을 보여줌
XCiT
Background
Token Self-Attention (기존 방식)
⇒ (N x d) x (d x N) ⇒ N x N ⇒
⇒ 각 토큰 사이의 covariance를 계산
⇒ (N x N) x (N x d) ⇒ N x d ⇒
⇒ 토큰의 개수 N에 대해 quadratic 증가
XCiT
XCiT Layer는 총 3개의 모듈로 구성되어있으며, 각각 서로의 부족한 점을 보완해주는 역할을 한다.
Cross-Covariance Attention (XCA)
기존 Self-Attention을 대체하는, 토큰 사이의 관계 대신 채널 사이의 관계성을 계산하는 레이어
저자들의 가정은 특정 채널이 특정한 특징을 표현하는 역할을 수행하도록 학습하는것
(예: 어떤 채널은 눈을 표현하고, 어떤 채널은 입을 표현하는 것)
XCA의 방식을 요약하면 행렬 연산의 순서만 바꿔서 Attention의 대상을 변경한거라고 보면 된다.
⇒ (d x N) x (N x d) ⇒ d x d ⇒
⇒ 각 채널 사이의 covariance를 계산
⇒ (N x d) x (d x d) ⇒ N x d ⇒
⇒ 토큰의 개수 N에 대해 선형적으로 증가
Local Patch Interaction (LPI)
토큰 사이의 covariance를 계산하는 방식에서 채널 사이의 covariance를 계산하게 되면서 XCA에는 patch 사이의 명시적인 상관관계에 대한 정보가 부족하게 되었다. 이러한 정보를 포함시키기 위해 간단한 depth-wise convolution 레이어를 통과시킨다.
결과적으로 XCA + LPI 모듈이 기존 Self-Attention 모듈을 대체한다.
Feed-Forward Network (FFN)
feature 사이의 interaction을 표현할 수 있도록 만들어주는 역할을 수행한다.
Handling images of varying resolution
XCA 모듈은 covariance를 계산하는 과정에서 이미지의 해상도와 의존관계가 형성되지 않는다. 다시말하면, 어떠한 크기의 입력이 들어오더라도 동일한 크기의 출력을 낼 수 있게 된다. 이러한 특징은 이미지 해상도의 변화에 대해 robust 한 성능을 보여줄 수 있도록 만들어준다.
Experiment
Image Classification
Result on ImageNet
- 비슷한 수준의 파라미터 수와 FLOPs 를 가지는 모델과 비교했을 때 더 좋은 성능을 보여줌
- 16 x 16 패치 사이즈에 384 x 384 해상도의 이미지를 사용하는 경우, 다른 state-of-the-art 모델들과 비교될만한 경쟁력있는 결과를 보여줌
- 384 x 384 이미지에 8 x 8 패치 사이즈를 적용하는 경우 FLOPs 가 크게 늘지만 성능이 증가함
결국에는 이미지 해상도가 크고, 작은 패치사이즈를 사용할수록 결과가 좋고 대신 FLOPs 가 크게 늘어난다는 내용이다.
XCiT는 더 큰 이미지, 더 작은 패치에 대한 비용의 증가가 상대적으로 적기 때문에 이러한 방식을 통해 성능을 향상시키기 더 수월하다.
Robustness to resolution changes
- ViT에서 공통적으로 이미지 크기의 변화에 따라 어느 정도까지는 성능이 좋아지다가 감소하는 현상이 나타남
- XCiT의 경우 성능의 감소의 기울기가 더 작다.
⇒ XCiT에서 토큰 대신 채널을 활용하기 때문에 이미지 크기에 의존성이 없어 조금 더 robust한 결과를 가짐
Class Attention Visualization
- 각 head 마다 서로 다른 부분에 집중하는 결과를 확인할 수 있음
- 우연의 일치인지는 모르겠지만 사람의 얼굴을 집중하는 head 에서 새의 얼굴에 집중하는 현상이 나타남.
Analysis and Ablations
- Patch Projection 을 할때 Linear Projection 보다 Convolutional Projection이 더 좋은 성능을 보여줌
- LPI Layer의 유무에 따라 작은 성능의 차이가 발생
- XCA Layer의 유무는 성능에 크게 영향을 미친다고 주장
⇒ Attention을 안하기 때문에 당연한 결과로 보임
- normalization 의 경우 없으면 학습이 안됨
Object Detection and Instance Segmentation
위 결과를 근거로 dense prediction에 XCiT가 좋은 성능을 이끌어낼 수 있다고 주장
- 이미지 해상도에 대해 메모리 사용량이 선형적으로 증가
- 이미지 해상도가 커져도 성능이 어느정도 Robust 하게 유지
COCO dataset 에 대해 평가를 진행했을 때 좋은 성능을 보여주는 것을 확인
- ConvNet 기반 backbone에 비해서는 확실히 좋은 성능올 보여줌
- 비슷한 파라미터 수를 가지는 ViT 모델과 비교했을 때 가장 성능이 좋은 모델과 비슷하거나 더 좋은 성능을 보여줌 (Swin Transformer 와 경쟁이 가능한 성능이다 라는 부분을 특히 강조)
Semantic Segmentation
- ConvNet 기반 backbone에 비해서는 확실히 좋은 성능을 보여줌
- 작은 patch를 사용하는 경우 비슷한 크기의 Swin Transformer와 비교하여 더 좋은 성능을 보여줌
Share article