이전 게시글 : 

이전 주제들은 현재는 공통 과학처럼 일반적인 내용이여서 한꺼번에 다뤘지만, 이번에는 호흡을 조금은 길게 가져가보고자 한다.

최근 ViT(Vision Transformer) 관련 연구를 2건이나 했지만 아직도 트랜스포머에 대해 잘안다고 대답하지는 못하기 때문이다.

(그래서 이런 시리즈 포스팅을 시작한 이유도 있다.)

Architecture

모델 아키텍처는 상태를 내부에 저장하지 않는(stateless) 함수로 구현되어 있다.

이 함수는 다음을 입력으로 받는다.

• 현재 토큰 ID
• 현재 위치(position) 정보
• 모델의 파라미터들
• 이전 위치들에서 계산되어 캐시된 key/value 값들

그리고 시퀀스에서 다음에 어떤 토큰이 올지에 대한 점수(logits)를 반환한다.

microgpt 의 구현은 GPT-2 구조를 따르되, 몇 가지 단순화를 적용했다.

• LayerNorm 대신 RMSNorm
• bias 항 제거
• GeLU 대신 ReLU 활성화 함수 사용

이제 본격적인 모델 정의에 앞서, 세 개의 간단한 보조(helper) 함수를 먼저 살펴보자. (linear, softmax, rmsnorm)

linear

def linear(x, w):
    return [sum(wi * xi for wi, xi in zip(wo, x)) for wo in w]

linear는 행렬-벡터 곱(matrix-vector multiplication)이다.

입력 벡터 x와 가중치 행렬 w를 받아, w의 각 행(row)과 x의 내적(dot product)을 계산한다. (유사도)

이는 신경망의 가장 기본적인 연산이다.

선형 변환(linear transformation)은 신경망이 학습하는 핵심 연산 단위다.

모든 복잡한 모델도 결국 선형 연산의 반복 위에 쌓여 있다.

softmax

def softmax(logits):
    max_val = max(val.data for val in logits)
    exps = [(val - max_val).exp() for val in logits]
    total = sum(exps)
    return [e / total for e in exps]

softmax는 로짓(logits)을 확률 분포로 변환한다.

로짓은 음의 무한대부터 양의 무한대까지 아무 값이나 가질 수 있다.
softmax를 통과하면,  모든 값이 0과 1 사이로 변환되고 전체 합이 1이 된다. 즉, 확률 분포가 된다.

위 코드에서 최대값(max_val)을 따로 저장해두는 이유는 수치적 안정성(numerical stability) 때문이다.
지수 함수(exp)는 값이 조금만 커도 폭발적으로 증가하기 때문에, overflow를 방지하기 위해서 최대값을 빼준다.

수학적으로는 결과가 동일하지만 계산 과정이 훨씬 안정적이다.

rmsnorm (Root Mean Square Normalization)

def rmsnorm(x):
    ms = sum(xi * xi for xi in x) / len(x)
    scale = (ms + 1e-5) ** -0.5
    return [xi * scale for xi in x]

rmsnorm은 벡터의 크기를 정규화(normalize)하는 함수다.

입력 벡터의 평균 제곱(mean square)을 계산한 뒤, 그 값의 역제곱근으로 스케일링한다.

이 과정은 활성값이 지나치게 커지거나 작아지는 것을 방지하기 위함이다.

원래 GPT-2에서는 LayerNorm을 사용하지만, microGPT는 더 단순한 RMSNorm을 사용한다.

둘의 핵심 목적은 같다 : 네트워크 내부의 값이 폭주(explosion)하지 않도록 안정화한다.

이제 모델 함수(gpt) 자체를 살펴 보자. [:3]

def gpt(token_id, pos_id, keys, values):
    tok_emb = state_dict['wte'][token_id]         # token embedding
    pos_emb = state_dict['wpe'][pos_id]           # position embedding
    ...

아래 항목들을 입력으로 특정 시점의 하나의 토큰을 처리한다. (return logits)

• token_id : 현재 토큰의 ID
• pos_id : 현재 위치 ID
• keys, values : 이전 위치들에서 저장해 둔 key/value 캐시, KV cache
• state_dict : 그리고 전역적으로 정의된 모델 파라미터들

출력은 현재까지의 문맥을 고려했을 때, 다음에 어떤 토큰이 올지에 대한 점수(logits)이다.

이 함수의 진행 순서대로 살펴보자.

Embeddings

tok_emb = state_dict['wte'][token_id]
pos_emb = state_dict['wpe'][pos_id]
x = [t + p for t, p in zip(tok_emb, pos_emb)]
x = rmsnorm(x)

모델은 단순한 정수 ID를 처리할 수 없기에 먼저 임베딩(embedding)을 수행한다.

• state_dict['wte']는 토큰 임베딩 테이블
• state_dict['wpe']는 위치 임베딩 테이블

각 토큰은 하나의 벡터로 변환되고, 여기에 위치 정보 벡터를 더한다.

이렇게 하면 모델은 “무슨 단어인가?”, “어디에 위치했는가?” 를 동시에 표현할 수 있다.

그 다음 rmsnorm을 통해 값을 정규화한다.

(단, 최신 LLM은 일반적으로 위치 임베딩을 생략하고 RoPE와 같은 다른 상대 기반 위치 지정 방식을 도입한다.)

Attention Block

q = linear(x, state_dict[f'layer{li}.attn_wq'])
k = linear(x, state_dict[f'layer{li}.attn_wk'])
v = linear(x, state_dict[f'layer{li}.attn_wv'])

현재 토큰은 세 개의 벡터로 변환된다 : Q (Query), K (Key), V (Value)

직관적으로 보면,

• Query → “나는 무엇을 찾고 있는가?”
• Key → “나는 무엇을 담고 있는가?”
• Value → “선택되면 무엇을 전달하는가?”

이전 위치의 key/value는 keys, values에 저장되어 있다. (KV cache)

각 attention head는 다음을 계산한다.

1. 현재 Query와 이전 Key들의 내적
2. softmax를 적용해 attention 가중치 계산
3. 가중치를 이용해 Value들의 가중합 계산

이 과정을 통해, 현재 토큰이 과거 토큰들 중 어디에 얼마나 집중할지를 결정한다.

MLP block

Attention 이후에는 MLP(Multi-Layer Perceptron)가 따라온다.

x = linear(x, state_dict[f'layer{li}.mlp_fc1'])
x = [xi.relu() for xi in x]
x = linear(x, state_dict[f'layer{li}.mlp_fc2'])

MLP는 단순한 두 층짜리 신경망이다.

• 차원을 4배로 확장
• ReLU 활성화
• 다시 원래 차원으로 축소

Attention이 “토큰 간 소통”이라면, MLP는 “각 토큰이 개별적으로 계산하는 부분”이다.

Residual Connection

각 블록의 출력은 입력에 더해진다.

x = [a + b for a, b in zip(x, x_residual)]

이러한 방법은 깊은 네트워크 학습을 안정화하고 gradient 흐름을 개선하는 효과가 있다.

Output

마지막 hidden state를 어휘 크기(vocab size) 차원으로 투영한다. 어휘의 각 토큰당 하나의 로짓 값을 생성한다.

우리 예시에서는 총 27개의 숫자가 사용됩니다. 

마지막 hidden state를 어휘 크기(vocab size) 차원으로 투영한다. (현재 프로젝트에서 27개 숫자)
각 숫자는 해당 토큰이 다음에 올 가능성에 대한 점수(logit)로, 이 값에 softmax를 적용하면 확률 분포가 된다.

KV cache에 대한 보충 설명

흥미로운 점은, microGPT는 학습 과정에서도 KV cache를 사용한다는 것이다.

보통 KV cache는 추론(inference)에서만 사용하는 것으로 알려져 있다.
그래서 이 부분이 다소 낯설게 느껴질 수 있다.

하지만, 개념적으로 보면 KV cache는 학습 중에도 항상 존재한다.

실제 서빙 환경에서는 attention 연산이 고도로 벡터화(vectorized)되어 있기 때문에, 시퀀스의 모든 위치를 동시에 처리한다.

따라서 KV cache가 명시적으로 보이지 않을 뿐, 계산 내부에서는 동일한 정보가 사용되고 있다.

microGPT는 한 번에 하나의 토큰만 처리한다는 점을 주목해야한다.

• 배치 차원 없음
• 시간 축 병렬 처리 없음

그래서 이전 위치의 key와 value를 직접 리스트에 저장하는 방식으로 KV cache를 명시적으로 구성한다.

keys[li].append(k)
values[li].append(v)

또 하나 중요한 차이가 있다.

일반적인 추론 환경에서는 KV cache에 저장된 텐서들은 그래디언트 계산에서 분리(detached)되어 있다.
즉, 역전파가 그 안을 통과하지 않는다.

하지만, microGPT에서는 다르다.

여기서 캐시된 key와 value는 단순한 숫자 배열이 아니라 계산 그래프에 연결된 Value 객체들이다.

따라서 역전파(backpropagation)가 이 캐시를 통과해 이전 토큰의 계산까지 거슬러 올라간다.

즉, microGPT의 KV cache는 단순한 속도 최적화용 저장소가 아니라, 완전히 살아 있는 계산 그래프의 일부다.