LLM Research2026.10.099 min read

GPT 모델 조립 & 생성: 블록을 쌓고 LM Head를 달면 logits와 loss가 나온다

지금까지 만든 Transformer 블록을 N개 쌓고, 마지막에 final norm과 LM Head를 붙여 작은 GPT를 완성한다. forward가 logits와 loss를 돌려주게 만들고, generate()로 greedy·temperature·top-k·top-p 방식으로 글자를 하나씩 뽑아 본다. 같은 프롬프트로 샘플링 설정만 바꿔 출력이 어떻게 달라

By 송민성

오늘의 목표

작은 GPT를 처음부터 끝까지 직접 만든다. 토큰 번호를 입력받아 "다음 토큰 점수(logits)"와 loss를 돌려주는 GPT 클래스를 완성하고, 그 위에 한 토큰씩 이어 쓰는 generate() 함수를 얹는다. 오늘이 끝나면 내가 만든 모델이 (학습이 충분하지 않더라도) 글자를 생성하는 것까지 눈으로 확인한다.

> 전제: 앞선 Day에서 만든 Transformer 블록(Block, 내부에 self-attention + MLP + 정규화 + 잔차 연결)이 있다고 가정한다. 없다면 아래 코드의 Block 자리에 직접 만든 블록을 넣는다. 아래에는 오늘 실행에 필요한 최소 Block도 함께 적어 둔다.

개념과 직관

1. 전체 흐름: 숫자 → 벡터 → 블록 N번 → 점수

GPT는 "공장 라인"과 비슷하다.

text

토큰 번호 (B, T)
  → 토큰 임베딩 + 위치 임베딩   (번호를 벡터로)
  → Block × N                   (문맥을 섞으며 벡터를 다듬음)
  → final norm                  (마지막 정리)
  → LM Head                     (벡터를 어휘 크기의 점수로)
logits (B, T, vocab_size)

임베딩(embedding, 토큰 번호를 고정 길이 벡터로 바꿔 주는 표): 번호 5를 "5번 행의 벡터"로 바꾼다. 위치 임베딩은 "몇 번째 자리인가"를 같은 방식으로 더해 준다.
블록(block): 같은 구조를 N번 쌓는다. 층이 깊을수록 더 복잡한 관계를 표현할 수 있지만 느려진다.
final norm: 블록을 다 지난 벡터의 크기를 마지막으로 정돈한다. 이걸 빼면 출력 점수의 크기가 들쭉날쭉해져 학습이 불안정해질 수 있다.
LM Head(Language Modeling Head, 벡터를 "어휘 안의 각 토큰이 다음에 올 점수"로 바꾸는 선형층): 출력 크기가 vocab_size다. 이 점수가 logits(확률로 바꾸기 전의 원점수)이다.

2. 손실(loss): 한 칸씩 밀어서 맞히기

입력이 [A, B, C, D]이면 정답(target)은 한 칸 민 [B, C, D, E]이다. 각 위치에서 "다음 토큰"을 맞히게 하고, 정답 토큰에 높은 점수를 줄수록 loss가 작아진다. 이때 쓰는 것이 교차 엔트로피(cross entropy, 정답에 준 확률이 낮을수록 커지는 손실)이다. 최소 수식:

text

loss = -log( 정답 토큰의 확률 )

정답 확률이 1이면 loss 0, 0.01이면 약 4.6이다. 모델이 어휘 크기 V개를 완전히 찍어서 맞힌다면 loss는 약 log(V)다. 학습을 시작하기 직전 loss가 이 값 근처인지 확인하면 "초기화가 정상인가"를 빠르게 점검할 수 있다.

3. Weight tying: 입력 표와 출력 표를 공유

토큰 임베딩은 (vocab_size, d_model) 모양의 표이고, LM Head의 가중치도 같은 모양 (vocab_size, d_model)이다. 가중치 공유(weight tying, 입력 임베딩과 출력층이 같은 파라미터를 쓰는 것)를 하면 파라미터가 줄고, "이 토큰의 의미 벡터"가 입력과 출력에서 일관되게 쓰인다. 코드로는 한 줄이다.

python

self.lm_head.weight = self.tok_emb.weight

작은 모델일수록 어휘 표가 전체 파라미터에서 차지하는 비중이 커서 절약 효과가 크다.

4. 초기화(initialization)

학습 시작 시점의 가중치를 어떻게 정하느냐의 문제다. 너무 크면 출력이 폭주하고, 0이면 모든 뉴런이 똑같이 움직인다. 흔히 쓰는 방식은 평균 0, 표준편차 0.02 정도의 작은 정규분포 난수로 선형층과 임베딩을 채우고 bias는 0으로 두는 것이다. (0.02는 GPT-2 계열 구현에서 자주 보이는 관례 값이며, 정답이 하나로 정해진 수치는 아니다.)

5. 생성(generation): 점수를 확률로, 확률에서 한 토큰 뽑기

모델은 한 번에 한 토큰만 예측한다. 그래서 "예측 → 뽑기 → 입력 끝에 붙이기"를 반복한다. 마지막 위치의 logits만 쓰고, 이를 소프트맥스(softmax, 점수들을 합이 1인 확률로 바꾸는 함수)에 통과시킨다. 어떻게 뽑느냐가 네 가지로 갈린다.

방식	직관	특징
greedy	항상 1등만 고른다	결정적(같은 입력이면 같은 출력), 같은 말을 반복하기 쉬움
temperature	점수를 T로 나눈 뒤 softmax. T<1이면 1등이 더 두드러지고, T>1이면 평평해진다	낮으면 보수적, 높으면 엉뚱해짐
top-k	상위 k개 후보만 남기고 나머지는 버린 뒤 그중에서 뽑는다	이상한 꼬리 후보 차단
top-p (nucleus)	확률이 높은 순으로 더해서 합이 p를 넘는 최소 집합만 남기고 뽑는다	후보 수가 상황에 따라 자동 조절

temperature의 최소 수식:

text

확률 = softmax(logits / T)

비유: 주사위를 굴려 다음 글자를 정하는데, temperature는 "주사위가 얼마나 편향되게 만들어졌나"를, top-k/top-p는 "애초에 주사위 면에 올릴 후보를 몇 개로 제한하나"를 조절한다.

코드로 직접 해보기

Step 0. 최소 Block (이미 만들었다면 건너뜀)

python

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

class Block(nn.Module):
    def __init__(self, d_model, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(d_model, n_head, batch_first=True)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )
        mask = torch.triu(torch.ones(block_size, block_size), diagonal=1).bool()
        self.register_buffer("mask", mask)

    def forward(self, x):
        T = x.size(1)
        h = self.ln1(x)
        a, _ = self.attn(h, h, h, attn_mask=self.mask[:T, :T], need_weights=False)
        x = x + a
        x = x + self.mlp(self.ln2(x))
        return x

직접 만든 Attention이 있다면 그것을 쓰는 편이 학습에 낫다. 위 코드는 "미래 토큰을 못 보게 가리는 마스크"가 들어간 동작 확인용이다.

Step 1. GPT 클래스

python

class GPT(nn.Module):
    def __init__(self, vocab_size, block_size, d_model=128, n_head=4, n_layer=4):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(block_size, d_model)
        self.blocks = nn.ModuleList(
            [Block(d_model, n_head, block_size) for _ in range(n_layer)]
        )
        self.ln_f = nn.LayerNorm(d_model)               # final norm
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

        self.lm_head.weight = self.tok_emb.weight       # weight tying
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        assert T <= self.block_size, "입력이 block_size보다 깁니다"
        pos = torch.arange(T, device=idx.device)
        x = self.tok_emb(idx) + self.pos_emb(pos)       # (B, T, d_model)
        for blk in self.blocks:
            x = blk(x)
        x = self.ln_f(x)
        logits = self.lm_head(x)                        # (B, T, vocab_size)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
            )
        return logits, loss

Step 2. 시작 loss 점검

python

import math
V = 65
model = GPT(vocab_size=V, block_size=64)
x = torch.randint(0, V, (8, 64))
y = torch.randint(0, V, (8, 64))
logits, loss = model(x, y)
print(logits.shape)            # torch.Size([8, 64, 65])
print(loss.item(), math.log(V))  # 둘 다 4.17 근처여야 정상

두 값이 크게 다르면 초기화나 loss 계산에 문제가 있다는 신호다.

Step 3. generate()

python

@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None, top_p=None, greedy=False):
    model.eval()
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -model.block_size:]           # 너무 길면 뒤쪽만 사용
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :]                       # 마지막 위치만 (B, V)

        if greedy:
            next_id = logits.argmax(dim=-1, keepdim=True)
        else:
            logits = logits / temperature
            if top_k is not None:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = float("-inf")
            if top_p is not None:
                sorted_logits, sorted_idx = torch.sort(logits, descending=True)
                probs = F.softmax(sorted_logits, dim=-1)
                cum = probs.cumsum(dim=-1)
                remove = cum - probs > top_p            # 1등은 항상 남도록
                sorted_logits[remove] = float("-inf")
                logits = torch.full_like(logits, float("-inf")).scatter(1, sorted_idx, sorted_logits)
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)

        idx = torch.cat([idx, next_id], dim=1)
    return idx

사용 예(학습 전이라 무작위 글자가 나오는 게 정상):

python

start = torch.zeros((1, 1), dtype=torch.long)
out = generate(model, start, 50, temperature=0.8, top_k=20)
print(out[0].tolist())

글자로 보려면 앞선 Day에서 쓴 decode 함수(토큰 번호 → 문자열)를 연결한다. 의미 있는 문장을 보려면 작은 텍스트(예: 앞선 Day의 데이터셋)로 먼저 몇 천 스텝 학습한 뒤 실험한다.

오늘의 실험

질문: 같은 프롬프트에서 샘플링 설정만 바꾸면 출력이 어떻게 달라지는가?

앞선 Day의 학습 루프로 작은 GPT를 짧게라도 학습한다(loss가 초기값보다 확실히 내려간 상태). 학습된 체크포인트를 쓰면 더 좋다.
프롬프트를 하나로 고정하고, 시드도 torch.manual_seed(0)으로 맞춘 뒤 아래 설정을 각각 돌린다. 길이는 모두 max_new_tokens=200.

번호	설정
A	greedy
B	temperature=0.5
C	temperature=1.0
D	temperature=1.5
E	temperature=1.0, top_k=10
F	temperature=1.0, top_p=0.9

기록할 것(표로 정리):

- 설정, 출력 앞 200자 - 같은 구절이 반복되는가(반복 정도를 "없음/조금/심함"으로) - 철자·단어가 깨지는가 - 시드를 바꿔 3번 돌렸을 때 출력이 달라지는가(greedy는 안 달라져야 정상)

관찰 포인트: temperature가 올라갈수록 다양성과 오류가 함께 늘어나는지, top-k·top-p가 높은 temperature의 "이상한 글자"를 얼마나 막아 주는지 비교한다. 결론은 한두 문장으로 직접 써 본다. 작은 모델·적은 데이터라 일반적인 결론으로 단정하지 말고 "내 실험에서는"이라고 적는다.

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

오늘 만든 구조(임베딩 → 블록 N개 → final norm → LM Head, 다음 토큰 loss)는 규모만 다를 뿐 현대 디코더 기반 LLM이 공유하는 뼈대다. 사전학습의 loss도, 이후 SFT(지도 미세조정)에서 다루는 loss도 이 logits에서 출발하므로 forward가 무엇을 돌려주는지 정확히 아는 것이 출발점이다. 또한 Post-training 평가와 데이터 생성 단계에서는 temperature·top-p 같은 샘플링 설정이 결과 품질과 재현성에 직접 영향을 주기 때문에, 설정별 차이를 직접 체감해 두면 실험 결과를 해석할 때 도움이 된다.

초보자가 막히는 지점

shape 에러 `view` 관련: logits.view(-1, V)와 targets.view(-1)의 앞 차원이 같아야 한다. 입력과 정답이 한 칸 밀려 있는지, (B, T) 모양인지 먼저 출력해 본다. view가 안 되면 reshape를 쓰거나 .contiguous()를 붙인다.
시작 loss가 log(V)와 크게 다름: 초기화를 빼먹었거나 std가 너무 큰 경우가 많다. self.apply(self._init_weights) 호출 여부를 확인한다. 단, weight tying 상태에서는 임베딩과 head가 같은 파라미터라는 점도 기억한다.
weight tying 후 파라미터 수가 안 줄어듦: sum(p.numel() for p in model.parameters())는 공유 파라미터를 한 번만 센다. tying 줄을 지웠을 때와 비교해 보면 차이가 보인다.
`block_size` 초과 에러: 생성이 길어지면 위치 임베딩 범위를 넘는다. idx[:, -block_size:]로 자르는 줄을 빼먹지 않는다.
top-p에서 후보가 전부 사라짐: "누적합이 p를 넘는 순간까지 포함"해야 한다. 위 코드의 cum - probs > top_p는 1등이 항상 남게 하려는 장치다.
temperature=0 나누기 에러: 0으로 나눌 수 없다. 결정적 출력이 필요하면 greedy=True를 쓴다.
학습 전인데 출력이 엉망이라 버그로 착각: 학습 전 모델은 무작위 글자를 내는 것이 정상이다. 실험은 학습 후에 한다.
`model.eval()`을 빼먹음: dropout을 쓰는 구조라면 생성 결과가 흔들린다. 생성 함수 안에서 eval()을 호출하고, 다시 학습할 때는 model.train()을 부른다.

알고리즘 (병행 1시간)

LeetCode Medium 2문제를 구현 중심으로 푼다. 문제 선택 기준은 배열·문자열·해시맵·스택·시뮬레이션처럼 "자료구조를 올바르게 코드로 옮기는 연습"이 되는 유형이다(구체적인 문제는 본인이 골라도 좋다). 문제당 30분을 넘기면 힌트나 풀이를 보고, 이후 아무것도 보지 않고 다시 구현한다.

풀이 후 노트에 남길 것:

접근: 처음 떠올린 방법과 최종 방법, 그 사이에 바꾼 이유
시간복잡도 / 공간복잡도: 입력 크기 n 기준으로 한 줄씩
틀린 이유: 어떤 입력에서 실패했는지(빈 입력, 길이 1, 중복, 경계값 등)와 고친 방법
다시 풀 때 기억할 한 줄