Enhancing Layer Interaction Using Key-Correlated Layer Attention

기술 요약: 키 상관관계 레이어 어텐션(Key-Correlated Layer Attention)을 통한 레이어 상호작용 강화

문제 정의
최근 심층 신경망(DNN) 설계의 발전은 레이어 출력을 시퀀스 토큰으로 취급하여 레이어 간 상호작용을 강화하는 **레이어 어텐션(Layer Attention)**을 활용해 왔습니다. 그러나 표준 레이어 어텐션은 네트워크 깊이 $D$에 대해 이차 복잡도($O(D^2)$)를 가지며, 모든 이전 레이어의 출력을 캐싱해야 하기 때문에 상당한 메모리 증가 문제를 겪습니다.

이러한 문제를 해결하기 위해 기존 연구들은 **재귀적 레이어 어텐션(Recurrent Layer Attention, RLA)**과 MRLA-L과 같은 선형 근사법을 도입했습니다. RLA는 키($K$)와 값($V$) 표현을 캐싱함으로써 계산 오버헤드를 줄이지만, 초기 레이어의 특징이 생성된 후 고정되는 정적 메커니즘으로 작동하여 동적인 정보 업데이트를 제한합니다. 또한, MRLA-L과 같은 선형 근사법은 어텐션 출력을 선형적으로 근사하여 복잡도를 줄이려 시도합니다. 저자들은 이러한 접근 방식이 이전 레이어의 기여도를 지수적으로 감쇄시킨다는 점을 관찰했습니다(저자들의 $\beta_l$ 감쇄 계수 분석에서 확인됨). 이는 효과적인 **장거리 교차 레이어 의존성(long-range cross-layer dependencies)**을 구축하는 것을 어렵게 만듭니다. 추가적으로, 기존의 선형화 방법들은 이전 레이어 출력을 캐싱하는 메모리 병목 현상이나 계층적 DNN 구조에 선형 어텐션을 적용할 때 발생하는 특정 과제들을 해결하지 못하는 경우가 많습니다.

방법론: 키 상관관계 레이어 어텐션 (Key-Correlated Layer Attention, KCLA)
제안된 KCLA는 동적 업데이트와 장거리 의존성 모델링을 유지하면서 선형 계산 복잡도를 달urut 달성하도록 설계된 새로운 선형화 레이어 어텐션 방법입니다. 이 방법론은 세 가지 핵심 통찰 및 메커니즘을 기반으로 합니다:

1. 키 상관관계 관찰 (Key Correlation Observation): 경험적 분석 결과, 표준 레이어 어텐션에서 서로 다른 레이어 간의 키 표현은 높은 코사인 유사도를 보입니다. 이를 활용하여 KCLA는 값($V$)의 선형 변환을 분리하고 키 표현이 방향적으로 상관되어 있다고 가정함으로써 표준 레이어 어텐션 출력을 근사합니다. 이를 통해 전체 $K$ 및 $V$ 행렬의 이력을 저장하지 않고도 레이어 출력들을 재귀적 상태($U^t$)로 집계할 수 있습니다.
2. 동적 소프트맥스 복원 (Dynamic Softmax Restoration): 어텐션을 선형화할 때 직면하는 결정적인 과제는 어텐션 점수를 정규화하고 동적인 분포 조정을 가능하게 하는 소프트맥스(softmax) 연산의 손실입니다. KCLA는 이 기능을 복원하기 위해 변형된 멀티 헤드 어텐션 메커니즘을 도입합니다:
   • 로그 균등 분포(logarithmically uniform sequence)에서 유도된 **서로 다른 스케일링 인자(온도, temperatures)**를 가진 여러 헤드를 사용합니다.
   • 쿼리의 노름($\|Q^t\|$)과 상관각($\cos \theta$)을 헤드별 가중치로 매핑하기 위해 학습 가능한 선형 변환 $\Phi$를 활용합니다.
   • 이를 통해 모델은 각 레이어에서 어텐션 분포를 동적으로 조정할 수 있으며, 전체 $N \times N$ 행렬을 계산하지 않고도 소프트맥스의 온도 파라미터 역할을 효과적으로 시뮬레이션할 수 있습니다.
3. 단계별 적응 (Cross-Stage Adaptation): 스테이지마다 다양한 특징 맵 크기와 채널 차원을 갖는 DNN에 대응하기 위해, KCLA는 특정 적응 절차(Algorithm 1)를 구현합니다. 이는 누적된 특징을 정규화하고, 다운샘플링을 위한 스트라이드가 포함된 $1 \times 1$ 컨볼루션을 적용하며, 스테이지 간의 차원을 연결하기 위해 정규화 인자를 재계산하는 과정을 포함합니다.

주요 기여

• 관찰 및 증명: 저자들은 표준 레이어 어텐션의 서로 다른 레이어 간 키 표현이 방향 측면에서 매우 상관되어 있음을 입증하여, 선형 근사의 이론적 근거를 제공합니다.
• 새로운 선형화: KCLA는 RLA의 정적 특성과 MRLA-L의 지수적 감쇄 문제를 피하는 선형화된 레이어 어텐션 방법으로 제안되었습니다. 이는 선형 계산 복잡도와 네트워크 깊이와 무관한 고정된 공간 복잡도를 달성합니다.
• 동적 소프트맥스 근사: 멀티 헤드 스케일링과 학습 가능한 변환을 사용하여 소프트맥스 연산을 근사하는 새로운 방법을 도입하여, 각 레이어에서 동적이고 적응적인 어텐션 분포 조정을 가능하게 합니다.

실험 결과
저자들은 세 가지 핵심 비전 작업인 이미지 분류, 객체 탐지, 의료 영상 분할에서 KCLA를 평가했습니다.

• 이미지 분류: CIFAR-100에서 KCLA는 ResNet 깊이 20, 56, 110 전반에 걸쳐 MRLA-L, DLA-L 및 표준 레이어 어텐션을 능가하는 최고의 성능을 기록했습니다. ImageNet-1K에서 KCLA를 ResNet-50 및 ResNet-101에 통합한 결과, 베이스라인 대비 각각 +1.7% 및 +1.4%의 Top-1 정확도 향상을 보였습니다. 또한, 동적 방식인 DLA-L보다 훨씬 적은 파라미터를 사용하면서도 MRLA-L에 대해 일관되게 우세한 우위를 유지했습니다.
• 객체 탐지: Faster R-CNN 및 Mask R-CNN을 사용한 COCO2017 데이터셋에서 KCLA는 경쟁력 있는 성능을 보여주었습니다. ResNet-101을 사용했을 때, KCLA는 더 복잡한 DLA-L 아키텍처보다 2.2M 개의 파라미터를 적게 사용하면서도 여러 지표(AP$_{75}$, AP$_S$, AP$_M$)에서 이를 앞질렀습니다.
• 의료 영상 분할: ISIC2017, ISIC2018, Kvasir-SEG 데이터셋에서 KCLA는 MRLA-L 및 베이스라인 ResNet-UNet을 지속적으로 능가하며, 최소한의 파라미터 오버헤드로 우수한 mIoU 및 DSC 점수를 달성했습니다.
• 복잡도 분석: 메모리와 FLOPs가 깊이에 따라 이차적으로 증가하는 MRLB와 달리, KCLA는 계산 부하와 메모리 비용 모두에서 선형적 증가를 보입니다.

의의 및 주장
본 논문은 KCLA가 선형 어텐션의 효율성과 표준 어텐션의 표현력 사이의 간극을 성공적으로 메웠다고 주장합니다. 키 상관관계를 기반으로 표준 레이어 어텐션 메커니즘을 직접 근사함으로써, KCLA는 RLA의 정적 특징 표현 문제와 MRLA-L의 장거리 의존성 저하 문제를 방지합니다.

저자들은 제안된 방법이 효율적인 방식으로 장거리 교차 레이어 의존성을 모델링할 수 있는 능력을 유지하며, 효율성과 성능 사이의 바람직한 균형을 제공한다는 점을 강조합니다. 또한, 소프트맥스 함수를 복원하는 방법이 이러한 결과를 얻는 데 핵심적인 요소라고 언급합니다. 현재 구현은 미리 설정된 온도 파라미터에 의존하고 있으나, 저자들은 완전한 자기 적응형 동적 온도 멀티 헤드 선형 레이어 어텐션 설계가 향후 최적화를 위한 방향임을 밝히고 있습니다. 재현성을 높이기 위해 코드는 공개되어 있습니다.