채널 인식 비전 트랜스포머로 동적 특징 융합 구현

초록

CAViT은 기존 ViT의 정적 MLP를 대체해 채널 차원에서도 자기‑주의(self‑attention)를 적용함으로써 입력 이미지 전체 컨텍스트에 기반한 동적 채널 재조정을 가능하게 한다. 공간‑주의와 채널‑주의를 순차적으로 수행하는 이중‑주의 구조는 파라미터와 FLOPs를 30% 이상 절감하면서도 CIFAR‑10, Cats‑vs‑Dogs, 의료 영상 3종 데이터셋에서 최고 3.6%p의 정확도 향상을 달성한다.

상세 분석

CAViT의 핵심 아이디어는 ViT 블록 내에서 기존의 MLP‑기반 채널 혼합을 완전히 제거하고, 차원 전치를 통해 채널을 토큰처럼 취급한 뒤 단일 헤드 자기‑주의(SHSA)를 적용하는 것이다. 이 설계는 두 가지 중요한 장점을 제공한다. 첫째, 채널 혼합이 입력 이미지의 전역 컨텍스트에 조건화되므로, 서로 다른 시각 구조(예: 텍스처, 객체 부위, 병변)마다 최적의 채널 가중치를 동적으로 학습한다. 기존 정적 MLP는 모든 입력에 대해 동일한 선형 변환을 수행하므로, 채널 간 상관관계를 충분히 활용하지 못한다는 한계가 있었는데, CAViT은 이를 완화한다. 둘째, 채널‑주의를 단일 헤드로 구현함으로써 연산량을 크게 늘리지 않는다. 전치 후 토큰 수가 C(채널 수)로 변하고, 각 토큰이 전체 이미지 정보를 내포하므로 다중 헤드가 필요 없으며, 이는 FLOPs 절감에 직접 기여한다.

구조적 흐름을 살펴보면, 입력 이미지 I∈ℝ^{W×W×3}를 w×w 패치로 나누어 N개의 패치 토큰을 만든 뒤, CLS 토큰을 추가해 B×(N+1)×C 형태의 텐서를 얻는다. 첫 단계에서 기존과 동일하게 다중 헤드 공간 자기‑주의(MHSA)를 적용한다. 이후 텐서를 (B, C, N) 형태로 전치하고, CLS 토큰을 (B,1,N)으로 변형해 채널 토큰 시퀀스에 삽입한다. 이렇게 구성된 B×(C+1)×N 텐서에 SHSA를 수행하면, 각 채널 토큰이 다른 채널과의 상관관계를 학습한다. 마지막으로 전치를 역전시켜 원래 형태로 복원하고, 다음 블록으로 전달한다.

실험에서는 ViT‑tiny와 동일한 하이퍼파라미터(학습률 0.001, 100 epoch, SGD) 하에 5개의 데이터셋(CIFAR‑10, Cats‑vs‑Dogs, Malaria, PneumoniaMNIST, BreastMNIST)에서 성능을 비교하였다. 파라미터는 5.75M→3.91M(≈32% 감소), FLOPs는 2.267G→1.52G(≈33% 감소)로 크게 줄였음에도 불구하고, 정확도는 CIFAR‑10에서 +3.64%, BreastMNIST에서 +2.50% 등 전반적으로 향상되었다. 특히 의료 영상에서는 작은 해상도와 제한된 데이터량에도 불구하고 채널‑주의가 병변 부위를 더 정확히 강조하는 시각화 결과를 보여, 모델의 해석 가능성도 개선되었다.

추가적인 Ablation 연구에서는 (1) 공간 MHSA만 사용, (2) 채널‑주의를 다중 헤드로 확장, (3) CLS 토큰을 채널‑주의에 포함시키지 않은 경우 등을 비교하였다. 결과는 SHSA와 CLS 토큰 포함이 가장 좋은 성능을 보였으며, 다중 헤드 채널‑주의는 연산량 증가 대비 성능 향상이 미미함을 확인했다. 이는 채널 차원에서 이미 전역 정보를 충분히 담고 있기 때문에 복수 헤드가 필요 없다는 가설을 뒷받침한다.

전반적으로 CAViT은 “채널도 공간처럼 동적으로 혼합한다”는 간단하지만 강력한 설계 원칙을 제시한다. 기존 ViT 구조에 최소한의 변경만으로 파라미터 효율성을 높이고, 다양한 도메인에서 일반화 능력을 강화한다는 점에서, 향후 대규모 비전 파운데이션 모델에 채널‑주의를 기본 모듈로 채택할 가능성을 시사한다.