[논문 리뷰] Replacing softmax with ReLU in Vision Transformers

Mitchell Wortsman, Jaehoon Lee외 2명|arXiv (Cornell University)|2023. 09. 15.

Advanced Neural Network Applications인용 수 8

한 줄 요약

본 논문은 비전 트랜스포머 주의에서 소프트맥스를 ReLU로 대체하는 것을 조사하며, 시퀀스 길이 확장과 함께 ReLU 주의가 ImageNet-21k에서 소프트맥스 주의의 스케일링 성능에 도달할 수 있음을 보이고, 병렬화를 개선한다.

ABSTRACT

Previous research observed accuracy degradation when replacing the attention softmax with a point-wise activation such as ReLU. In the context of vision transformers, we find that this degradation is mitigated when dividing by sequence length. Our experiments training small to large vision transformers on ImageNet-21k indicate that ReLU-attention can approach or match the performance of softmax-attention in terms of scaling behavior as a function of compute.

연구 동기 및 목표

• 비전 트랜스포머에서 소프트맥스 주의의 점별 대안을 탐구한다.
• 시퀀스 길이로 스케일링된 ReLU 기반 주의가 소프트맥스 성능에 도달할 수 있는지 평가한다.
• 정확도와 학습 효율성에 대한 시퀀스 길이 스케일링과 qk-layernorm의 영향을 평가한다.
• ReLU 주의에서 게이팅(gating) 및 대체 활성화의 영향을 분석한다.
• 대규모에서 비소프트맥스 주의를 사용하는 비전 트랜스포머 학습에 관한 실용적 지침을 제공한다.

제안 방법

• phi = L^{-alpha} h 로 주의를 수식화하고, alpha가 대략 1인 상황에서 ReLU에 초점을 둔다.
• ReLU 주의는 phi = L^{-1} ReLU 로 정의하고 소프트맥스 주의와 비교한다.
• 다양한 비선형성 및 시퀀스 길이 스케일링을 실험하여 정확도를 평가한다.
• 하이퍼파라미터를 변경하지 않고 BigVision의 ImageNet-21k 및 ImageNet-1k 학습 설정을 변경 없이 사용한다.
• qk-layernorm 제거의 효과와 게이트된 주의 유닛의 영향을 평가한다.

실험 결과

연구 질문

• RQ1 시퀀스 길이로 스케일링했을 때 비전 트랜스포머 주의에서 소프트맥스를 ReLU로 대체하면 비슷한 정확도를 달성할 수 있는가?
• RQ2 다양한 비선형성에서 시퀀스 길이 스케일링(L^{-alpha})이 성능에 미치는 영향은?
• RQ3 확대 시 ReLU 주의의 안정성 및 정확도에 있어 qk-layernorm의 역할은 무엇인가?
• RQ4 게이팅 메커니즘을 추가하면 ReLU 주의에서 시퀀스 길이 스케일링의 필요성이 줄어드는가?
• RQ5 이러한 대체들이 학습 계산량과 병렬화에 어떤 영향을 미치는가?

주요 결과

• L^{-1} 스케일링을 갖는 ReLU 주의는 ImageNet-21k 학습에서 softmax 주의의 스케일링 추세와 일치할 수 있다.
• 알파가 1에 가까운 시퀀스 길이 스케일링은 일반적으로 테스트 모델들에서 최상의 결과를 낸다.
• qk-layernorm은 이들 모델에 큰 영향을 주지 않으나 스케일에 따라 영향은 달라질 수 있다.
• 게이팅은 시퀀스 길이 스케일링의 이점을 제거하지 못하며, 최고 정확도는 여전히 L^{-alpha} 스케일링에서 얻어진다.
• ReLU 주의를 사용하면 전통적인 softmax 주의보다 적은 수의 gather 연산으로 시퀀스 길이에 대해 병렬화를 가능하게 한다.