셀프 어텐션 파이토치 구현 해설, 쿼리·키·값 행렬과 소프트맥스로 이해하는 어텐션 층

핵심 메시지

• 셀프 어텐션 층은 입력 차원과 내부 차원만 정하면 선형 레이어 세 개(Q·K·V)로 구성되는 단순한 구조다.
• 점수 행렬은 쿼리와 전치된 키의 배치 행렬 곱으로 만들어지며, 전치를 빼면 차원이 맞지 않아 연산 자체가 성립하지 않는다.
• 내부 차원의 제곱근으로 점수를 나누는 스케일링은 벡터 원소 개수에 따라 점수 크기가 달라지는 문제를 막기 위한 장치다.
• 소프트맥스를 마지막 차원에 적용하면 점수 행렬의 각 행 합이 1이 되어, 점수가 곧 확률적 가중치로 바뀐다.
• 어텐션 층의 출력 텐서는 입력 텐서와 같은 크기를 유지하기 때문에 층을 계속 쌓아 올릴 수 있다.

쉽게 이해하기

어텐션이 무엇인지 설명하는 글은 많지만, 여전히 추상적으로 느껴지는 사람이 적지 않다. 이 영상의 발표자도 같은 문제의식에서 출발한다. 개념 설명은 앞선 영상에 맡기고, 이번에는 오직 코드로만 셀프 어텐션 층을 처음부터 만들어 보이며 텐서들이 서로 어떻게 맞물리는지를 눈으로 확인시킨다.

구현은 파이토치의 nn.Module을 상속한 Attention 클래스로 시작한다. 생성자는 들어오는 벡터의 크기인 입력 차원과 어텐션 내부에서 쓰일 쿼리·키·값 벡터의 크기인 출력 차원, 두 개의 인자만 받는다. 그리고 입력 차원을 출력 차원으로 옮기는 선형 레이어를 똑같이 세 번 만든다. 슬라이드에서 WQ·WK·WV로 불리던 행렬들이 여기서는 그냥 Q·K·V 선형 레이어가 된다.

순전파 함수는 입력 텐서를 세 선형 레이어에 각각 통과시켜 쿼리·키·값을 만드는 것으로 시작한다. 이어 키의 차원을 전치한 뒤 쿼리와 배치 행렬 곱을 취해 점수 행렬을 얻는다. 배치 행렬 곱을 쓰는 이유는 입력 텐서의 첫 번째 차원이 배치이고 나머지가 시퀀스 길이와 벡터 크기이기 때문이다. 발표자는 전치를 하지 않으면 이 곱셈이 아예 성립하지 않는다는 점을 실제로 차원을 찍어 보여 준다.

점수는 곧바로 쓰이지 않는다. 출력 차원의 제곱근으로 나눠 정규화하는데, 이는 각 벡터에 들어 있는 원소 개수와 무관하게 점수 규모를 유지하기 위한 조치라고 설명한다. 그다음 마지막 차원을 기준으로 소프트맥스를 취해 어텐션 가중치를 만들고, 이 가중치와 값 행렬을 다시 배치 행렬 곱으로 결합해 은닉 상태를 만들어 반환한다.

동작 확인을 위해 아주 작은 장난감 예제가 동원된다. 문장 몇 개에서 소문자화와 공백 분리로 고유 토큰을 뽑아 인덱스-단어 사전을 만들고, 이를 뒤집어 단어-인덱스 사전까지 준비한다. 문장을 받아 인덱스 목록으로 바꾼 뒤 torch.long 타입 텐서로 만들고 맨 앞에 배치 차원을 하나 끼워 넣는 변환 함수도 함께 만든다. 은닉 크기를 10으로 두고 임베딩을 통과시키면 5개 토큰짜리 문장이 1×5×10 텐서가 된다.

주요 인사이트

• 입력 차원과 출력 차원을 같은 값으로 두면 세 선형 레이어가 모두 정방 행렬이 되는데, 이는 어텐션 층을 반복해서 쌓기 좋게 만드는 실용적인 선택이다.
• 소프트맥스 결과를 특정 차원으로 합해 1이 나오는지 확인하는 습관은, 어텐션이 제대로 확률 가중치로 작동하는지를 검증하는 가장 값싼 디버깅 방법이다.
• 5개 토큰 입력에 대해 점수 행렬이 5×5 정방 행렬로 나온다는 사실은, 어텐션 연산량이 시퀀스 길이의 제곱에 비례한다는 잘 알려진 특성을 코드 차원에서 직접 확인시켜 준다.
• 개념 설명과 코드 구현을 분리해 두 편으로 나눈 접근은, 수식으로는 이해했다고 생각했지만 막상 구현하지 못하는 학습자의 간극을 정확히 겨냥한다.

자주 묻는 질문

관련 AI 소식