LLM Research2026.10.106 min read

주간 리뷰 (Week 6): Transformer를 백지에서 다시 쓰면 어디서 막히는지 확인한다

이번 주에 배운 Transformer의 부품(멀티헤드 어텐션, FFN, LayerNorm, residual)을 아무것도 보지 않고 다시 조립해 보는 날입니다. 부품을 하나씩 작성한 뒤 GPT 모델로 묶고, 생성 파라미터(temperature, top-k 등)를 바꿔 가며 출력이 어떻게 달라지는지 재현합니다. 막히는 곳이 곧 아직 체화되지 않은 부분이므로,

By 송민성

오늘의 목표

Transformer 전체를 "읽어서 아는 것"에서 "빈 파일에서 직접 쓸 수 있는 것"으로 바꾼다. 새로 배우는 내용은 없다. Transformer Block(MHA + FFN + LayerNorm + residual)과 이를 쌓은 GPT를 참고 코드 없이 다시 작성하고, 막힌 곳을 기록해서 다음 주 이전에 메운다.

개념과 直감

오늘은 새 개념이 없으므로, 이번 주 내용을 한 장의 지도로 정리한다. 백지에서 쓰려면 이 지도가 머릿속에 있어야 한다.

한 줄 요약: GPT는 "토큰 임베딩 + 위치 임베딩 → Transformer Block을 N번 반복 → 마지막 LayerNorm → 어휘 크기만큼의 점수(logits)를 내는 선형층"이다.

Transformer Block 하나는 두 단계다. 회의에 비유하면 이렇다.

어텐션(attention, 각 토큰이 앞쪽 토큰들 중 무엇을 얼마나 참고할지 정해 정보를 섞는 연산): 참석자끼리 서로 발언을 듣는 시간
FFN(feed-forward network, 각 토큰을 따로따로 2층 신경망에 통과시키는 부분): 각자 자리로 돌아가 들은 내용을 혼자 정리하는 시간

그리고 각 단계를 감싸는 장치가 있다.

residual(잔차 연결, 입력을 출력에 그대로 더해 주는 지름길): 원래 노트를 지우지 않고 그 위에 수정 사항만 덧붙이는 것
LayerNorm(층 정규화, 한 토큰의 벡터를 평균 0, 분산 1 근처로 맞추는 것): 발표 전에 모두의 목소리 크기를 맞추는 것

수식은 이것만 기억한다.

text

attention(Q, K, V) = softmax(Q K^T / sqrt(d_head) + mask) V
x = x + MHA(LN(x))
x = x + FFN(LN(x))

두 번째, 세 번째 줄은 LayerNorm을 서브층 앞에 두는 pre-LN 방식이다. 원 논문은 뒤에 두는 post-LN 방식이라 순서가 다르다. 오늘은 자신이 이번 주에 구현한 방식 하나로 통일한다.

shape 지도. 백지 구현의 절반은 shape 추적이다.

입력 토큰: (B, T) — B는 배치 크기, T는 시퀀스 길이
임베딩 후: (B, T, C) — C는 임베딩 차원
헤드 분리 후: (B, n_head, T, C // n_head)
attention 점수: (B, n_head, T, T)
최종 logits: (B, T, vocab_size)

코드로 직접 해보기

규칙: 이전 코드를 열지 않는다. 막히면 5분은 혼자 고민하고, 그래도 안 되면 그 줄에 # STUCK이라고 표시한 뒤 참고한다. 이 표시가 오늘의 가장 중요한 산출물이다.

1단계: 멀티헤드 어텐션 (causal)

python

import torch
import torch.nn as nn
import torch.nn.functional as F

class MHA(nn.Module):
    def __init__(self, C, n_head, block_size):
        super().__init__()
        assert C % n_head == 0
        self.n_head = n_head
        self.qkv = nn.Linear(C, 3 * C)
        self.proj = nn.Linear(C, C)
        mask = torch.tril(torch.ones(block_size, block_size))
        self.register_buffer("mask", mask.view(1, 1, block_size, block_size))

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        # (B, T, C) -> (B, n_head, T, C//n_head)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) / (k.size(-1) ** 0.5)
        att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf"))
        att = F.softmax(att, dim=-1)
        y = att @ v                                  # (B, n_head, T, d_head)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(y)

2단계: FFN과 Block

python

class FFN(nn.Module):
    def __init__(self, C):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(C, 4 * C),
            nn.GELU(),
            nn.Linear(4 * C, C),
        )
    def forward(self, x):
        return self.net(x)

class Block(nn.Module):
    def __init__(self, C, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(C)
        self.attn = MHA(C, n_head, block_size)
        self.ln2 = nn.LayerNorm(C)
        self.ffn = FFN(C)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

3단계: GPT 조립

python

class GPT(nn.Module):
    def __init__(self, vocab_size, block_size, C=128, n_head=4, n_layer=4):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, C)
        self.pos_emb = nn.Embedding(block_size, C)
        self.blocks = nn.Sequential(*[Block(C, n_head, block_size) for _ in range(n_layer)])
        self.ln_f = nn.LayerNorm(C)
        self.head = nn.Linear(C, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)
        x = self.blocks(x)
        logits = self.head(self.ln_f(x))
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

4단계: 스스로 검증하는 세 가지 테스트

python

model = GPT(vocab_size=50, block_size=16)
idx = torch.randint(0, 50, (2, 16))
logits, loss = model(idx, idx)
print(logits.shape)   # torch.Size([2, 16, 50])
print(loss.item())    # 학습 전에는 ln(50) ≈ 3.91 근처여야 정상

# 인과성(causality) 테스트: 뒤쪽 토큰을 바꿔도 앞쪽 logits는 변하면 안 된다
model.eval()
idx2 = idx.clone()
idx2[:, -1] = (idx2[:, -1] + 1) % 50
a, _ = model(idx)
b, _ = model(idx2)
print(torch.allclose(a[:, :-1], b[:, :-1], atol=1e-5))  # True여야 함

세 가지를 모두 통과해야 "동작하는 구현"이다. shape이 맞아도 마스크가 틀리면 인과성 테스트에서 걸린다.

오늘의 실험

이번 주에 했던 생성 파라미터 실험을 그대로 재현한다. 먼저 작은 텍스트(예: 이번 주에 쓴 문자 단위 데이터)로 위 모델을 몇 분 학습시킨 뒤, 같은 프롬프트로 아래 생성 함수의 설정만 바꿔 본다.

python

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None):
    model.eval()
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -model.block_size:]
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :] / temperature
        if top_k is not None:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = float("-inf")
        probs = F.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, next_id], dim=1)
    return idx

바꿀 것과 관찰할 것

실험	설정	관찰
A	temperature = 0.2 / 1.0 / 1.5	낮을수록 반복적·안전, 높을수록 다양하지만 엉뚱해지는가
B	top_k = 1 / 5 / 50	top_k=1일 때 항상 같은 출력이 나오는가
C	시드 고정 vs 미고정	같은 설정에서 출력이 재현되는가

기록 방법: 설정마다 같은 프롬프트로 출력 2~3개를 표에 붙이고, 한 줄 관찰을 적는다. 시드(torch.manual_seed), 학습 스텝 수, 최종 loss도 함께 적는다. 이 값이 없으면 나중에 결과를 재현할 수 없다.

왜 이걸 배우나 (LLM / Post-training 연결)

이번 주에 만든 구조는 규모만 다를 뿐 현대 LLM의 뼈대와 같은 계열이다. 이후 사전학습 데이터, 파인튜닝, 평가를 다룰 때 모든 것이 이 모델의 logits 위에서 일어난다. 생성 파라미터는 모델을 평가하거나 후속 학습용 응답을 샘플링할 때 결과를 직접 좌우하는 설정이므로, 샘플링 방식이 출력 품질과 재현성에 주는 영향을 몸으로 알아 두면 실험 해석이 쉬워진다. LLM Post-training 업무에서는 모델 내부를 열어 보고 디버깅하는 일이 잦기 때문에, 백지 구현 경험이 그 기초 체력이 된다.

초보자가 막히는 지점

view 에러 (`view size is not compatible`): transpose 뒤에는 메모리가 연속적이지 않다. .contiguous()를 붙인 뒤 view를 쓴다.
마스크를 softmax 뒤에 적용: 반드시 softmax 전에 -inf로 채운다. 뒤에 하면 확률 합이 1이 아니게 된다.
sqrt(d_head)를 빼먹음: 에러는 안 나지만 학습이 불안정해질 수 있다. 눈에 안 띄는 버그라 인과성·loss 테스트와 별개로 수식을 한 번 더 대조한다.
헤드 수와 차원 불일치: C % n_head != 0이면 shape이 맞지 않는다. assert를 넣어 두면 원인을 바로 찾는다.
초기 loss가 ln(vocab_size)보다 훨씬 큼: 초기화나 targets 정렬을 의심한다. 다음 토큰 예측이라면 입력과 정답이 한 칸 어긋나 있어야 한다. 위 4단계 테스트는 단순 동작 확인용으로 targets를 입력과 같게 뒀을 뿐이니, 실제 학습에서는 x = data[i:i+T], y = data[i+1:i+T+1]로 만든다.
eval 모드를 잊음: 이 코드에는 dropout이 없어 영향이 작지만, dropout을 넣으면 model.eval() 여부에 따라 결과가 달라진다. 습관으로 굳혀 둔다.
자신감 착각: 코드를 읽을 때는 다 아는 것 같아도 쓰면 막힌다. # STUCK 표시가 많다고 낙담하지 말고, 그 목록을 내일의 복습 대상으로 삼는다.

알고리즘 (병행 1시간)

이번 주에 틀린 문제를 다시 푼다. 새 문제는 풀지 않는다.

풀이마다 아래 세 가지를 기록한다.

접근: 어떤 아이디어로 풀었는지 두세 줄로 (예: "정렬 후 투 포인터")
시간복잡도: 시간과 공간을 각각 O(·)로, 입력 크기 기준으로 적는다.
틀린 이유: 처음에 왜 틀렸는지 한 줄. 개념을 몰랐는지, 엣지 케이스를 놓쳤는지, 구현 실수인지 분류해 둔다.

재풀이는 풀이를 보지 않고 시작한다. 30분 안에 못 풀면 힌트만 보고 다시 시도하고, 그래도 안 되면 풀이를 읽은 뒤 이틀 후에 한 번 더 풀 문제로 따로 표시한다.