Model AI paling kuat apa yang bisa dilatih di laptop dalam 5 menit?

• Melatih model transformer bergaya GPT dengan sekitar 1.8M parameter di MacBook Pro dalam 5 menit menggunakan sekitar 20M token TinyStories, dan mencapai sekitar 9.6 perplexity
• Batas utama untuk pelatihan dalam 5 menit adalah ukuran model dan jumlah token yang bisa diproses; jika model terlalu besar, konvergensi melambat dan efektivitas dengan data terbatas menurun
• Dalam optimasi performa, penggunaan MPS, kompilasi/kuantisasi/akumulasi gradien, atau pengganti PyTorch ternyata kurang efektif dibanding memilih model yang lebih kecil
• Dataset yang sederhana dan konsisten seperti TinyStories memberi dampak lebih positif pada performa model kecil dibanding data bergaya ensiklopedia
• Arsitektur transformer menunjukkan hasil yang lebih baik daripada pendekatan LSTM atau diffusion dalam kondisi ukuran kecil dan waktu pelatihan singkat

---

Gambaran umum

Tulisan ini menyajikan hasil eksperimen tentang model bahasa AI dengan performa tertinggi yang bisa dilatih di laptop (MacBook Pro) dalam 5 menit, serta insight mengenai strategi training terbaik, pemilihan dataset, dan arsitektur model

Ringkasan hasil eksperimen

• Model transformer bergaya GPT dengan sekitar 1.8M parameter dilatih menggunakan sekitar 20M data TinyStories, mencatat 9.6 perplexity
• Contoh generasi memang pendek, tetapi berbentuk cerita yang konsisten, dengan tata bahasa Inggris yang umumnya tetap benar
• Ditekankan bahwa hasil model pada tingkat yang praktis dalam waktu 5 menit tergolong melampaui ekspektasi

Latar belakang eksperimen dan keterbatasan

• Eksperimen ini berawal dari rasa ingin tahu yang kurang realistis: melatih model yang kuat dengan cepat di lingkungan laptop
• Dalam praktiknya, model yang lebih kuat tentu bisa dilatih di cloud dengan GPU berperforma tinggi (seperti H100), tetapi batasan utama eksperimen ini adalah waktu, yaitu 5 menit
• Semakin besar ukuran model, semakin lambat kecepatan pemrosesan token, sehingga sulit menghasilkan hasil yang baik dalam 5 menit
  • Model yang terlalu kecil (misalnya 10K parameter) tidak mampu mempelajari kompleksitas yang memadai
  • Rentang yang praktis adalah model dengan sekitar 1M~2M parameter

Optimasi throughput

• Penggunaan MPS (Metal Performance Shaders milik Apple) adalah yang paling efektif
• Berbagai optimasi matematis seperti torch.compile, float16, MLX, dan lainnya memberi peningkatan performa yang kecil, lebih rendah dari harapan, atau bahkan justru menurunkannya
• Akumulasi gradien berguna untuk pengelolaan memori, tetapi dalam praktiknya menyebabkan penurunan kecepatan yang signifikan
• Agar efisien, model harus bisa melakukan update bobot dengan cepat di memori internal

Pemilihan dataset

• Dengan jumlah token yang terbatas (sekitar 10~20M), percobaan awal menggunakan data wiki bahasa Inggris sederhana seperti Simple English Wikipedia menghasilkan konsistensi gramatikal, tetapi kurang konsisten secara makna
  • Karena berpusat pada nama diri dan daftar fakta yang terasa dipaksakan, ada batasan dalam menghasilkan konten yang bermakna
• Saat menggunakan dataset TinyStories, hasilnya jauh lebih konsisten dan bermakna karena struktur ceritanya jelas dan bahasanya sederhana
  • Ini adalah cerita setingkat anak usia 4 tahun, sehingga bahkan model kecil pun dapat mempelajarinya dengan baik

Tokenizer dan tokenisasi

• Training tokenizer tidak dihitung dalam 5 menit, dan karena skala datanya kecil, kebutuhan optimasinya juga rendah
• Mempelajari token multibyte lebih mudah bagi model

Eksperimen arsitektur model

• Menggunakan arsitektur transformer (gaya GPT-2)

  • Hyperparameter seperti 2~3 layer, fungsi aktivasi SwiGLU, positional embedding, dan lainnya disesuaikan
  • LSTM memiliki performa yang cukup dekat, tetapi transformer lebih unggul dari sisi perplexity
  • Dropout, mixture-of-experts, dan sejenisnya tidak efisien pada ukuran sekecil ini
  • Curriculum learning hampir tidak memberi efek karena waktu pelatihannya terlalu singkat

• Percobaan dengan model diffusion (D3PM)

  • Karena bahasa alami terdiri dari token diskret, proses difusi hanya menghasilkan token acak yang tidak bermakna dan akhirnya gagal
  • Dibanding transformer atau LSTM, model ini sulit membentuk struktur kalimat dengan cepat

Hubungan ukuran model dan throughput token/detik

• Model dengan 1M~2M parameter adalah sweet spot yang paling ideal
  • Jika terlalu besar, model tidak bisa konvergen dalam 5 menit; jika terlalu kecil, batas performanya langsung tercapai saat training
• Hukum penskalaan Chinchilla secara umum sejalan dengan hasil eksperimen
  • Ukuran model ideal adalah total token pelatihan/20, dan hal ini juga terkonfirmasi dalam eksperimen tersebut

Kesimpulan dan implikasi

• Bahkan dengan waktu yang sangat singkat dan perangkat keras kecil, melatih model storytelling yang konsisten itu mungkin dilakukan
• Pelatihan 5 menit memang tidak cocok untuk mengembangkan model yang kuat, tetapi tetap bermakna untuk eksperimen desain model kecil dan ultra-ringan serta optimasi hardware dan arsitektur
• Seiring perkembangan GPU laptop dan struktur model di masa depan, ada kemungkinan performa model yang bisa dilatih hanya dalam beberapa menit akan terus meningkat