트랜스포머 마스크드 셀프 어텐션 정리 — 디코더가 미래 단어를 보지 못하게 막고 학습을 병렬화하는 원리

핵심 메시지

• 마스크드 셀프 어텐션은 디코더가 학습 중 모든 위치의 다음 단어 확률을 한꺼번에 계산할 수 있게 해주는 장치다.
• 정답 문장 전체를 디코더에 한 번에 넣는 방식은 학습할 때만 가능하다. 실제 번역 시점에는 정답이 없으므로 이 병렬화는 학습 속도에만 기여한다.
• 핵심 제약은 '커닝 금지'다. i번째 단어의 다음 단어 확률을 계산할 때 신경망은 인코더 출력과 앞선 단어들까지만 볼 수 있어야 한다.
• 구현은 정규화 전 어텐션 점수 가운데 뒤쪽 토큰에 해당하는 값을 0으로 만든 뒤 다시 정규화하는 것으로, 결과적으로 앞쪽 원소들에 대해서만 소프트맥스를 취한 것과 같아진다.
• 마스크를 씌우면 각 임베딩이 자기 앞의 입력에만 의존하게 되지만, 그 대가로 입력 순서가 의미를 갖게 되어 더 이상 집합에서 집합으로의 사상이 아니게 된다.

쉽게 이해하기

트랜스포머의 디코더는 인코더처럼 입력의 형태를 그대로 유지하지만, 층 사이를 오가는 행렬의 크기는 임베딩된 출력 시퀀스에 맞춰져 있다. 즉 벡터의 개수는 언제나 출력 시퀀스의 길이와 같고, 이 값은 인코더가 처리하는 벡터 개수와 일반적으로 다르다. 이 강의는 디코더에 있는 두 어텐션 층 가운데 첫 번째인 마스크드 멀티헤드 셀프 어텐션을 다룬다.

왜 마스크가 필요한지는 학습을 빠르게 하려는 요구에서 출발한다. 다음 단어 예측 확률을 위치마다 하나씩 순서대로 구하는 대신 전부 병렬로 계산할 수 있다면 좋을 것이고, 그러려면 원하는 번역문 전체를 디코더에 한꺼번에 넣어야 한다. 이는 정답 문장을 이미 알고 있는 학습 단계에서만 가능한 방식이며, 실제로 문장을 번역할 때는 쓸 수 없다.

여기서 문제가 생긴다. 두 번째 단어의 확률을 계산할 때 신경망이 그 단어를 이미 볼 수 있다면 과제가 너무 쉬워지고, 보지 못한 단어를 예측하는 능력을 학습하지 못한다. 그래서 두 번째 단어를 계산할 때는 인코더 출력과 첫 토큰만, 네 번째 단어를 계산할 때는 인코더 출력과 앞의 세 토큰만 접근하도록 만들어야 한다.

계산 자체는 일반적인 셀프 어텐션과 크게 다르지 않다. 각 입력 토큰마다 쿼리·키·값을 구하고 키와 쿼리의 내적을 벡터 길이의 제곱근으로 나눠 점수를 얻는다. 다만 소프트맥스를 통과시키기 전에 뒤쪽 토큰에 해당하는 점수를 0으로 만들고 나머지를 합이 1이 되도록 정규화한다. 강의는 이 조작이 결국 앞쪽 점수들에 대해서만 소프트맥스를 적용한 것과 같음을 보이며, 따라서 뒤쪽 점수는 애초에 계산할 필요조차 없다고 지적한다.

행렬 형태로 쓰면 가중치 행렬 세 개로 쿼리·키·값과 점수 행렬을 구한 다음, i번째 단어에 대한 가중치 중 그 뒤 단어들의 몫을 0으로 만드는 마스크를 적용하고 열마다 합이 1이 되게 정규화한다. 그 결과 i번째 열에는 0이 아닌 원소가 i개만 남고, 새 임베딩은 앞의 값 벡터들의 가중합이 된다. 멀티헤드 구조는 인코더와 동일해서, 여러 헤드의 결과 행렬을 이어붙인 뒤 다시 가중치 행렬을 곱해 입력과 같은 차원의 출력을 얻는다.

주요 인사이트

• 마스킹은 모델의 표현력을 키우는 장치가 아니라 학습을 병렬화하기 위한 공학적 타협이다. 순차적으로 계산하면 마스크 없이도 같은 결과를 얻지만 학습이 느려진다.
• '뒤쪽 가중치를 0으로 만들고 다시 정규화한다'는 설명과 '앞쪽 점수에만 소프트맥스를 적용한다'는 설명은 수학적으로 같은 연산이다. 후자로 보면 계산하지 않아도 되는 항이 무엇인지가 분명해진다.
• 마스크드 셀프 어텐션은 순열 불변성을 잃는다. 일반 셀프 어텐션이 집합을 집합으로 옮기는 연산인 반면, 마스크가 들어가면 입력 벡터의 순서가 결과를 바꾼다.
• 강의는 조건부 확률 표기에서 조건 변수를 생략해 표기를 단순화하는데, 실제로는 모든 확률이 인코더 출력과 앞선 토큰들에 조건부라는 점을 놓치지 않아야 한다.

자주 묻는 질문

관련 AI 소식