트랜스포머 어텐션의 낭비를 고치는 XSA — 자기 값 벡터 제거 두 줄 코드 해설

핵심 메시지

• 어텐션 출력은 다른 토큰에서 온 '맥락 정보'와 현재 토큰 자신의 값 벡터라는 두 성분으로 나뉜다. 후자는 이미 잔차·피드포워드 경로가 표현할 수 있어 어텐션 안에서는 낭비다.
• 이 뒤섞임은 어텐션이 맥락 관계 모델링과 개별 특징 표현 사이에서 불필요하게 경쟁하게 만든다.
• 현재 토큰을 단순히 무시하는 방법은 통하지 않는다. 비슷한 토큰끼리 값 벡터가 상관돼 자기 값이 새어 들고, 어텐션 싱크의 유연성도 잃기 때문이다.
• 해법은 직교 투영으로 어텐션 출력에서 자기 값 벡터 성분을 제거하는 것이며, 쿼리·키·값 행렬은 그대로 두고 마지막에 기하학적 보정 한 번만 적용한다.
• 코드 두 줄이면 되고 연산 부담이 거의 없으며, 모델 크기·학습률·시퀀스 길이 전반에서 성능이 일관되게 개선되고 특히 긴 시퀀스일수록 이득이 커진다.

쉽게 이해하기

트랜스포머는 어텐션 층과 피드포워드 신경망(FFN) 층을 쌓아 입력 시퀀스를 처리한다. 어텐션 층은 각 토큰에 필요한 맥락을 모은다. 토큰 임베딩에 값 행렬 WV를 곱해 값 벡터를 얻고, 쿼리와 키의 스케일드 닷프로덕트에 소프트맥스를 적용한 가중치로 값 벡터를 가중합한 뒤 출력 행렬 WO로 투영해 잔차 연결로 임베딩을 갱신한다. 그런데 이 합을 뜯어보면 문제가 드러난다.

어텐션 출력의 합은 이전 토큰들의 기여(우리가 원하는 풍부한 맥락 정보)와 현재 토큰 자신의 기여(자기 값 벡터 VI)로 나뉜다. 그런데 VI는 임베딩 XI에 선형 변환을 적용한 방향이라, 이미 잔차와 FFN 경로가 표현할 수 있는 방향이다. 결국 어텐션이 맥락 관계와 개별(점별) 특징을 한꺼번에 담으려다 둘 사이에서 불필요한 경쟁이 생긴다.

그렇다면 현재 토큰을 아예 무시하면 될까? 두 가지 이유로 통하지 않는다. 첫째, 한 시퀀스 안의 값 벡터들은 양의 상관을 갖는 경향이 있다. '네 차로 갈까 내 차로 갈까'에서 두 번째 'car'는 첫 번째 'car'에 강하게 어텐션하며 그 값 방향을 흡수해, 정작 필요한 맥락을 밀어낸다. 가중치를 0으로 둬도 비슷한 토큰을 통해 자기 값 방향이 새어 든다. 둘째, 현재 토큰 어텐션을 막으면 무의미한 토큰에 과한 가중치를 피하게 해주는 '어텐션 싱크' 유연성을 잃어 성능이 떨어진다.

실제로 시퀀스 길이 2048에서 1000억 토큰으로 학습한 모델을 보면, 현재 토큰에 대한 평균 어텐션 가중치가 높고, 층이 깊어질수록 어텐션 출력이 자기 값 벡터 방향과 정렬되는 정도(코사인 유사도)가 커진다. 논문은 이를 '어텐션 유사도 편향'이라 부른다. 어텐션이 다른 토큰의 맥락을 활용하기보다 자기 특징 변환에 용량을 낭비한다는 뜻이다. 그래서 핵심 아이디어는 어텐션 출력에서 자기 값 벡터 성분을 제거하는 것이다.

제거는 직교 투영으로 한다. 어떤 벡터든 값 벡터 V 방향 성분과 그에 수직인 성분으로 정확히 분해되므로, 출력 Y를 Z + α·V로 쓰고 스칼라 α를 구해 α·V를 빼면 V에 수직인 Z가 남는다. 이 방식(XSA, 배타적 자기 어텐션)은 어텐션 가중치를 어떻게 계산했든 상관없이 작동하고, 쿼리·키·값 행렬을 바꿀 필요 없이 마지막에 보정 한 번만 하면 된다. 코드로는 값 벡터를 단위 길이로 정규화하고, 출력에서 그 방향 성분을 빼는 두 줄이면 충분하다. 0.7·1.4·2.7B 모델 모두에서 XSA가 기준선보다 낮은 손실을 일관되게 유지했고, 다운스트림 성능 향상은 모델이 클수록 커졌으며, 연산·메모리 부담은 미미했다. 시퀀스가 512에서 1만 6천 토큰으로 길어질수록 격차가 벌어졌다. 다만 훨씬 더 큰 규모에서도 이득이 지속될지는 아직 열린 물음이다.

주요 인사이트

• 어텐션의 진짜 역할은 다른 토큰에서 맥락을 모으는 것인데, 자기 값 벡터가 섞이면 그 용량이 개별 특징 표현에 낭비된다.
• 자기 토큰을 단순히 배제하는 직관적 해법이 실패하는 이유는, 값 벡터의 상관성 때문에 유사 토큰을 통해 자기 값이 새어 들고 어텐션 싱크의 이점도 사라지기 때문이다.
• '어텐션 유사도 편향'은 층이 깊어질수록 커지므로, 문제는 깊은 모델에서 더 두드러진다.
• XSA는 학습된 가중치 계산 방식을 건드리지 않고 마지막에 기하학적 보정만 더하므로, 기존 트랜스포머에 그대로 끼워 넣을 수 있다.
• 이득이 시퀀스 길이에 따라 커지는 것은, 맥락이 길수록 먼 위치의 정보를 정확히 통합하는 일의 가치가 커지고 자기 정보 중복 인코딩의 낭비가 더 커지기 때문이다.

자주 묻는 질문

관련 AI 소식