Dot-Product Attention의 숨겨진 Norm Bias와 RBF 기반 해결책

Context

Standard Dot-Product Attention 구조 내 Key의 큰 Norm이 스코어를 지배하는 Magnitude Bullying 현상 발생. 특정 토큰이 확률 질량을 독점하는 Attention Sinks 현상은 이러한 스코어링 규칙의 부작용으로 분석됨.

Technical Solution

• 거리 기반 스코어링 방식인 RBF Attention 도입을 통한 유클리드 공간 내 실제 근접도 기반 유사도 측정
• RBF Attention을 Dot-Product Attention과 Key에 대한 L2 Penalty 결합 형태로 재정의하여 수식적 등가성 확보
• Key Norm의 무분별한 증가를 억제하여 모델이 메트릭적으로 의미 있는 클러스터 표현을 학습하도록 유도
• Selectivity 제어 권한을 Key에서 Query 중심으로 이동시켜 특정 토큰의 과도한 지배력 제거
• FlexAttention의 score_mod API를 활용해 커스텀 스코어 로직을 작성하고 FlashAttention 수준의 fused kernel로 컴파일하는 최적화 전략
• Triton 커널 직접 구현을 통해 $N \times N$ 거리 행렬의 materialized 메모리 오버헤드를 방지하는 타일링 로직 적용

Key Takeaway

Attention 메커니즘의 성능은 단순한 수식 선택이 아니라 하드웨어 가속기(Fused Kernels)와 위치 임베딩(RoPE) 등 전체 스택의 최적화 방향에 종속됨. 새로운 지표 도입보다 기존 SDPA에 특정 Inductive Bias를 주입하는 진단적 접근이 엔지니어링 비용 효율적임.