LLM 어텐션 메커니즘 직접 구현하기: 자기어텐션과 쿼리·키·값, 인과 마스크와 멀티헤드까지

핵심 메시지

• 저자는 어텐션을 자동차의 엔진에 비유한다. 여기서 만드는 언어모델은 수천억이 든 페라리가 아니라 혼자서도 손볼 수 있는 옛 머스탱에 가깝지만, 쓰이는 방법 자체는 같다는 것이다.
• 어텐션이 필요한 이유는 순환 신경망처럼 입력 전체를 중간 상태 하나로 압축하면 정보를 잃기 때문이다. 어텐션은 매 생성 단계마다 입력 전체를 참조하되 토큰마다 다른 중요도를 매긴다.
• 문맥 벡터는 입력 임베딩들의 가중합이며 그 가중치가 어텐션 가중치다. 기준이 되는 토큰을 쿼리라 부르고, 이를 모든 입력으로 일반화하면 6×6 형태의 어텐션 행렬이 나온다.
• 학습 가능한 가중치를 넣으면 쿼리·키·값 세 행렬이 생기고, 어텐션 점수를 키 차원의 제곱근으로 나눈 뒤 소프트맥스를 거쳐 값 벡터에 곱해 문맥 벡터를 만든다.
• 미래 토큰을 가리는 인과 마스크와 과적합을 줄이는 드롭아웃 마스크를 얹으면 인과 어텐션이 되고, 이를 여러 개 병렬로 두면 멀티헤드 어텐션이 된다.

쉽게 이해하기

이 영상은 '밑바닥부터 만드는 대규모 언어모델' 책의 3장을 코드와 함께 따라가는 보조 강의다. 저자는 이 장을 책에서 가장 기술적인 부분이라 부르면서 자동차 엔진에 비유한다. 여기서 만드는 모델은 수백만 달러와 정밀 장비가 필요한 페라리가 아니라, 혼자서도 복원하고 고칠 수 있는 1960년대 머스탱에 가깝다. 성능은 최신 모델에 못 미쳐도 작동 원리를 이해하는 데는 충분하고, 쓰이는 기법 자체는 실제 언어모델과 같다는 것이 요지다.

어텐션이 왜 필요한지는 예전 구조의 약점에서 나온다. 순환 신경망으로 독일어를 영어로 옮길 때는 입력 문장을 끝까지 읽어 하나의 압축된 표현으로 저장한 뒤 그것만 보고 출력을 만들어야 했고, 이 과정에서 정보가 새어 나가 번역 품질이 떨어졌다. 자기어텐션의 발상은 출력을 만드는 매 순간 입력 전체를 다시 볼 수 있게 하는 것이다. 게다가 모든 토큰을 똑같이 보는 것이 아니라, 지금 만들려는 단어에 중요한 토큰에는 큰 가중치를 준다. 이 중요도가 곧 어텐션 점수이자 가중치다.

코드에서는 '너의 여정은 한 걸음에서 시작된다'라는 짧은 문장을 각 단어가 3차원 임베딩인 텐서로 두고 출발한다. 두 번째 토큰을 기준(쿼리)으로 삼아 모든 입력과의 내적을 구해 점수를 만들고, 소프트맥스로 합이 1이 되게 정규화한 뒤, 그 가중치로 입력 벡터들을 가중합해 문맥 벡터를 얻는다. 이렇게 얻은 벡터는 모든 입력의 정보를 담되 두 번째 토큰에 특화돼 있다. 이어서 기준을 하나씩 바꾸는 대신 반복문을 하나 더 두어 6×6 어텐션 행렬로 일반화하는데, 입력 자신에 대한 가중치를 구한다는 점에서 '자기'어텐션이라는 이름이 붙었다.

그다음이 학습 가능한 가중치를 넣는 단계다. 입력을 각각 다른 행렬로 선형 변환해 쿼리·키·값을 만들고, 쿼리와 키의 내적으로 점수를 구한 뒤 키 차원의 제곱근으로 나눠 값이 지나치게 커지지 않게 한다. 이 정규화는 트랜스포머를 소개한 'Attention Is All You Need' 논문에서 제안된 방식이다. 소프트맥스를 거친 가중치는 이제 원래 입력이 아니라 값 벡터에 곱해져 문맥 벡터가 된다. 저자는 행렬을 꼭 세 개 둘 필요는 없고 수를 줄여 실험한 독자도 있었다며, 이 설계가 이론적 필연이라기보다 경험적으로 잘 작동해 자리 잡은 것임을 덧붙인다.

마지막으로 두 종류의 마스크와 멀티헤드가 붙는다. 인과 마스크는 대각선 위쪽, 즉 미래 토큰을 가려 모델이 다음 단어를 미리 훔쳐보지 못하게 한다. 드롭아웃 마스크는 어텐션 가중치 행렬에서 임의의 위치를 지워 특정 위치에 덜 의존하게 만들고, 지워진 만큼 남은 값을 키워 합을 비슷하게 유지한다. 멀티헤드는 인과 어텐션을 여러 개 만들어 결과를 이어 붙이는 방식으로 먼저 구현한 뒤, 헤드끼리 서로 기다릴 이유가 없다는 점을 살려 큰 행렬 곱 한 번으로 처리하는 효율적인 방식으로 다시 짠다. 저자가 아홉 가지 구현을 벤치마크한 결과 파이토치 내장 함수가 가장 빨랐지만, 직접 짠 구현도 컴파일을 거치면 큰 차이가 나지 않았다.

주요 인사이트

• 쿼리·키·값을 세 개의 행렬로 두는 설계는 수학적 필연이 아니라 경험적 선택이다. 저자는 행렬 수를 줄여도 결과가 꽤 괜찮았다는 독자 실험을 언급하면서, 원 논문 구현을 그대로 따르는 이유가 성능이 조금 더 잘 나오기 때문이라고 밝힌다.
• 드롭아웃은 요즘 언어모델 학습에서는 거의 쓰이지 않는다. GPT-2 시절 오픈AI 팀이 썼기 때문에 원래 구조를 재현하는 차원에서 넣은 것이고, 직접 학습할 때는 빼도 된다고 저자가 명시한다.
• 파이토치의 드롭아웃은 하드웨어나 운영체제에 따라 완전히 결정적이지 않아 같은 시드를 써도 결과가 조금 달라질 수 있다. 저자가 이 문제를 파이토치 팀에 제보했지만 드롭아웃 자체가 덜 쓰이다 보니 우선순위가 낮다는 답을 받았다고 한다.
• 멀티헤드를 여러 모듈을 쌓아 순차 호출하는 방식으로 짜면 파이썬 반복문이 병목이 된다. 헤드는 서로 독립적이라 기다릴 이유가 없고, 행렬 곱으로 한 번에 처리하면 GPU는 물론 CPU에서도 유리하다.
• 아홉 가지 구현 벤치마크에서 파이토치의 스케일드 닷프로덕트 어텐션이 플래시 어텐션을 쓸 때 가장 빨랐고, 플렉스 어텐션은 컴파일을 하지 않으면 오히려 매우 느렸다. 직접 짠 구현도 컴파일을 적용하면 내장 구현과 격차가 크게 줄었다.

자주 묻는 질문

관련 AI 소식