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

커뮤니티 반응

작성자의 엔지니어링적 시도에 대해 대체로 긍정적인 반응이며, 특히 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)를 도입했다. 미리 계산된 회전하지 않는 사인파를 벡터에 직접 더함으로써 위치적 거리가 유클리드 공간에서 명시적인 페널티로 작용하게 만들었다. 이는 어텐션 메커니즘의 변화가 위치 인코딩 전략까지 연쇄적으로 변경해야 함을 시사한다.

용어 해설

RBF 커널(RBF Kernel): — Radial Basis Function 커널의 약자로, 두 데이터 포인트 사이의 거리를 기반으로 유사도를 측정하는 함수이다. AI 모델에서는 주로 유클리드 거리가 가까울수록 높은 가중치를 부여하는 방식으로 작동하며, 벡터의 방향뿐만 아니라 크기 차이까지 고려할 수 있게 한다. 표준 Dot-Product의 한계를 보완하는 대안적 유사도 측정 지표로 활용된다.
어텐션 싱크(Attention Sink): — LLM이 문맥상 중요하지 않은 정보를 처리할 때 주의력(Attention)을 집중시키는 특정 토큰이나 메커니즘을 의미한다. 주로 첫 번째 토큰이나 특수 토큰이 이 역할을 수행하며, 모델의 안정적인 추론과 긴 문맥 유지에 필수적이다. 거리 기반 어텐션에서는 벡터 크기를 키우는 기존 방식이 작동하지 않아 별도의 레지스터 토큰이 필요하다.
소프트맥스 이동 불변성(Softmax Shift-Invariance): — 입력 벡터의 모든 요소에 동일한 상수를 더하거나 빼도 Softmax 함수의 출력 결과가 변하지 않는 수학적 성질이다. 이 성질 덕분에 어텐션 계산 시 모든 쿼리에 공통된 항(예: 쿼리 벡터의 노름)을 제거하여 연산량을 줄일 수 있다. 메모리 최적화를 위한 수식 전개에서 핵심적인 역할을 하는 원리이다.

언급된 도구

Triton추천

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

FlashAttention중립

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

PyTorch중립

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

언급된 리소스

문서Scaled RBF Attention Blog Post

GitHubRBF Attention GitHub Repository