Discrete Cosine Transform Based Decorrelated Attention for Vision Transformers

"Discrete Cosine Transform-Based Decorrelated Attention for Vision Transformers" 논문에 대한 상세한 기술적 요약입니다.

1. 문제 제기

비전 트랜스포머 (ViT) 는 이미지 패치에 자기-주의 (self-attention) 메커니즘을 적용하여 컴퓨터 비전을 혁신하였습니다. 그러나 두 가지 중요한 과제가 남아 있습니다:

• 초기화 민감성: 표준 접근 방식은 쿼리 ($W_Q$), 키 ($W_K$), 값 ($W_V$) 투영 행렬을 무작위 숫자 (예: Xavier 또는 He 초기화) 로 초기화합니다. 이는 가중치를 '백색 잡음'으로 취급하여 스펙트럼을 균일하게 덮지만 구조적 사전 지식이 부족합니다. 이는 특히 작은 데이터셋에서 불안정한 학습 동역학을 초래할 수 있으며, 수렴하기 위해 상당한 계산 자원을 필요로 합니다.
• 계산 효율성: 자기-주의 메커니즘은 계산 비용이 많이 듭니다 (시퀀스 길이에 대한 $O(N^2)$ 복잡도 및 높은 파라미터 수). 기존 압축 방법들은 종종 정확도를 희생하거나 시각 데이터 및 주의 행렬에 내재된 특정 상관 구조를 활용하지 못합니다.

2. 방법론

저자들은 이러한 문제를 해결하기 위해 **이산 코사인 변환 (DCT)**을 활용하는 두 가지 상호 보완적인 방법을 제안합니다:

A. DCT 기반 초기화

무작위 초기화 대신, 저자들은 투영 가중치 행렬 ($W_Q$, $W_K$, 또는 $W_V$) 중 하나를 직교 Type-II DCT 행렬을 사용하여 초기화할 것을 제안합니다.

• 메커니즘: DCT 행렬은 구조화된 특징 추출기로 작용합니다. 무관한 잡음인 무작위 가중치와 달리, DCT 기저 벡터는 전체 주파수 스펙트럼을 체계적으로 커버합니다.
• 근거:
  • 스펙트럼 커버리지: 무작위 초기화는 모든 가중치 벡터를 동일한 스펙트럼 분포로 시작합니다. DCT 초기화는 각 가중치 벡터가 서로 다른 대역폭 (저주파에서 고주파까지) 을 커버하도록 하여, 다양한 특징을 학습하기 위한 더 견고한 시작점을 제공합니다.
  • 비상관화: DCT 기저 벡터는 고도로 상관된 신호 (자연 이미지 등) 에 대한 카를룬 - 로브 변환 (KLT) 을 근사하여 학습 시작 시 입력 특징을 효과적으로 비상관화합니다.
• 구현: 가중치 대칭성 문제 (여러 행렬을 동일한 방식으로 초기화하면 학습이 저하되는 현상) 를 피하기 위해 하나의 행렬만 DCT 로 초기화합니다. 나머지는 무작위로 초기화된 상태로 둡니다.

B. DCT 기반 압축 주의

이 방법은 차원 축소와 계산 비용 감소를 위해 주파수 영역을 활용합니다.

• 메커니즘:
  1. 변환: 입력 패치를 채널 축을 따라 DCT 를 사용하여 주파수 영역으로 변환합니다.
  2. 자르기: 고주파 성분 (일반적으로 잡음과 덜 중요한 세부 사항에 해당) 을 잘라냅니다. 상위 $\tau$ 비율의 계수 (예: 25%, 50%, 75%) 만 유지합니다.
  3. 투영: 축소된 차수의 텐서 ($\tau C$ 대신 $C$) 에 대해 쿼리, 키, 값 투영을 계산합니다.
  4. 재구성: 주의 출력은 원래 차원으로 패딩된 후 역 DCT (IDCT) 를 통해 다시 변환됩니다.
• 최적화: 그 사이에 비선형성이 없으므로, IDCT 와 최종 선형 출력 레이어 ($L_O$) 를 단일 선형 레이어로 융합하여 불필요한 연산을 제거합니다.
• 효율성: 이는 선형 레이어의 계산 복잡도를 $O(C^3)$에서 $O(\tau^2 C^3)$로 줄이고 Softmax 행렬 크기를 크게 감소시킵니다.

3. 주요 기여

1. 새로운 초기화 전략: ViT 주의 가중치에 대한 최초의 DCT 기반 초기화를 도입하여, 구조화된 주파수 영역 사전 지식이 무작위 초기화보다 우수함을 입증했습니다.
2. 효율적인 압축: 고주파 잡음을 제거하면서도 필수적인 저주파 의미 정보를 보존하여 파라미터와 FLOPs 를 줄이는 DCT 기반 압축 기법을 제안했습니다.
3. 이론적 정당성: DCT 와 고도로 상관된 이미지 패치의 비상관화 (KLT 근사) 사이의 수학적 연결을 제공하여, 왜 DCT 가 무작위 잡음이나 단순 풀링보다 주의 압축에 우월한지 설명했습니다.
4. 포괄적인 평가: 분류, 감지, 세그멘테이션 등 여러 작업에서 Swin-T, Swin-S, ViT-B 등 여러 아키텍처에 걸쳐 방법을 검증했습니다.

4. 실험 결과

이 방법들은 Swin 트랜스포머를 백본으로 사용하여 CIFAR-10, ImageNet-1K, COCO 2017 벤치마크에서 평가되었습니다.

• 이미지 분류 (CIFAR-10 및 ImageNet-1K):

  • 초기화: $W_K$(키) 를 DCT 로 초기화하는 것이 Top-1 정확도를 일관되게 향상시켰습니다. ImageNet-1K 의 Swin-T 의 경우 정확도가 **81.47% 에서 81.73%**로 증가했습니다.
  • 압축: DCT 계수의 75% ($\tau=0.75$) 를 유지하는 것은 Swin-T 의 경우 파라미터를 약 13%, Swin-S 의 경우 약 10% 줄이면서도 기준선 대비 정확도를 유지하거나 약간 향상시켰습니다.
  • 비교: 제안된 방법은 저주파 성분을 평균내는 대신 정확하게 유지함으로써 DCFormer(이전 DCT 기반 방법) 보다 우수한 성능을 발휘하여 판별력 있는 특징을 보존했습니다.

• 객체 감지 및 세그멘테이션 (COCO 2017):

  • DCT 초기화 ($W_K$) 와 75% 압축을 사용하면, 파라미터 수를 Swin-T 는 3M, Swin-S 는 5M 줄이면서도 일반 Swin 트랜스포머보다 높은 평균 정밀도 (AP) 점수를 달성했습니다.

• 절대 실험 (Ablation Studies):

  • 단일 vs 다중 초기화: 여러 행렬 (예: $W_Q$와 $W_K$ 모두) 을 DCT 로 초기화하면 대칭성 깨짐 문제로 인해 성능이 저하되었으며, 단일 행렬을 초기화하는 것이 가장 좋은 결과를 보였습니다.
  • DCT 의 필요성: 압축 파이프라인에서 DCT 변환을 제거하고 공간 가지치기만 남기면 정확도가 크게 떨어졌는데, 이는 방법의 성공을 위해 주파수 영역 비상관화가 필수적임을 확인시켜 주었습니다.

5. 의의

이 논문은 주파수 영역 분석이 비전 트랜스포머를 최적화하는 강력한 도구임을 보여줍니다.

• 학습 안정성: 심층 주의 네트워크를 초기화하는 원칙적이고 비무작위적인 방식을 제공하여, 데이터가 부족한 작업에서 학습 시간을 단축하고 수렴을 개선할 수 있습니다.
• 손실 없는 효율성: 고주파 잡음을 주의 메커니즘에서 안전하게 제거할 수 있음을 증명하여, 성능을 훼손하지 않고 (심지어 향상시키기도 하며) 파라미터가 적고 FLOPs 가 낮은 경량 모델을 가능하게 합니다.
• 일반화 가능성: 이 접근법은 아키텍처에 구애받지 않으며 (Swin 과 ViT 에서 테스트됨), 작업에도 구애받지 않아 (분류, 감지, 세그멘테이션 모두 작동) ViT 도구에 다재다능한 추가 요소가 됩니다.

구현 코드는 공개되어 있어 주파수 기반 주의 메커니즘에 대한 추가 연구를 용이하게 합니다.