파이토치로 디코더 전용 트랜스포머 직접 구현하기: 위치 인코딩부터 마스크 어텐션까지

핵심 메시지

• 챗GPT 계열 모델의 뼈대인 디코더 전용 트랜스포머는 임베딩·위치 인코딩·마스크 셀프 어텐션·완전연결층이라는 네 조각으로 이뤄져 있다.
• 파이토치의 nn.Embedding 이 단어 임베딩을, 미리 계산해 둔 사인·코사인 행렬이 위치 정보를 담당한다.
• 쿼리·키·값은 편향 없는 nn.Linear 세 개로 만들고, 어텐션은 행렬 곱과 소프트맥스 몇 줄로 계산된다.
• 마스크는 상삼각 부분을 사실상 음의 무한대로 채워 앞쪽 토큰이 뒤쪽 토큰을 미리 훔쳐보지 못하게 막는 장치다.
• 학습 전에는 엉뚱한 답을 내던 모델이 30에폭 학습만으로 의도한 답을 정확히 생성한다.

쉽게 이해하기

조시 스타머는 챗GPT의 토대가 되는 디코더 전용 트랜스포머를 파이토치로 처음부터 구현하며 각 줄이 무슨 일을 하는지 설명한다. 예제는 아주 작다. '왓 이즈 스탯퀘스트?'와 '스탯퀘스트 이즈 왓?'이라는 두 질문에 '어썸'이라고 답하게 만드는 것이 전부다. 어휘는 what, is, StatQuest, awesome, EOS 다섯 개뿐이고, nn.Embedding 이 숫자만 받기 때문에 토큰마다 ID를 붙인 사전과 그 역방향 사전을 함께 만든다.

학습 데이터를 만드는 방식이 흥미롭다. 각 토큰이 '다음에 올 토큰'을 만들어내야 하므로, 입력은 프롬프트를 처리하는 구간과 응답을 생성하는 구간에서 모두 나온다. 그래서 입력 텐서는 프롬프트에 EOS와 정답 토큰까지 이어 붙인 형태가 되고, 정답 라벨은 그것을 한 칸씩 밀어놓은 형태가 된다. 이렇게 만든 입력과 라벨을 TensorDataset 과 DataLoader 에 넘기면 학습 데이터 준비가 끝난다.

위치 인코딩은 사인과 코사인 곡선을 번갈아 쓴다. 토큰의 위치를 뜻하는 pos 와 임베딩 값의 순번을 뜻하는 i 를 공식에 넣어 값을 얻는데, 매번 계산하면 낭비이므로 max_len 행과 d_model 열짜리 행렬에 미리 계산해 저장한다. 짝수 열에는 사인 값을, 홀수 열에는 코사인 값을 채우고 register_buffer 로 등록해 GPU 로 함께 옮겨지도록 한다. forward 는 이 행렬을 단어 임베딩에 더하기만 하면 된다.

마스크 셀프 어텐션은 다이어그램 대신 행렬 곱으로 바뀐다. 인코딩된 토큰에 세 개의 가중치 행렬을 곱해 Q, K, V 를 만드는데, 원 논문이 어텐션 계산에 편향을 쓰지 않으므로 nn.Linear 도 bias=False 로 둔다. Q 와 K 의 전치를 곱해 유사도를 구하고 키 차원의 제곱근으로 나눈 뒤, 마스크의 True 자리에 마이너스 10억을 더해 소프트맥스에서 사실상 0이 되게 만든다. 마지막으로 어텐션 비율에 V 를 곱하면 어텐션 값이 나온다.

이 조각들을 모아 라이트닝 모듈을 상속한 디코더 전용 트랜스포머 클래스를 만든다. forward 는 토큰 ID를 임베딩으로 바꾸고 위치 인코딩을 더한 뒤 torch.triu 로 만든 마스크와 함께 어텐션을 계산하고, 잔차 연결을 더해 완전연결층으로 내보낸다. 손실 함수로 쓰는 교차 엔트로피가 소프트맥스를 대신 처리하므로 완전연결층의 출력을 그대로 반환한다. 학습 전 모델은 엉뚱한 답을 내지만, 30에폭을 학습시키면 두 질문 모두에 '어썸 EOS' 를 정확히 생성한다.

주요 인사이트

• 트랜스포머의 어려움은 개념보다 자료구조에 있다. 개념도의 화살표를 행렬 곱으로 옮겨 적는 순간 코드는 몇 줄로 줄어든다.
• 위치 인코딩을 미리 계산해 버퍼로 등록하는 습관은 성능뿐 아니라 GPU 이동 같은 실무 문제까지 함께 해결한다.
• 마스킹을 '음의 무한대 더하기'로 구현하는 이유는 소프트맥스를 거치면 그 자리의 확률이 0에 수렴하기 때문이다. 값을 지우는 대신 확률을 죽이는 방식이다.
• 스케일링(키 차원의 제곱근으로 나누기)은 2017년 원 논문 이후 관행이 됐지만 필수는 아니며, 모델이 클 때 주로 도움이 된다고 설명한다.
• 어휘 다섯 개짜리 장난감 예제라도 전체 파이프라인이 동일하다는 점이 중요하다. 규모만 키우면 그대로 대형 모델이 된다.

자주 묻는 질문

관련 AI 소식