RBF-Attention으로 Transformer를 다시 짜 본 실험, r/MachineLearning 토론 정리

r/MachineLearning에서 165점과 23개 댓글을 받은 이 글은, Transformer의 핵심인 scaled dot-product attention을 distance-based RBF-Attention으로 바꿔 보면 어떤 일이 벌어지는지를 기록한 실험 보고서다. 작성자는 dot product가 벡터 방향뿐 아니라 크기에도 크게 좌우되기 때문에, norm이 큰 key가 softmax를 "밀어붙이는" 문제가 있다고 보고 출발한다. 대안으로 제시한 것은 Euclidean distance 기반 점수화이며, 직관은 간단하다. query와 key가 실제로 가까울 때만 높은 점수를 주자는 것이다.

흥미로운 부분은 아이디어보다 구현이다. 작성자는 단순한 torch.cdist 접근이 즉시 메모리 폭발을 일으킨다고 설명하고, 식을 다시 정리해 사실상 2(Q·K)-||K||^2 형태로 계산할 수 있다고 적는다. softmax의 shift-invariance 덕분에 query norm은 버릴 수 있으므로, 완전히 새로운 연산처럼 보이지만 기존 matmul 친화적 경로를 일부 유지할 수 있다는 주장이다. 그래도 PyTorch의 SDPA에는 key norm penalty를 끼워 넣을 수 없어서, 결국 Triton으로 custom kernel을 작성했다.

• attention sink 문제를 해결하기 위해 register tokens를 별도로 두었다. 거리 기반 공간에서는 거대한 key가 sink가 될 수 없기 때문이다.
• RoPE는 회전 기반 기하라 Euclidean distance와 잘 맞지 않는다고 보고, SuSiE라는 additive positional embedding으로 바꿨다.
• TinyStories 기반의 작은 causal model에서는 기준선보다 약간 빠른 수렴을 봤지만, 작성자 스스로도 이것이 FlashAttention을 대체할 단계는 아니라고 선을 그었다.

댓글에서도 분위기는 비슷했다. 일부는 gaussian-process 계열 kernel attention과의 연결점을 언급했고, 일부는 결국 hardware lottery가 실전 채택을 좌우한다고 지적했다. 요약하면 이 글의 가치는 즉시 상용화될 대체 attention을 제시했다는 데보다, attention의 기본 가정을 다시 실험 가능한 형태로 분해해 보여줬다는 데 있다.

원문은 Reddit 스레드, 상세 설명은 기술 블로그, 코드는 GitHub 저장소에서 볼 수 있다.