Dot-Product Attention을 거리 기반 RBF-Attention으로 교체한 실험기

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

표준 Dot-Product Attention의 벡터 크기 편향 문제를 해결하기 위해 RBF 기반 거리 어텐션을 구현하고 Triton 커널과 레지스터 토큰으로 기술적 한계를 극복한 사례이다.

배경

표준 Dot-Product Attention의 'Magnitude Bullying' 문제를 해결하기 위해 RBF(Radial Basis Function) 기반 거리 어텐션으로 교체하고, 이 과정에서 발생한 기술적 난관과 해결책을 공유했다.

의미 / 영향

이 실험은 표준 어텐션 메커니즘의 대안을 탐색하며 하드웨어 최적화와 수학적 정합성 사이의 균형을 맞추는 과정을 잘 나타냈다. 현재의 ML 스택이 Dot-Product에 고도로 최적화되어 있어 새로운 연산 도입이 어렵지만, 특정 도메인에서는 거리 기반 방식이 수렴 속도 향상 등의 이점을 제공할 가능성을 확인했다.

커뮤니티 반응

작성자의 엔지니어링적 시도에 대해 대체로 긍정적인 반응이며, 특히 Triton 커널 구현과 수학적 최적화 과정이 유익하다는 평가가 많다.

실용적 조언

커스텀 어텐션 구현 시 Softmax의 이동 불변성을 활용해 수식을 단순화하면 메모리 사용량을 획기적으로 줄일 수 있다.
거리 기반 어텐션을 사용할 때는 Attention Sink 역할을 할 별도의 Register Token을 추가해야 모델이 안정적으로 수렴한다.

섹션별 상세

표준 Dot-Product Attention에서는 키 벡터의 크기가 크면 정렬이 완벽하지 않아도 Softmax 점수를 독점하는 'Magnitude Bullying' 현상이 발생한다. RBF 어텐션은 쿼리와 키 사이의 고차원 공간상 유클리드 거리를 측정하여 실제 근접성을 기준으로 점수를 산출한다. 이를 통해 단순히 벡터 크기만 키워 점수를 가로채는 행위를 방지하고 기하학적 정합성을 확보한다. 실무적으로는 모델이 특정 토큰에 과도하게 의존하는 현상을 완화하는 효과를 기대할 수 있다.

torch.cdist를 사용하여 유클리드 거리를 직접 계산하면 N x N 거리 행렬이 메모리에 생성되어 즉시 OOM이 발생한다. 작성자는 제곱 거리 공식을 전개하여 -||Q||2 - ||K||2 + 2(Q · K) 형태로 변환하는 수학적 트릭을 적용했다. Softmax의 이동 불변성 덕분에 쿼리 노름인 -||Q||2을 제거할 수 있으며, 결과적으로 RBF 어텐션은 키 벡터에 L2 페널티를 부여한 Dot-Product 연산으로 귀결된다. 이 최적화를 통해 기존 행렬 곱셈 기반의 효율적인 연산 구조를 유지하면서 거리 기반 어텐션을 구현했다.

PyTorch의 표준 SDPA는 퓨즈드 루프 내부에서 임의의 키 노름 페널티를 차감하는 기능을 지원하지 않는다. 작성자는 이를 해결하기 위해 FlashAttention의 타일링 로직을 모방한 커스텀 Triton 커널을 직접 작성했다. 이 커널은 SRAM에서 실시간으로 키의 제곱 L2 노름을 계산하고 Softmax 직전에 이를 차감하여 선형적인 메모리 사용량을 유지한다. 하드웨어 가속기 최적화가 Dot-Product에 편중된 상황에서 커스텀 연산 도입을 위한 엔지니어링적 노력을 입증했다.

모델은 무의미한 토큰에 주의력을 집중시키는 'Attention Sink'를 생성하기 위해 의도적으로 벡터 크기를 키우는 경향이 있다. 하지만 거리 기반 수학에서는 벡터 크기가 커질수록 거리가 무한해져 확률값이 0에 수렴하므로 기존 방식의 싱크 생성이 불가능하다. 작성자는 학습 가능한 더미 벡터인 'Register Tokens'를 시퀀스 앞에 추가하고 이를 0으로 초기화하여 해결했다. 쿼리가 유용한 정보를 찾지 못할 때 이 레지스터 토큰으로 주의력을 안전하게 분산시켜 실제 토큰의 정보가 오염되는 것을 방지했다.

현대 모델에서 널리 쓰이는 RoPE는 벡터를 회전시켜 상대적 각도를 학습하지만, 이는 절대적 유클리드 거리를 측정하는 RBF 어텐션의 기하학적 구조와 충돌한다. 작성자는 RoPE를 제거하고 대신 가산적 방식인 SuSiE(Subspace Sinusoidal Embeddings)를 도입했다. 미리 계산된 회전하지 않는 사인파를 벡터에 직접 더함으로써 위치적 거리가 유클리드 공간에서 명시적인 페널티로 작용하게 만들었다. 이는 어텐션 메커니즘의 변화가 위치 인코딩 전략까지 연쇄적으로 변경해야 함을 시사한다.

실무 Takeaway

RBF 어텐션은 수학적 전개를 통해 키 벡터에 L2 페널티를 부여한 형태로 변환 가능하며, 이를 통해 메모리 효율적인 구현이 가능하다.
표준 라이브러리가 지원하지 않는 커스텀 어텐션 연산을 위해 Triton을 활용한 FlashAttention 스타일의 커널 작성이 필수적이다.
거리 기반 어텐션 도입 시 RoPE와 같은 회전 기반 위치 인코딩은 기하학적 정합성이 깨지므로 가산적 방식인 SuSiE 등으로 대체해야 한다.

언급된 도구

Triton추천

커스텀 GPU 커널 작성 및 최적화

FlashAttention중립

효율적인 어텐션 연산 아키텍처

PyTorch중립

딥러닝 모델 구현 및 실험 프레임워크

언급된 리소스

문서Scaled RBF Attention Blog Post

GitHubRBF Attention GitHub Repository