학습 없이 손으로 만든 트랜스포머로 간단한 시퀀스 예측하기

 (vgel.me)
2P by GN⁺ | ★ favorite | 댓글 1개
  • GPT-2와 비슷한 디코더 전용 트랜스포머를 학습 없이 직접 가중치만 설계해 (aab)* 패턴을 예측하게 만든 실험임
  • "aabaabaabaab..." 예측은 직전 두 토큰을 봐야 하므로, 단순 교대 패턴보다 어텐션 동작을 드러내기 좋음
  • 모델은 N_CTX=5, N_VOCAB=2, N_EMBED=8로 작게 구성하고, a=0, b=1 토큰화와 원핫 임베딩을 사용함
  • 단일 어텐션 헤드는 최근 두 토큰에 0.5씩 주의를 배분하고, a=1, b=-1 인코딩의 상쇄를 이용해 다음 토큰을 계산함
  • 모호하지 않은 문맥에서는 100.0% (27/27) 정확도를 보였지만, 5토큰 문맥 한 번 예측에 약 4,000 FLOPs가 필요해 직접 구현한 규칙보다 훨씬 비효율적임

학습 없이 가중치를 직접 지정한 미니 GPT-2

  • 목표는 트랜스포머와 어텐션의 각 구성요소가 실제로 무엇을 하는지 직관적으로 이해하는 것임
  • 모델은 학습하지 않고 사전학습 가중치도 쓰지 않으며, 저녁 한 번에 각 가중치를 직접 지정하는 방식으로 만들어짐
  • 구조는 GPT-2와 비슷한 디코더 전용 트랜스포머이며, 구현은 jaymody의 picoGPT 구현을 바탕으로 단순화됨
    • layer norm 제거
    • multi-head attention 대신 단일 헤드 사용
    • transformer block의 mlp feed-forward layer 제거

(aab)* 시퀀스를 고른 이유

  • 처음에는 "ababababab" 같은 시퀀스를 예측하려 했지만, transformer가 shifted sequence를 예측하기 때문에 너무 쉬운 작업이 됨
    • a이면 b, 아니면 a를 예측하면 되어 위치 임베딩을 쓸 필요가 없음
  • 최종 작업은 "aabaabaabaab...", 즉 (aab)* 시퀀스 예측으로 정해짐
    • 이전 두 토큰이 ab 또는 ba이면 다음 토큰은 a
    • 이전 두 토큰이 aa이면 다음 토큰은 b
    • bb는 작업 범위 밖의 경우로 취급됨
  • 토큰화는 두 기호만 다루는 단순한 방식임
    • a는 0
    • b는 1

모델 차원과 계산 흐름

  • 선택한 모델 파라미터는 세 가지임
    • N_CTX = 5: 모델이 한 번에 보는 최대 문맥 길이
    • N_VOCAB = 2: a, b 두 토큰
    • N_EMBED = 8: 토큰·위치·계산용 공간을 담는 임베딩 크기
  • 실제 작업은 이전 두 토큰만 필요하지만, N_CTX=5로 잡아 관련 없는 토큰을 무시해야 하는 상황을 포함함
  • gpt 함수는 다음 순서로 동작함
    • wte[inputs] + wpe[range(len(inputs))]로 토큰 임베딩과 위치 임베딩을 더함
    • 하나의 transformer block을 통과함
    • 마지막에 x @ wte.T로 어휘 공간의 logits를 만듦

위치와 토큰을 원핫으로 담는 임베딩

  • wpe는 위치 임베딩이며, 앞의 5개 임베딩 차원을 위치 원핫으로 사용함
    • position 0은 [1, 0, 0, 0, 0]
    • position 4는 [0, 0, 0, 0, 1]
  • wte는 토큰 임베딩이며, 다음 2개 차원을 토큰 원핫으로 사용함
    • token a는 해당 토큰 차원에서 [1, 0]
    • token b는 [0, 1]
  • 8번째 임베딩 위치는 처음에는 쓰지 않고, transformer block 안에서 scratch space로 사용함
  • 예를 들어 "aabaa"는 5 x 8 임베딩 행렬로 표현되며, 각 행은 위치 원핫과 토큰 원핫을 함께 포함함

어텐션 헤드가 최근 두 토큰을 고르는 방식

  • transformer block은 하나의 attention head와 attention 결과를 다시 임베딩 공간으로 돌리는 c_proj 선형층으로 구성됨
  • c_attn은 embed_size x (embed_size * 3) 크기의 선형층이며, 입력 임베딩을 qkv 행렬로 바꾼 뒤 q, k, v로 나눔
  • k는 위치 임베딩을 분리해 각 토큰이 가진 위치 정보를 나타냄
  • q는 각 위치가 찾고 싶은 위치 범위를 나타내며, q @ k.T를 통해 attention score 행렬을 만듦
  • softmax와 causal mask 이후 attention 행렬은 다음 성질을 가짐
    • 첫 행은 첫 토큰에만 100% attention
    • 이후 행들은 접근 가능한 최신 두 토큰에 각각 0.5씩 attention
  • causal mask는 미래 토큰 위치에 매우 작은 값, 실제 코드에서는 -1e10 같은 값을 더해 미래 토큰을 보지 못하게 함
    • 이 손으로 만든 모델은 미래를 훔쳐보도록 설계되지 않았지만, GPT-2 구조에 가깝게 유지하기 위해 mask를 남김
  • np.sqrt(q.shape[-1])로 나누는 scaling은 실제 학습에서 기울기 개선에 도움이 되지만, 이 수제 모델에는 영향을 주지 않음

v 인코딩과 더하기 상쇄로 만든 예측

  • v는 토큰 원핫을 a=1, b=-1 인코딩으로 바꿈
  • attention 결과가 최근 두 토큰을 0.5씩 평균내기 때문에, 이 인코딩은 다음 규칙을 계산함
    • a, b → 0.5 * 1 + 0.5 * (-1) = 0
    • b, a → 0.5 * (-1) + 0.5 * 1 = 0
    • a, a → 0.5 * 1 + 0.5 * 1 = 1
  • 결과적으로 행의 7번째 위치에는 다음 값이 생김
    • 0이면 a를 예측해야 하는 경우
    • 1이면 b를 예측해야 하는 경우
  • 입력 "aabaa"에서 첫 행은 정보가 부족해 b 예측이 생길 수 있지만, 이후 예측은 (aab)* 규칙과 맞음

예측값을 다시 어휘 공간으로 보내기

  • c_proj는 attention 결과의 7번째 위치 값을 토큰 원핫 형식으로 되돌림
  • 단순히 [..., 1, 0, ...] 또는 [..., 0, 1, ...]를 만들지 않고, 1024로 스케일된 원핫을 만듦
    • embedding[row, 5] = 1024 + (-1024) * prediction
    • embedding[row, 6] = 0 + 1024 * prediction
  • transformer block에는 residual connection이 있어 x = x + causal_self_attention(...)로 원래 임베딩이 더해짐
  • residual signal이 불필요하게 남기 때문에, 1024 스케일을 사용해 이 신호를 압도함
  • 마지막으로 x @ wte.T를 계산해 logits를 만들고 softmax를 적용함
    • "aabaa" 문맥에서는 최종 예측 행이 b를 가리킴
    • 학습 때는 모든 행의 예측이 유용하지만, 추론에서는 마지막 행만 필요함

생성 결과와 정확도

  • complete 함수는 마지막 최대 5토큰을 모델에 넣고, softmax 결과의 마지막 행에서 argmax로 다음 토큰을 선택함
  • 예시 생성 결과는 다음과 같음
    • complete("a") → a :: baabaabaab
    • complete("ba") → ba :: abaabaabaa
    • complete("abaab") → abaab :: aabaabaaba
  • 범위 밖 입력에서도 반복 패턴으로 회복하는 경우가 있음
    • complete("ababa") → ababa :: abaabaabaa
    • complete("bbbbb") → bbbbb :: aabaabaaba
  • "aab" * 10 테스트에서 모호하지 않은 문맥만 평가하면 정확도는 100.0% (27/27) 임

4,000 FLOPs와 8개 명령의 차이

  • 5토큰 문맥 전체를 사용할 때, 이 모델은 단일 토큰 예측에 약 4,000 floating point operations가 필요함
    • 대부분은 attention 계산에서 사용됨
    • context window 축소, fused multiply-add, kv caching 등으로 줄일 수 있음
    • 그래도 단일 토큰 예측에는 수백 개의 머신 명령이 필요함
  • 같은 (aab)* 규칙을 직접 작성한 x64 어셈블리는 8개 명령으로 다음 토큰을 계산함
  • 자연어 생성에서 현재 모델보다 1000배 더 효율적인 언어 모델을 학습할 수 있는지에 대한 질문이 남음

함께 보면 좋은 글 β

GN⁺   [-]
Hacker News 의견들

  • 관련 작업으로 "Thinking Like Transformers"가 있음
    RASP라는 원시 프로그래밍 언어를 소개하는데, Transformer 구성 요소로 모델링 가능한 연산들로 이루어져 있고 히스토그램이나 정렬 같은 프로그램을 작성할 수 있음을 보여줌
    Sasha Rush와 Gail Weiss의 훌륭한 블로그 글도 있으며, 후속 연구에서는 RASP류 프로그램을 학습 없이 실제 모델 가중치로 컴파일할 수 있음을 보였음
    [1] https://arxiv.org/abs/2106.06981
    [2] https://srush.github.io/raspy/
    [3] https://arxiv.org/abs/2301.05062

    • RASP 계열을 정말 좋아함
      이 분야가 재미있다면, Transformer 모델의 가중치를 직접 골라 사람이 초등학교에서 배우는 방식과 비슷하게 긴 덧셈을 하도록 만든 내 HandCrafted Transformers 작업도 한번 볼 만함
      [1] https://colab.research.google.com/github/newhouseb/handcraft...
    • 이런 작업에는 Haskell 같은 함수형 언어가 잘 맞을 것 같음
      신경망에서 코드로 가는 방향도 설명 가능성 측면에서 아주 흥미로울 듯함

  • Transformer를 잘 이해한다고 생각했지만 직접 구현해 본 적은 없었음
    어느 날 직접 구현해 보니 표준 PyTorch Transformer만큼 잘 동작하거나 학습되지 않았고, 결국 dropout을 무시한 게 원인임을 깨달음
    숫자 덧셈을 학습시켰고 같은 쌍을 두 번 보여준 적이 없어서 과적합은 불가능하다고 생각했는데, dropout의 역할이 생각보다 훨씬 컸음
    요약하면 그냥 Transformer를 직접 구현해 보는 게 좋고, 밑바닥부터 할수록 더 좋음
    해 본 사람들은 모두 예상 못 한 걸 배웠고, 토큰 단위 학습 병렬화부터 역전파가 실제로 어떻게 동작하는지까지 사람마다 깨닫는 부분이 달랐음

    • 이 작업에 접근하는 데 도움이 될 만한 참고 문헌이 있을지 궁금함

  • Karpathy 쪽 자료도 좋아하지만, 내게 Transformer가 드디어 이해되게 만든 건 이 영상이었음: https://youtu.be/kWLed8o5M2Y?si=SJT5_lCJ0hSR7Z_k

  • 비슷한 생각을 한동안 해 왔음
    도메인 전문가가 손으로 조정해서 학습을 가속할 수 있는, 모델 가중치에 대한 직관적 인터페이스를 만들 수는 없을까?
    예를 들어 시각 모델에서 교통 콘을 감지할 때 "주황색다움"에 해당하는 가중치 묶음을 올리는 식임
    그러면 "주황색다움"을 제대로 보정하기 위해 수천·수백만 개의 예시를 더 요구하는 대신 사람이 가속할 수 있음
    물론 어려운 점은 이 인터페이스가 서로 다른 의미를 갖는 가중치 묶음에 매핑되어야 한다는 건데, 기술적으로 불가능한 이유가 있는지 궁금함

    • "도메인 전문가가 모델의 가중치를 손으로 조정한다"는 건 딥러닝 이전의 이미지 인식 방식과 비슷하게 들림
      [1] https://www.youtube.com/watch?v=8SF_h3xF3cE&t=1358s
    • 찾는 이유는 The Bitter Lesson이라고 불림
      짧게 말하면, AI에 사람이 도와주는 방식은 더 많은 컴퓨팅 파워로 돌리는 것보다 거의 항상 비용 효율이 낮음
      사람이 주황색 교통 콘을 감지하도록 가중치 층을 보정하는 동안, GPU 클러스터는 교통 콘, 신호등, 나무, 다른 차, 약간 다른 주황색의 교통 콘까지 감지하도록 AI를 학습시켜 버림
    • 층과 가중치 수가 사람이 수동으로 갱신할 수 있는 규모가 아니고, 설령 가능하더라도 가중치를 바꿨을 때의 하류 효과를 관리하기가 너무 어려움
      주황색을 더 잘 보도록 그림을 조정한다고 해도, 동시에 다른 모든 색의 정확성을 감시할 수 없다면 자신도 모르게 다른 색에서 문제를 만들 가능성이 큼
    • 기술적으로 불가능하거나 매우 어려운 이유는 가중치가 보통 해석하기 매우 어렵기 때문임
      특정 뉴런 클러스터가 특정 개념에 대응하는 식이 아니라, 모든 것이 대체로 모든 일을 조금씩 함
    • Transformer의 어텐션 메커니즘은 사람이 이해할 수 있는 의미론에 쉽게 매핑되지 않는 것 같음
      관여하는 매개변수가 너무 많음

  • Transformer 논문은 너무 기술적이라 늘 얕게라도 이해하고 싶었지만 어려웠음
    이 글은 작동 방식을 이해하는 데 정말 도움이 됐고, 적어도 예시는 아주 명확했음
    덕분에 대학 때 배운 행렬도 다시 떠올릴 수 있었음

  • Turing 기계나 정규식을 파싱하는 기계처럼 일종의 추상 기계 아닌가?

    • 조금 단순화하면, 입력 집합을 다음 출력의 확률 집합으로 매핑하는 "기계"임
      먼저 토큰 목록을 정의하고, 예를 들어 쉽게 24개 글자라고 해 보자
      이 기계는 토큰 입력 시퀀스를 받아 결정론적인 행렬 연산들을 수행한 뒤, 모든 토큰의 확률 목록을 출력함
      "학습"은 그 연산에 쓰이는 행렬 안의 숫자 일부를 설정하는 과정일 뿐임
      최종 코드에는 if 문이 하나뿐이고, 그것도 결과의 정확도를 평가하기 위한 것임에 주목할 만함
      모든 "논리"는 행렬 연산의 결과에서 나옴
    • 이런 것들을 보통 떠올리는 의미의 오토마타로 해석하기는 꽤 어려움
      신경망에서는 모든 것이 대체로 약간 흐릿하고, if/else 같은 것은 거의 없지만 Transformer 예시처럼 값을 0이나 -∞로 "마스킹"하는 경우는 있음
      출력도 거의 항상 점수나 확률의 묶음이라서, 고양이와 개 사진을 구분하는 모델이 dog:0.95 cat:0.05 같은 결과를 내면 개 점수가 더 높기 때문에 개를 예측했다고 말함
      Transformer의 핵심인 어텐션 메커니즘은 일종의 부드러운 조회 연산에 기반함
      흐릿하지 않은 시스템이라면 시퀀스의 각 토큰을 순회하며 현재 토큰과 관련 있는지 확인하고 관련 있으면 어떤 동작을 하겠지만, Transformer에서는 관련성이 이진 결정이 아님
      대신 시퀀스의 모든 토큰 쌍 사이에 연속적인 관련성 점수를 계산하고, 그 점수를 이용해 다음 동작을 함
      다만 어떤 것들은 이진 결정 기반 시스템에서 직접 일반화하기 쉽지 않음
      예를 들어 그런 관련성 점수는 어휘 토큰들에 대한 가중 평균을 계산하는 가중치로 쓰이고, 현재 위치에 대한 "평균 토큰"을 얻음
      이걸 분기 논리 기반 과정의 확장으로 쉽게 해석할 방법은 없어 보임
    • 선형대수 더미에 AllSpark가 닿은 것 아닌가?
    • 맞음
      Linear Transformers가 사실은 Fast Weight Programmers라는 걸 설명하는 이 논문을 보면 좋음: https://arxiv.org/abs/2102.11174
    • 신경망은 Turing 기계임
      가중치를 신중하게 설정하면 어떤 계산이든 수행하게 만들 수 있음
      다만 근사에 기반하지 않는 컴파일러가 있으면 좋겠음

  • "직접 모델을 만들어 보고 싶어질 수도 있다"는 말은 호기심을 만족시키는 학습 연습 말고는 어디에 쓰려는 건지 궁금함
    복잡한 머신러닝 모델은 집에서 블로그 읽는 사람이 다루기엔 비현실적이라는 느낌이 들기 시작함

    • nanoGPT에서는 Shakespeare로 모델을 사전학습시키면 3분 만에 원본 자료에 대해 Lewis Carroll의 Jabberwocky 수준 충실도에 도달함
      그럴듯해 보이는 옛 영어 단어를 많이 만들어 내고, 영어 문법의 기초와 희곡의 형식 등을 배움
      그렇게 짧은 시간에 그 정도까지 가는 게 꽤 놀라웠음
      로컬에서 여러 모델을 Shakespeare-from-Wish.com 수준의 충실도로 학습해 보면, 좋은 아키텍처를 찾았는지와 확장해 볼 때가 됐는지 판단하는 데 도움이 될 것 같음
    • 글 첫 문단에 목적이 나와 있음
      Transformer와 어텐션을 더 잘 이해하고 싶었고, The Illustrated Transformer를 읽었지만 어텐션의 여러 조각이 실제로 무엇을 하는지 직관적으로 와닿지 않았다는 내용임
      q와 k의 차이, 그리고 v는 더 말할 것도 없었다는 식임
    • 훌륭한 학습 연습임
      단순히 호기심을 만족시키는 것을 넘어 이해를 만들고 깊게 하는 데 도움이 됨
    • 그냥 이런 프로젝트를 해킹하는 걸 실제로 즐길 수도 있지 않을까? 이상하긴 하지만

  • 제목에 neural network 같은 표현을 넣을 수 있으면 좋겠음
    이건 두 회로를 전자기적으로 결합하는 코일 묶음이 아니라, 머신러닝의 "Transformer" 아키텍처와 관련된 내용임

답변달기