명시적 Attention 없이 선형 시간 복잡도를 갖는 글로벌 시각 모델링

Transformer의 핵심인 Attention 메커니즘은 시퀀스 길이의 제곱에 비례하는 연산량 때문에 고해상도 이미지 처리에 한계가 있습니다. 이 논문은 Attention을 입력에 따라 가중치가 변하는 동적 MLP로 재해석하여, 연산량을 선형적으로 줄이면서도 글로벌 정보를 효과적으로 처리할 수 있는 새로운 설계 패러다임을 제시합니다.

기존 Transformer의 Self-Attention은 모든 토큰 쌍의 유사도를 계산하는 Softmax(QKᵀ) 과정을 거치며, 이 과정에서 시퀀스 길이 N의 제곱(N²)에 비례하는 비용이 발생한다. 연구진은 이 수식을 다시 들여다보아, Q를 입력으로 받고 Kᵀ와 V를 각각 첫 번째와 두 번째 층의 가중치로 사용하는 2층 MLP 구조와 동일하다는 점을 발견했다. 즉, Attention의 본질은 '토큰 간의 직접적인 섞임'이 아니라 '입력에 최적화된 가중치를 실시간으로 만들어내는 과정'에 있다는 것이다.

이러한 통찰을 바탕으로, 굳이 N×N 크기의 거대한 유사도 행렬을 만들지 않고도 글로벌 정보를 압축한 가중치를 직접 예측하는 방식을 고안했다. 입력 이미지 전체의 통계적 특성을 추출하여 작은 크기의 가중치 텐서를 생성하고, 이를 통해 입력을 변환하면 Attention과 유사한 글로벌 모델링 효과를 얻을 수 있다.

결과적으로 WeightFormer는 시퀀스 길이에 따라 연산량이 폭증하지 않으면서도 이미지 전체를 한 번에 조망하는 능력을 갖추게 되었다. 이는 고해상도 이미지(1248x1248 등)에서 DeiT 대비 7.7배 빠른 속도와 91%의 메모리 절감 효과를 보여주며, 효율성과 성능의 새로운 균형점을 제시한다.

Attention 수식 O = Softmax(QKᵀ/√d)V에서 Q를 입력 x로 보면, Kᵀ와 V는 입력에 의존하여 생성되는 동적 가중치 W₁, W₂로 해석된다. 연구진은 이를 일반화하여 입력 X로부터 동적 가중치 업데이트 값인 ΔW를 예측하는 함수 G(X)를 정의했다. 구체적으로 W(X) = W₀ + ΔW(X) 형태를 취하며, 여기서 W₀는 학습된 정적 가중치이고 ΔW는 입력의 글로벌 컨텍스트를 반영한다.

가중치 예측을 위해 두 가지 주요 전략을 사용한다. 첫째는 Pooling 기반 전략으로, Global Average Pooling(GAP)을 통해 시퀀스를 고정된 크기의 벡터로 압축한 뒤 MLP를 통과시켜 가중치를 생성한다. 둘째는 Correlation 기반 전략으로, XᵀX 연산을 통해 채널 간의 상관관계를 추출하고 이를 비선형 변환하여 가중치를 얻는다. [입력 X(N×d) → XᵀX 연산 → d×d 상관 행렬 생성 → 비선형 층 통과 → 동적 가중치 ΔW 도출] 과정을 통해 토큰 수 N에 무관한 연산이 수행된다.

WeightFormer 아키텍처는 이러한 동적 가중치 층을 전략적으로 배치한다. 모든 블록에 적용하는 대신 3개 블록마다 하나씩 동적 블록을 삽입하는 Sparse Distribution 전략을 채택하여 성능과 효율성의 최적점을 찾았다. 동적 블록 내부에서는 Spatially Adaptive Prediction을 적용한 Depthwise Convolution과 Bilateral Activation을 적용한 MLP를 결합하여 로컬 특징과 글로벌 컨텍스트를 동시에 학습한다.

ImageNet-1K 분류 실험에서 WeightFormer-S는 81.3%의 정확도를 기록하며 유사한 파라미터 규모의 DeiT-S(79.8%)와 ConvNeXt-S(79.7%)를 능가했다. 특히 연산량(FLOPs) 측면에서 DeiT-S(4.6G)보다 낮은 4.4G를 사용하면서도 더 높은 성능을 보였다.

고해상도 환경에서의 효율성 분석 결과, 1248×1248 해상도에서 WeightFormer는 DeiT 대비 7.7배 높은 처리량(Throughput)을 기록했으며, GPU 메모리 사용량은 91% 감소했다. 이는 Attention의 제곱 복잡도 문제를 선형 복잡도로 성공적으로 대체했음을 증명한다.

객체 탐지(COCO) 및 의미론적 분할(ADE20K) 작업에서도 WeightFormer는 강력한 성능을 보였다. ADE20K에서 WeightFormer-T는 40.7 mIoU를 달성하여 DeiT-T(39.2 mIoU)보다 우수한 성능을 보였으며, 백본 FLOPs는 11G에서 7G로 크게 줄어들었다.

WeightFormer는 Attention을 '입력 조건부 파라미터 생성기'로 보는 관점에서 출발한다. 표준 Attention의 복잡도가 O(N²d)인 이유는 N×N 크기의 어텐션 맵을 명시적으로 생성하기 때문인데, WeightFormer는 이를 O(N)으로 줄이기 위해 공간적 차원을 고정된 크기 M으로 압축하는 ϕ(X) 함수를 도입한다.

핵심 메커니즘 중 하나인 Bilateral Activation은 가중치 예측 과정을 두 개의 상호 보완적인 비선형 브랜치로 분리한다. ΔW(X) = W₁σ(W₂Xᵀ)σ(XW₃)W₄ 수식을 통해 입력 X와 그 전치 행렬 Xᵀ에 독립적으로 작용하는 브랜치를 구성하며, 이는 낮은 연산 비용으로도 풍부한 표현력을 확보하게 한다. [입력 X → 두 갈래의 비선형 변환 σ → 요소별 곱셈 → 최종 가중치 생성] 순으로 연산이 이루어진다.

Effective Receptive Field(ERF) 분석을 통해, 학습 전에는 로컬에 국한되었던 수용 영역이 학습 후에는 이미지 전체로 확장됨을 확인했다. 이는 동적 가중치가 실제로 글로벌 컨텍스트를 캡처하여 입력 토큰들에 적용하고 있음을 시각적으로 증명하는 결과이다.

현재 연구는 시각 작업(Vision Tasks)에 국한되어 평가되었으며, 언어 모델링 등 다른 도메인으로의 일반화 가능성은 아직 명확하지 않다. 또한 동적 파라미터의 표현력과 유도 편향(Inductive Bias)에 대한 이론적 이해가 더 필요하며, 입력 조건부 가중치 생성 과정에서 발생할 수 있는 그래디언트 흐름의 불안정성 문제를 해결하기 위한 추가 연구가 요구된다.