5분 만에 노트북에서 학습할 수 있는 가장 강력한 AI 모델은 무엇일까?

• MacBook Pro에서 5분 만에 약 1.8M 파라미터의 GPT 스타일 트랜스포머 모델을 약 20M TinyStories 토큰으로 학습하여, 약 9.6 퍼플렉시티 달성
• 5분 내 학습의 주요 제약은 모델 크기와 처리 가능한 토큰 수로, 모델 크기가 크면 느린 수렴과 적은 데이터로 효과 감소
• 성능 최적화에서는 MPS 사용, 컴파일/양자화/그래디언트 누적 및 PyTorch 대체보다는 작은 모델 선택이 효과적임
• TinyStories 같은 단순하고 일관된 데이터셋이, 백과사전식 데이터보다 작은 모델 성능에 더 긍정적 영향 제공
• 트랜스포머 아키텍처가 작은 사이즈와 짧은 학습 시간 조건에서 LSTM이나 diffusion 방식보다 뛰어난 결과 보여줌

개요

이 글은 노트북(MacBook Pro)에서 5분 만에 학습 가능한 최대 성능 AI 언어 모델을 실험한 결과와, 최적의 트레이닝 전략 및 데이터셋 선정, 모델 아키텍처에 대한 인사이트를 제공함

실험 결과 요약

• 약 1.8M 파라미터의 GPT 스타일 트랜스포머 모델을 약 20M TinyStories 데이터로 학습, 9.6 퍼플렉시티 기록
• 생성 예시는 짧지만 일관된 이야기 형태로, 영어 문법이 대체로 올바르게 유지되는 수준
• 5분 내 실용적인 수준의 모델 결과는 기대 이상인 점을 강조

실험 배경 및 한계

• 노트북 환경에서 강력한 모델을 빠르게 학습하는 것은 현실적이지 않은 호기심에서 출발한 실험
• 실제로는 클라우드에서 고성능 GPU(H100 등) 로 더 강력한 모델을 학습할 수 있으나, 실험의 한계 조건은 시간(5분)임
• 모델 크기가 커질수록 토큰 처리 속도가 느려져 5분 내 좋은 결과를 내기 어려움
  • 너무 작은 모델(예: 10K 파라미터)은 충분한 복잡도를 학습하지 못함
  • 실용적인 범위는 약 1M~2M 파라미터 모델임

처리량 최적화

• MPS(Apple의 Metal Performance Shaders) 사용이 가장 효과적임
• torch.compile, float16, MLX 등 다양한 수학적 최적화는 기대보다 성능 개선 효과 미미하거나 오히려 저하
• 그래디언트 누적은 메모리 관리 목적이 있으나, 실제로는 속도 저하 심각
• 모델이 내부 메모리에서 빠르게 weight update 할 수 있어야 효율적

데이터셋 선택

• 제한된 토큰 수(약 10~20M)로 Simple English Wikipedia 등 단순 영문 위키 데이터를 먼저 사용한 결과, 문법적 일관성은 잡았으나 의미 일관성 부족
  • 고유명사 중심, 억지로 만든 듯한 사실 나열로 유의미한 내용 생성에 한계
• TinyStories 데이터셋 사용 시, 이야기 구조가 명확하고 언어가 단순하여 결과가 훨씬 더 일관적이고 의미 있음
  • 4세 수준 스토리로, 작은 모델에도 학습이 잘 이루어짐

토크나이저 및 토큰화

• 토크나이저 트레이닝은 5분 내 포함되지 않으며, 데이터 규모가 작으므로 최적화 필요성 낮음
• 멀티바이트 토큰 학습이 모델 학습에 더 쉬움

모델 아키텍처 실험

• 트랜스포머(GPT-2 스타일) 아키텍처 사용

  • 2~3 레이어, SwiGLU 등 활성화 함수, positional embedding 등 하이퍼파라미터 조정
  • LSTM은 성능이 근접하지만 트랜스포머가 퍼플렉시티 측면에서 더 우수
  • Dropout, mixture-of-experts 등은 크기가 작아 비효율적
  • Curriculum learning은 학습 시간이 너무 짧아 효과 미비

• Diffusion 모델(D3PM) 시도

  • 자연어는 이산 토큰이므로 확산 과정에서 무의미한 랜덤 토큰만 생성되어 실패
  • 트랜스포머나 LSTM 대비 빠른 문장 구조 형성이 어려움

모델 크기와 토큰/초 처리량 관계

• 1M~2M 파라미터 모델이 가장 이상적인 sweet spot
  • 너무 크면 5분 내 수렴 불가, 너무 작으면 학습 즉시 성능 한계 도달
• Chinchilla scaling law와 실험 결과가 대체로 일치함
  • 전체 훈련 토큰/20이 이상적인 모델 크기로, 해당 실험에서도 확인됨

결론 및 시사점

• 매우 짧은 시간, 작은 하드웨어로도 일관된 스토리텔링 모델 학습이 가능
• 5분 학습은 강력한 모델 개발에는 적합하지 않지만, 소규모·초경량 모델 설계와 하드웨어·아키텍처 최적화 실험에는 의의가 있음
• 향후 노트북 GPU 및 모델 구조 발전 시, 단 몇 분 내 훈련 가능한 모델 성능의 발전 가능성 있음