LLM 개발 과정 8부 - 학습 가능한 셀프 어텐션 기술

 (gilesthomas.com)
2P by GN⁺ | ★ favorite | 댓글 1개
  • GPT식 디코더 전용 Transformer에서 학습 가능한 셀프 어텐션은 각 토큰이 앞선 입력의 어떤 토큰에 주목해야 할지 계산해 문맥 벡터를 만듦
  • 핵심은 스케일드 닷 프로덕트 어텐션으로, 입력 임베딩을 query, key, value 공간으로 보내는 세 학습 행렬 Wq, Wk, Wv를 사용함
  • 입력 행렬 X는 Q=XWq, K=XWk, V=XWv로 변환되고, Ω=QKᵀ를 √c로 나눈 뒤 행 단위 softmax로 어텐션 가중치 A를 얻음
  • 문맥 벡터는 C=AV라는 한 번의 행렬 곱으로 만들어지며, 전체 계산은 5번의 행렬 곱과 전치 하나로 모든 토큰에 적용 가능함
  • 이 단계는 입력 임베딩끼리 직접 dot product를 하던 장난감 예제를 넘어, PyTorch nn.Module과 nn.Linear로 구현할 수 있는 훈련 가능한 어텐션으로 이어짐

LLM 처리 흐름에서 셀프 어텐션의 위치

  • GPT식 디코더 전용 Transformer 기반 LLM은 지금까지의 토큰을 보고 다음 토큰을 예측하는 구조임
  • 처리 흐름은 문자열을 토큰으로 나누고, 각 토큰을 토큰 임베딩으로 바꾼 뒤, 위치 정보를 나타내는 위치 임베딩을 더해 입력 임베딩을 만드는 순서임
  • 셀프 어텐션은 각 입력 임베딩에 대해 다른 토큰들에 얼마나 주목할지를 나타내는 어텐션 점수 목록을 생성함
    • 예시 문장 "the fat cat sat on the mat"에서 "cat"을 볼 때 "fat"은 중요할 수 있음
    • "mat"을 볼 때는 "fat"의 중요도가 상대적으로 낮을 수 있음
  • 어텐션 점수는 softmax를 거쳐 합이 1인 어텐션 가중치가 되고, 이 가중치로 입력 임베딩들을 가중합해 문맥 벡터를 만듦
  • 문맥 벡터는 각 토큰의 의미를 입력 전체의 문맥 속에서 표현하는 벡터로 다뤄짐

학습 가능한 셀프 어텐션의 목표

  • 이전 단계까지는 입력 임베딩끼리 직접 dot product를 계산하는 장난감 셀프 어텐션을 사용함
  • 이번 단계의 목표는 입력 벡터에서 어텐션 점수를 만들 수 있는 학습 가능한 어텐션 메커니즘을 구성하는 것임
  • Sebastian Raschka의 Build a Large Language Model (from Scratch) 3.4절은 이를 scaled dot product attention으로 구현함
  • 초점은 왜 이 구조가 효과적인지보다, 어떤 계산으로 동작하는지에 맞춰져 있음

Query, Key, Value 행렬과 공간 투영

  • 입력 시퀀스 길이를 n, 입력 임베딩 차원을 d, 문맥 벡터 차원을 c로 둠
  • 입력 임베딩 시퀀스는 x1, x2, x3, ... xn으로 표현되고, 각 입력 임베딩은 d차원 벡터임
  • 세 개의 학습 가능한 가중치 행렬을 정의함
    • query weights matrix: Wq
    • key weights matrix: Wk
    • value weights matrix: Wv
  • 각 행렬은 d×c 크기이며, d차원 입력 벡터를 c차원 공간으로 투영함
  • 입력 벡터 xm을 query 공간으로 보내는 계산은 qm=xmWq임
  • key 공간과 value 공간도 같은 방식으로 입력 임베딩을 각각 다른 c차원 공간으로 투영함

행렬을 투영으로 보는 방식

  • 행렬은 점을 회전시키는 등 기하학적 변환에 사용할 수 있음
  • 정사각 행렬은 같은 차원 안에서 변환을 수행하고, 정사각이 아닌 행렬은 벡터를 다른 차원 공간으로 보낼 수 있음
  • 예를 들어 3×2 행렬은 3차원 벡터들을 2차원 벡터들로 변환할 수 있음
  • 3D 그래픽에서 3D 점을 2D 화면의 점으로 변환하는 frustum 행렬도 이런 투영의 예로 사용됨
  • 셀프 어텐션은 입력 임베딩을 query, key, value라는 세 개의 서로 다른 투영 공간으로 보낸 뒤, 투영된 벡터들로 계산을 진행함
  • 이 투영 행렬들은 훈련 중 학습되므로, 단순 dot product 어텐션에는 없던 간접성이 생김

어텐션 점수 계산

  • 특정 입력 xm을 고려할 때, 다른 입력 xp에 대한 어텐션 점수는 query 투영과 key 투영의 dot product로 정의됨
  • 계산식은 다음과 같음
    • qm=xmWq
    • kp=xpWk
    • ωm,p=qm·kp
  • 모든 입력에 대해 이 계산을 반복문으로 처리할 수도 있지만, 행렬 곱을 쓰면 한 번에 계산할 수 있음
  • 입력 임베딩 전체를 행렬 X로 두면, X는 n×d 크기임
  • key 행렬은 K=XWk로 한 번에 계산됨
    • 결과 K는 n×c 크기
    • 각 행은 해당 입력 임베딩을 key 공간으로 투영한 벡터임
  • query 행렬도 같은 방식으로 Q=XWq로 계산됨
  • 모든 query와 모든 key 사이의 dot product는 QKᵀ로 얻음
    • Q는 n×c
    • Kᵀ는 c×n
    • 결과 Ω는 n×n
  • Ωm,p는 xm의 문맥 벡터를 만들 때 xp에 얼마나 주목할지를 나타내는 어텐션 점수임

스케일링과 softmax 정규화

  • 어텐션 점수는 이전 예제처럼 softmax를 거쳐 합이 1인 가중치로 바뀜
  • softmax는 큰 값은 더 키우고 작은 값은 낮추면서, 목록 전체의 합이 1이 되도록 조정함
  • 실제 LLM에서는 d와 c가 수천 단위가 될 수 있어, 순수 softmax만 쓰면 작은 gradient가 생길 수 있음
  • 이 경우 softmax가 “step function처럼” 동작할 수 있음
    • 가장 큰 값이 지배하고 나머지 값들이 매우 작아지는 상황으로 해석됨
  • 이를 완화하기 위해 어텐션 점수를 투영 공간 차원 c의 제곱근으로 나눈 뒤 softmax를 적용함
  • 행렬 표현은 다음과 같음
    • A=softmax(Ω/√c, axis=1)
  • axis=1은 PyTorch식 표기로, softmax를 행 단위로 적용한다는 뜻임
  • 결과 A는 정규화된 어텐션 점수, 즉 어텐션 가중치 행렬임

문맥 벡터 생성

  • value 공간 투영은 V=XWv로 계산함
  • A는 n×n 크기의 어텐션 가중치 행렬임
    • Am,p는 xm의 문맥 벡터를 만들 때 입력 p에 적용할 어텐션 가중치임
  • V는 n×c 크기이며, 각 행은 입력 임베딩을 value 공간으로 투영한 벡터임
  • 문맥 벡터 행렬은 C=AV로 계산됨
    • 결과 C는 n×c
    • C의 m번째 행은 입력 xm에 대한 문맥 벡터임
  • 이 계산은 각 토큰에 대해 value 벡터들을 어텐션 가중치로 곱해 더하는 작업을 행렬 곱 하나로 수행함

전체 계산 요약

  • 입력 행렬 X는 토큰 시퀀스의 입력 임베딩을 담고 있으며, 크기는 n×d임
  • 세 개의 학습 가능한 행렬로 입력을 각각 query, key, value 공간으로 투영함
    • Q=XWq
    • K=XWk
    • V=XWv
  • query와 key의 dot product로 어텐션 점수를 계산함
    • Ω=QKᵀ
  • 점수를 스케일링한 뒤 행 단위 softmax를 적용해 어텐션 가중치를 만듦
    • A=softmax(Ω/√c, axis=1)
  • value 투영과 어텐션 가중치를 곱해 문맥 벡터를 생성함
    • C=AV
  • 전체 셀프 어텐션 메커니즘은 5번의 행렬 곱과 하나의 전치로 모든 입력 토큰의 문맥 벡터를 만들 수 있음

PyTorch 구현과 다음 단계

  • 책의 3.4절은 위 계산을 PyTorch 코드로 구현하고, 동일한 행렬 연산을 수행하는 간단한 nn.Module 서브클래스를 만듦
  • 첫 번째 버전은 세 개의 가중치 행렬에 일반 nn.Parameter 객체를 사용함
  • 두 번째 버전은 더 효과적인 훈련을 위해 nn.Linear를 사용함
  • 이후 다룰 주제는 두 가지임
    • causal self-attention: 특정 토큰을 볼 때 이후 토큰에는 주목하지 않는 방식
    • multi-head attention: 처음 생각했던 것만큼 복잡하지 않은 주제로 예고됨
  • 배치 처리는 별도 고민거리로 남아 있음
    • 단일 입력 시퀀스에서도 어텐션 점수 행렬을 사용함
    • 여러 입력 시퀀스를 병렬 처리하려면 행렬보다 높은 차수의 텐서가 필요할 수 있음
  • 다음 글은 Writing an LLM from scratch, part 9 -- causal attention로 이어짐

함께 보면 좋은 글 β

GN⁺   [-]
Hacker News 의견들

  • Raschka의 명확한 설명을 여덟 번 다시 읽고서야 이해한 느낌이 들었다는 말처럼, 반복은 새 개념을 내면화하게 만드는 과정으로 보임
    단순 암기라기보다 뉘앙스를 알아차리고, 머릿속에서 여러 각도로 보며, 기존 세계관에 보강·대체·조정 형태로 통합하는 일에 가까움. 운동 능력이 필요한 활동을 연습할 때도 비슷하고, 역할 모델을 따라 하거나 사람·사물·아이디어와 상호작용하는 피드백 고리 속에서 사회에 맞춰 가는 법처럼 덜 명시적으로 배우는 개념도 있음

    • 반복 자체뿐 아니라 같은 개념에 대한 서로 다른 설명을 반복해서 접하는 게 중요함
      막혔을 때 예전에는 전혀 이해되지 않던 다른 설명이 실마리가 되어, 나머지 설명들도 갑자기 더 잘 들어오기 시작할 때가 있음. 중간에 주의 계층을 이해하려면 행렬 곱셈이나 텐서 연산을 알아야 한다는 식으로, 사실 다른 기초 개념이 부족했다는 걸 깨닫기도 함. 각 설명에는 여러 숙련도 단계의 작은 지식 조각이 들어 있고, 현재 수준보다 조금 높은 것만 붙잡을 수 있다는 느낌임. 너무 높으면 아무것도 안 보이고, 너무 낮으면 당연하게 느껴짐
    • 처음 자료를 훑은 지 한참 뒤에도 이런 일이 생김
      대학 졸업 몇 년 뒤 예전에는 도무지 이해되지 않던 전기공학 개념을 다시 보니, 이번에는 갑자기 모든 게 맞아떨어졌음
    • 항상 놀라운 건 여기서 필요한 재료가 수면이라는 점임
      반복도 도움이 되지만, 그 반복 사이에 잠을 넣으면 “와, 이제 알겠다!”가 자주 나옴. 수면이 일종의 램에서 SSD로 옮겨 적는 단계라는 글은 많지만, 실제로 그렇게 잘 작동할 때마다 여전히 신기함
    • 이 문제의 일부는 참고서 저자에게도 있다고 봄
      특정 주제의 입문을 아주 쉽게 썼다고 생각했는데, 한 달 뒤 다시 읽어 보면 논리적 비약, 정의하기 전에 쓰인 전문 용어 같은 것들 때문에 민망해질 때가 많음. 독자가 여러 번 읽어야 뉘앙스를 정리할 수밖에 없지 않게 새 개념을 풀어 쓰는 건 정말 어렵지만, 이건 학습 방식의 문제라기보다 글쓰기 방식의 문제일 수도 있음
    • 대학에서 컴퓨터과학과 전기공학을 섞어 들었는데, 특히 신호 처리가 아주 약했음
      첫 과목은 간신히 통과했지만, 다음 학기에 이전 학기 디지털 신호 처리 이론을 프로그래밍으로 구현하는 “응용 DSP” 같은 수업을 듣고는 반에서 1등을 했음. 그다음 학기 같은 교수의 DSP 2에서도 이론 기초를 다시 돌아보지 않았는데 거의 1등이었음. 처음 채점이 부당했을 수도 있지만, 두 번째 수업에서 실제 응용을 충분히 익히면서 빈틈이 메워지고 전체가 맞아떨어진 것 같기도 함

  • “바닥부터”라고 해도 한편으로는 직접 만들고, 다른 한편으로는
    from fancy_module import magic_functions
    같은 걸 보게 됨. 반쯤 진담으로, 어떤 것을 from scratch라고 부르려면 필요한 지식은 밑바닥부터 쌓아야 한다고 봄. 예를 들어 토크나이저는 직접 쓰고 싶지만, 계산이 일어나게 하는 양자물리 법칙까지 유도하고 싶지는 않음

    • 어디에 시간을 쓸지에 따라 선을 그어야 함
      Build an LLM from scratch는 300쪽이 넘고, 그 페이지들도 매우 빽빽함. 블로그 글은 그중 10쪽도 안 되는 분량만 다루지만, 공정하게 말하면 가장 어려운 부분임. 토크나이저를 바닥부터 깊게 추가하면 100쪽쯤 더 늘고, 쓸 만한 행렬 곱셈까지 넣으면 몇백 쪽, CUDA로 하면 몇천 쪽이 될 수도 있음. 학습용 기울기를 구하는 자동 미분, 최적화 알고리즘까지 더하면 끝이 없음. 결국 어딘가에서 끊어야 하고, 아니면 정말 깨끗한 모래를 구해 칩부터 만들라는 책이 되어야 함. 그런 점에서 tiktoken과 PyTorch를 경계선으로 삼는 건 책을 감당 가능한 크기로 유지하면서, 아래쪽 구조를 충분히 훑어보고 다음에 무엇을 파고들지 정할 수 있게 해 주는 좋은 선택으로 보임
    • 책을 읽어볼까 고민 중이었고 아직 그럴 수도 있지만, 성격상 계속 진행하기 전에 수학 함수가 어떻게 작동하는지 토끼굴로 빠질 것 같음
    • 비슷하게 느껴짐. 바닥부터 이해하고 싶다는 건 PyTorch나 NumPy 없이 가고 싶다는 뜻임
    • 1부와 8부를 봤는데, 글 어디에서도 import 문은 보지 못했음
    • 여기서 말한 fancy_module이 무엇인지 궁금함

  • Bay Area에 살고 이 책을 읽고 싶다면, 지역 도서관 온라인 계정으로 O'Reilly Media에 접근할 수 있고 이 책도 거기에서 이용 가능함

답변달기