Hugging Face 어텐션 메커니즘

Hugging Face 어텐션 메커니즘 (Attention Mechanisms)

대부분의 트랜스포머 모델은 어텐션 행렬이 정사각형인 'full attention'을 사용해요. 긴 텍스트에서는 이것이 큰 계산 병목이 될 수 있어요. Longformer와 Reformer 같은 모델은 이 문제를 해결하기 위해 희소(sparse) 어텐션 행렬 버전을 사용해서 학습을 가속화해요.

이 글에서는 Reformer의 LSH 어텐션과 Longformer의 로컬 어텐션을 중심으로 희소 어텐션 방식이 어떻게 동작하는지 살펴볼게요.

출처: Hugging Face Attention Mechanisms 문서

LSH 어텐션 (Reformer)

Reformer는 LSH(Locality Sensitive Hashing) 어텐션을 사용해요. softmax(QK^t)에서 QK^t 행렬의 softmax 차원에서 가장 큰 요소들만 유용한 기여를 해요. 그래서 각 쿼리 q에 대해 K 중 q와 가까운 키 k만 고려하면 돼요.

쿼리와 키가 가까운지 판별하는 데 해시 함수가 사용돼요. 어텐션 마스크는 현재 토큰(첫 번째 위치 제외)을 마스킹하도록 수정돼요. 그러면 쿼리와 키가 같아져(즉 서로 매우 유사해져) 어텐션이 자기 자신에게 과도하게 쏠리는 걸 막아줘요.

해시가 어느 정도 무작위적일 수 있기 때문에, 실제로는 여러 해시 함수(n_rounds 파라미터로 결정)를 쓰고 그 결과를 평균 내요. 이렇게 하면 해시의 무작위성으로 인한 편향을 줄일 수 있어요.

로컬 어텐션 (Longformer)

Longformer는 로컬 어텐션을 사용해요. 종종 로컬 컨텍스트(예: 주어진 토큰의 왼쪽·오른쪽 두 토큰)만으로도 그 토큰에 대한 행동을 결정하기에 충분해요. 그리고 작은 윈도우를 가진 어텐션 레이어를 여러 개 쌓으면, 마지막 레이어는 윈도우 안의 토큰 너머까지 수용 영역(receptive field)을 가지게 돼서 토큰의 더 넓은 표현을 만들 수 있어요.

즉 로컬 어텐션은 각 토큰이 가까운 이웃들만 바라보게 하여 계산량을 줄이면서도, 레이어를 쌓아 올리면 먼 토큰들 간의 관계까지 간접적으로 포착할 수 있게 해줘요.

왜 희소 어텐션이 필요한가

Full attention의 계산 비용은 시퀀스 길이의 제곱에 비례해서 커져요. 긴 문서나 긴 컨텍스트를 다루는 태스크에서는 이 비용이 감당하기 어려울 수 있어요. 희소 어텐션은 불필요한(또는 기여가 작은) 계산을 생략해 계산 비용을 줄이면서도, 핵심적인 원거리 관계는 레이어 구성으로 보완해요.

더 알아보기