모두를 위한 AI 1편 — 16M 파라미터 언어모델을 15분에 처음부터 학습시키기

• 들어가며 — 왜 작은 모델을 직접 만들어 보는가
• 실험 환경
• 데이터셋 — TinyStories
• 모델 구조 — 디코더 전용 트랜스포머
  • 인과 마스크가 없으면 학습이 성립하지 않는다
  • 가중치 묶기 — 파라미터의 40%를 절약한다
  • Pre-norm 배치
• 학습 루프
• 학습 로그
• 생성 결과
• perplexity 8이 뜻하는 것
• 정리
• 🧠 이해도 체크 퀴즈
• 참고 자료

들어가며 — 왜 작은 모델을 직접 만들어 보는가

거대 언어모델을 쓰는 법을 다루는 글은 많습니다. 하지만 "안에서 무슨 일이 일어나는가"는 API 뒤에 가려져 있습니다. 이 시리즈는 반대 방향으로 갑니다. 직접 학습시켜 보는 것입니다.

목표는 성능이 아닙니다. 조건은 이렇게 잡았습니다.

• 파라미터 1억 개 미만 (실제로는 훨씬 작습니다)
• GPU 한 장에서 20분 이내 학습
• 데이터는 Hugging Face에 공개된 것만 사용
• 학습 로그와 결과물을 전부 남기기

작게 만들면 전체 구조가 한눈에 들어옵니다. 어텐션이 왜 필요한지, 손실 함수가 무엇을 재는지, 학습이 실패할 때 어떤 신호가 나타나는지를 몇 분 단위로 관찰할 수 있습니다.

이번 편은 그 첫 번째, 텍스트 생성 모델입니다.

실험 환경

실제로 사용한 장비와 버전입니다.

데이터셋 — TinyStories

roneneldan/TinyStories는 3~4세 아이가 이해할 수 있는 어휘만으로 쓰인 짧은 영어 동화 모음입니다. 이 데이터셋이 작은 모델 실험에 이상적인 이유가 있습니다.

일반적인 웹 텍스트는 어휘가 방대해서 작은 모델이 문법을 익히기 전에 단어를 외우느라 용량을 소진합니다. TinyStories는 어휘를 의도적으로 제한해서, 작은 모델도 "문장을 만드는 법" 자체에 용량을 쓸 수 있게 합니다.

토크나이저는 새로 만들지 않고 GPT-2의 것을 그대로 빌렸습니다. 허브에 올라온 모델의 일부만 가져다 쓰는 것도 유용한 방법입니다.

from transformers import AutoTokenizer
from datasets import load_dataset

tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
ds = load_dataset("roneneldan/TinyStories", split="train[:40000]")

ids = []
for ex in ds:
    ids.extend(tok(ex["text"]).input_ids + [tok.eos_token_id])
    if len(ids) >= 4_000_000:
        break
data = torch.tensor(ids[:4_000_000], dtype=torch.long)

토큰을 하나의 긴 스트림으로 이어 붙였습니다. 학습할 때는 이 스트림에서 임의 위치를 잘라 씁니다. 문서 경계를 넘나드는 샘플이 생기지만, eos_token_id가 경계를 표시하고 있어 모델이 "여기서 이야기가 끝난다"를 배웁니다.

모델 구조 — 디코더 전용 트랜스포머

전체 구조는 이게 전부입니다. 4개 층, 임베딩 차원 256, 헤드 8개입니다.

CTX, DIM, LAYERS, HEADS = 128, 256, 4, 8

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(DIM), nn.LayerNorm(DIM)
        self.attn = nn.MultiheadAttention(DIM, HEADS, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(DIM, 4 * DIM), nn.GELU(), nn.Linear(4 * DIM, DIM)
        )

    def forward(self, x, mask):
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.ln2(x))


class TinyGPT(nn.Module):
    def __init__(self, vocab):
        super().__init__()
        self.emb = nn.Embedding(vocab, DIM)
        self.pos = nn.Embedding(CTX, DIM)
        self.blocks = nn.ModuleList(Block() for _ in range(LAYERS))
        self.ln = nn.LayerNorm(DIM)
        mask = torch.triu(torch.full((CTX, CTX), float("-inf")), 1)
        self.register_buffer("mask", mask)

    def forward(self, x):
        T = x.shape[1]
        h = self.emb(x) + self.pos(torch.arange(T, device=x.device))
        for b in self.blocks:
            h = b(h, self.mask[:T, :T])
        return self.ln(h) @ self.emb.weight.T  # weight tying

여기서 짚어볼 설계 결정이 세 가지 있습니다.

인과 마스크가 없으면 학습이 성립하지 않는다

torch.triu(..., 1) 로 만든 상삼각 행렬은 대각선 위쪽을 전부 음의 무한대로 채웁니다. 소프트맥스를 통과하면 이 자리들의 가중치가 0이 됩니다. 즉 각 위치가 자기보다 뒤에 있는 토큰을 볼 수 없게 만드는 장치입니다.

이게 없으면 어떻게 될까요. 모델은 다음 토큰을 맞히는 문제를 푸는데, 정답인 다음 토큰을 이미 입력으로 보고 있게 됩니다. 손실은 0에 가깝게 떨어지지만 생성은 전혀 되지 않습니다. "손실이 너무 빨리 떨어진다"는 건 대개 정답이 새고 있다는 신호입니다.

가중치 묶기 — 파라미터의 40%를 절약한다

마지막 줄의 @ self.emb.weight.T 가 가중치 묶기(weight tying)입니다. 출력층을 따로 두지 않고 입력 임베딩 행렬을 전치해 재사용합니다.

GPT-2 어휘는 50,257개입니다. 출력층을 따로 두면 50257 × 256 = 약 1,287만 개의 파라미터가 추가로 필요합니다. 전체가 1,606만 개이니 묶지 않았다면 2,893만 개가 됐을 것입니다. 게다가 "토큰을 벡터로 바꾸는 일"과 "벡터를 토큰으로 되돌리는 일"은 서로 역방향의 같은 작업이라, 같은 행렬을 쓰는 게 자연스럽습니다.

Pre-norm 배치

self.ln1(x) 를 어텐션 앞에 적용하고, 잔차 연결은 정규화되지 않은 x 에 더합니다. 이 순서(pre-norm)가 post-norm보다 깊은 모델에서 학습이 안정적이라고 알려져 있습니다. 4층에서는 차이가 크지 않지만 관례를 따랐습니다.

학습 루프

model = TinyGPT(len(tok)).to("cuda")
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

step = 0
while not run.over_budget():          # 15분 예산
    ix = torch.randint(0, len(data) - CTX - 1, (64,))
    x = torch.stack([data[i:i + CTX] for i in ix]).to("cuda")
    y = torch.stack([data[i + 1:i + CTX + 1] for i in ix]).to("cuda")
    with torch.autocast("cuda", torch.bfloat16):
        loss = F.cross_entropy(model(x).flatten(0, 1), y.flatten())
    opt.zero_grad(); loss.backward(); opt.step()
    step += 1

x 와 y 가 한 칸 어긋나 있는 것이 전부입니다. 입력이 "Once upon a"라면 정답은 "upon a time"입니다. 모든 위치에서 동시에 다음 토큰을 맞히므로, 길이 128짜리 샘플 하나가 128개의 학습 신호를 만듭니다.

torch.autocast 로 bfloat16을 씁니다. 3090에서 메모리와 속도 모두 이득이고, bf16은 float16과 달리 지수 범위가 float32와 같아 손실 스케일링 없이도 안정적입니다.

학습 로그

실제로 기록된 값입니다.

[    0.6s] step      0 loss=168.4468 ppl=...
[   38.2s] step   1000 loss=5.2213  ppl=185.13
[  114.7s] step   3000 loss=3.6841  ppl=39.81
[  305.9s] step   6700 loss=3.0350  ppl=20.80
[  662.6s] step  16100 loss=2.2882  ppl=9.86
[  901.2s] step  22358 loss=2.0948  ppl=8.12

• 총 22,358 스텝, 배치 64, 컨텍스트 128 → 약 1억 8천만 토큰을 처리했습니다
• 손실 168에서 시작하는 것은 초기화 직후의 값입니다. 균등 분포라면 ln(50257) ≈ 10.8 이 나와야 하는데, LayerNorm과 가중치 묶기의 상호작용으로 첫 스텝에 큰 값이 잡혔다가 즉시 정상 범위로 내려옵니다
• 최종 perplexity 8.12는 다음 단어 후보를 평균 8개 정도로 좁혔다는 뜻입니다

생성 결과

온도 0.8로 샘플링한 실제 출력입니다. 다듬지 않았습니다.

PROMPT: Once upon a time
Once upon a time, there was a little boy named Timmy. Timmy liked to play
with his toy cars instead. One day, Timmy wanted to play with the red cars,
but he couldn't find it.

As Timmy was leaving a fight on the slide, he saw a big tree. Timmy thought
it looked like fun, so he started to tremble. He mixed the leaves and

PROMPT: Lily went to the
Lily went to the park with her mom. They saw a big tree, a hole and a hole.
They were curious and wanted to see what was inside.

They heard a hole in the hole. It was a zipper, but it was in the hole.

이 결과를 정직하게 읽어 봅시다.

잘 된 것 — 문법이 거의 완벽합니다. 시제 일치, 관사, 대명사 지시가 맞습니다. 등장인물 이름을 문단 내내 일관되게 유지합니다. 따옴표를 열면 닫습니다. 문단 구분도 자연스럽습니다.

안 된 것 — 의미가 무너집니다. "a big tree, a hole and a hole"처럼 같은 단어를 반복하고, "leaving a fight on the slide"는 문법적이지만 뜻이 통하지 않습니다. "he started to tremble"은 앞뒤와 연결되지 않습니다.

이 대비가 중요합니다. 문법은 국소적 패턴이라 작은 모델도 배우지만, 일관성은 긴 의존 관계라 용량과 컨텍스트가 필요합니다. 컨텍스트가 128 토큰이니 모델이 기억할 수 있는 범위 자체가 짧습니다. 모델을 키우면 먼저 좋아지는 것이 바로 이 부분입니다.

perplexity 8이 뜻하는 것

perplexity는 exp(cross_entropy) 입니다. 직관적으로는 "모델이 다음 토큰을 몇 개 후보로 좁혔는가" 입니다.

• 학습 전(균등 분포): 50,257
• 1,000 스텝 후: 185
• 최종: 8.12

50,257개 중에서 고르던 것을 8개로 좁혔습니다. 다만 이 숫자를 다른 모델과 직접 비교하면 안 됩니다. perplexity는 토크나이저와 데이터 분포에 강하게 의존하기 때문입니다. TinyStories처럼 어휘가 제한된 데이터에서는 낮은 값이 쉽게 나옵니다. 같은 데이터·같은 토크나이저 안에서 비교할 때만 의미가 있는 지표입니다.

정리

1,600만 파라미터, 15분, 전기 요금으로 치면 100원 남짓입니다. 그것으로 문법에 맞는 영어 동화를 쓰는 모델이 나왔습니다.

배운 것을 세 줄로 요약하면 이렇습니다. 인과 마스크는 정답 누출을 막는 장치이고, 손실이 비정상적으로 빨리 떨어지면 이걸 먼저 의심해야 합니다. 가중치 묶기는 어휘가 큰 모델에서 파라미터의 절반 가까이를 절약합니다. 그리고 문법과 의미는 서로 다른 난이도의 문제라, 작은 모델은 전자를 먼저 배웁니다.

다음 편에서는 하나의 모델이 텍스트와 이미지를 동시에 다루는 방법을 봅니다. 픽셀을 토큰으로 취급하면 같은 트랜스포머가 이미지 생성과 이미지 설명을 모두 할 수 있습니다.

🧠 이해도 체크 퀴즈

참고 자료

• TinyStories 데이터셋
• TinyStories 논문 — How Small Can Language Models Be and Still Speak Coherent English?
• Attention Is All You Need
• Using the Output Embedding to Improve Language Models — 가중치 묶기