LG AI Research 293

기존 Transformer의 Multi-Head Attention은 시퀀스 길이의 제곱에 비례하는 계산 비용이 발생하여, 긴 입력 시퀀스를 처리하는 데 한계가 있다. 기존의 Sparse attention 방식은 과도한 Inductive bias로 인해 성능과 비용 간의 트레이드오프가 발생하는 문제가 있었다.

이를 해결하기 위해 제안된 SBM-Transformer는 각 attention head에 **mixed-membership Stochastic Block Model (SBM)**을 부여한다. 이를 통해 데이터의 특성에 맞춰 Dense 또는 Sparse attention 사이를 유연하게 조절하며, 계산 비용을 최적화할 수 있다.

SBM-Transformer의 핵심 동작 방식은 다음과 같다:

• 그래프 샘플링: 각 attention head는 입력 토큰을 노드로 보고, Query와 Key를 연결하는 이분 그래프(Bipartite graph)를 샘플링한다.
• 파라미터화: Query/Key representation을 MLP를 통해 노드 표현 공간으로 옮긴 후, Cluster embedding과 내적하여 클러스터 멤버십을 결정한다.
• fastRG 알고리즘: 샘플링 단계에서 fastRG 알고리즘을 사용하여 효율적으로 그래프를 생성하며, 이 과정은 병렬 실행이 가능해 가속화할 수 있다.

또한, 그래프 샘플링의 불연속성(Non-differentiability) 문제를 해결하기 위해 **Straight-Through Estimator (STE)**를 활용하여 역전파(Backpropagation)가 가능하도록 설계되었다.