[LLM] How Transformer LLMs Work : Self Attention

이전 강의에서 이해 못해서 머리 아팠던 부분 여기서 다 배운다ㅡㅡ

 

Architectual Overview

대규모 언어모델에서는 입력프롬프트와 출력 텍스트가 있다. 여기서 중요한 건 모델이 토큰을 하나씩 출력한다는 것이다. 

 

LLM 모델에서는 대표적으로 세가지 단계가 있는데 앞서 배운 토크나이저가 먼저 실행된다. 그리고 이 출력은 트랜스포머 블록 스택으로 전달된다. 이 부부분이 대부분의 연산이 일어나는 신경망이다. 그 다음 이 트랜스포머이 출력은 언어 모델링 헤드(LM Head) 신경망으로 들어간다. 

tokenizer은 다음과 같이 토큰 voca가 있고, 모델은 이 각각 토큰에 대한 토큰 임베딩 벡터를 가직고 있다.

그 다음 언어 모델링 헤드엣는 그 다음 토큰이 무엇이어야 하는지를 기반으로 확률계산이 이루어진다. 이 확률의 합은 100퍼여야 한다.

다음 토큰을 선택하는 방법을 디코딩 전략이라고 부른다. 가장 높은 확률을 가진 토큰을 고르는 것은 좋은 전략이고, greedy decoding이라고 한다. 이 경우는 temperature이 0인 경우이다.

반면 Top-p같은 전략도 있는데 이건 1등만 뽑는 것이 아니라, 자연스럽게 들리는 텍스트를 생성하기 위해서 p개의 토큰 중 랜덤하게 고르는 방식이다. 동일한 프롬프트로 다양한 답이 나오는 것도 이런 디코딩 전략과 연관이 있다.

 

트랜스포머의 장점은 병렬로 계산이 가능하다는 점이다. 트랜스포머는 입력 토큰을 모두 병렬로 처리한다. 이 병렬화 덕분에 시간 효율성이 좋아진다. 이 때 여러개의 트랙이 트랜스포머 블록 스택을 통과한다고 생각하면 된다. 여기서 트랙의 갯수는 모델의 컨텍스트 크기와 같다. 예를 들어 모델의 컨텍스트 크기가 100이면 100개의 토큰을 동시에 처리할 수 있다.

또한 토큰이 생성되면 토큰을 포함한 입력이 다시 트랜스포머로 들어간다. 이 과정은 루프이다. 

이 때 첫번째 스텝과 두번째 스텝의 차이가 있는데 이전 토큰들에 대한 계산은 매번 정확히 동일한 계산이기 때문에 이를 캐싱해서 모델의 생성을 가속할 수 있다. 이것을 kv 캐싱이라고 한다. 

 

또한 효율성과 관련된 지표가 있는데 Time-to-First-Token이라는 지표가 있다. 첫 토큰까지 걸리는 시간은 모델이 이 모든 초기 처리를 끝낼 때까지의 시간을 말한다.

 

The Transformer Block

예를 들어 the shawshank라는 입력이 있을 때 이 stack을 통과하는 track는 두가지이다. 이들의 임베딩 벡터는 transformer block으로 들어간다. 이 블록들은 이 임베딩 값에 여러가지 처리를 한다.

 

두번째 블록은 첫번째 블록의 출력을 입력으로 받아 여러 track에 대해서 병렬로 연산을 한다. 이러한 것을 반복하면 마지막 tocken에 대한 vactor가 LMHead에 전달된다.

 

Transformer block에는 self attention과 feed forward neural network 두가지 구성요소로 이루어져 있다.

 

feed forward neural network에 대해서 알아보자. 만약 self attention이 없었다면 이 모델은 단순히 shawshank 다음에 올 가능성이 높은  redemption 토큰을 예측할 것이다. 왜냐하면 인터넷 등에서는 shawshank 다음에 redemption이 자주 등장하기 때문이다. 이와 같이 이 단어 다음에 어떤 단어가 나올 확률이 높을까? 와 같은 통계적 패턴 정보를 저장하고 학습한다.

feed forward neural network는 입력층이 확장되고, 그 다음에 또 다른 층으로 축소되며 마지막 출력층으로 이어지는 구조이다. 

이 dense layer들이 모델이 알고있는 지식과 정보를 저장하고, 해석하고, 조합한다.

 

그런데 self attention은 모델이 이전 token들을 참고하고 그 문맥을 현재 token의 이해에 통합할 수 있게 해준다.

만약 위 문장에서 it라는 단어를 처리할 때 그 it이 dog를 의미하는지 llama를 의미하는지 판단해야한다. 이를 해주는 것이 self attention이다.

이건 NLP 분야에서 Coreference Resolution이라고 불리는 과제이다. 

self attention은 여러 블록을 거치면서 임베딩 값이 다른 위치로부터 가져온 정보를 포함한 형태로 변환된다.

크게 relevance scoring(특정 토큰과 얼마나 관련있는지)와 information combination(관련 있는 토큰의 정보 합치기)를 한다.

 

Self-Attention

self attenton은 attention head에서 일어난다. 현재 우리가 보는 토큰과 이 토큰에 앞서있는 sequence의 다른 위치들이 있다.

이들은 현재 토큰보다 앞선 다른 token들의 vactor representations이다. 

self attention은 query projection matrix,key projection matrix,value projection matrix로 이루어진다.

이 weight matrices를 이용해서 query, key, value 매트릭스를 계산한다.

쿼리에는 현재 위치를 나타내는 쿼리 백터가 있고, key, values에도 각각 토큰마다(row마다) 각각의 값을 가지고있다.

 

이건 현재 표현중인 토큰이 다른 토큰과 얼마나 관련있는지를 계산하는 과정이다. 쿼리와 키를 곱하면 relevance 점수를 얻게 되고, 이는 enriched vactor에 반영된다. 

이후 관련있는 토큰들로부터 정보를 결합한다.

앞서 키로 구한 점수를 values와 곱해 나온 값을 모두 더함으로서 information combination의 출력을 구할 수 있다. 

이 계산은 attention head에서 일어난다. 하지만, self attention에서는 여러 attention heads에서 병렬로 일어난다.

 

각 head가 벡터들에 할당하는 attention은 다르다. 각각 자신만의 key,queries,values weight matrices를 가지기 때문이다.

 

먼저 정보를 attention heads으로 분할한 후 다시 결합하여 최종출력을 형성한다.

하지만 이 방법은 연산이 너무 많이 수행되기 때문에 최근에는 Multi-Query Attention을 사용한다. 각 attention head가 같은 key와 value를 갖는 것이다.(쿼리만 다름) 이렇게 하면 계산해야할 파라미터 수가 줄어 효율적이다.

 

또 이런 Grouped-Qurey Attention이 있는데, 이는 key value를 여러세트 쓰되, 그룹 단위로 같은 세트를 쓰는 것이다. 이는 단일세트보다 더 나은 결과를 가져온다.

 

그래서 앞으로 Multi-Query Attention을 쓴다면 attention heads 수 뿐만 아니라 key, value heads 수도 함께 봐야한다.

 

최근 attention 효율을 높이기 위해서 Sparse Attention이라는 기법이 나왔다.

이제까지는 

full self attention을 사용해 모든 토큰을 봤다. 하지만 모델이 커질수록 모든 토큰을 보는 것은 비효율적이다.

그래서 모든 레이어(블럭)에서 모든 토큰을 보게 하는 것이 아니라 특정 레이어에서만 모든 토큰을 보게 하는 것이다.

이런식으로 

sparse attention은 위와 같은 방식으로 진행된다. a는 모든 토큰을 보는 것이고, b와 같이 ~만큼 떨어진 토큰을 띄엄띄엄 볼수도 있으며 c와 같이 특정 위치만 추가로 볼 수 있게 할수도 있다.

 

자 이제 논문에서 각각이 뭘 의미하는지 알 수 있게 되었다.

layers는 블록의 수, model dimension은 임베딩 크기(파라미터 수), ffn Dimension은 FFNN에서의 units 수, attention headers는 전체 attention headers이고, key/value heads는 key value 쌍의 개수, 그리고 vocabulary size는 토큰을 가짓수를 의미한다.