Memahami ELF (Efficient Language Framework): Cara Kerja Model Bahasa Difusi Kontinu

Industri model bahasa terus berkembang, dan inovasi baru selalu muncul untuk mendorong batasan yang ada. Salah satu yang patut kita perhatikan adalah ELF (Efficient Language Framework), sebuah kelas model bahasa difusi kontinu yang membawa pendekatan segar. ELF menonjol karena cara kerjanya yang unik, beroperasi sebagian besar dalam ruang embedding kontinu sebelum memetakan ke token diskrit pada langkah waktu terakhir.

Apa Itu ELF dan Mengapa Penting?

ELF adalah implementasi resmi JAX untuk sebuah makalah penelitian, yang dibangun di atas dasar Continuous-time Flow Matching. Berbeda dari Model Bahasa Difusi (DLM) yang ada, ELF dirancang untuk sebagian besar tetap berada dalam ruang embedding kontinu. Ini terjadi sampai langkah waktu akhir, di mana model akan memetakan ke token diskrit menggunakan jaringan dengan bobot bersama. Formulasi ini membuat ELF sangat adaptif. Kita bisa dengan mudah mengadopsi berbagai teknik yang sudah mapan dari model difusi domain gambar, seperti classifier-free guidance (CFG).

Implementasi ini dibuat dan diuji pada perangkat TPU. Versi PyTorch akan segera dirilis, jadi para pengembang dengan preferensi PyTorch juga bisa menantikan pembaruan ini.

Memulai dengan ELF: Instalasi dan Verifikasi

Untuk mencoba ELF, ada beberapa dependensi yang perlu kamu siapkan, yaitu JAX dan TPU. Jika kamu ingin melacak eksperimenmu, kamu bisa login ke WandB.

Kamu bisa memverifikasi pengaturanmu dengan cepat menggunakan checkpoint yang sudah disediakan. Penting untuk diingat bahwa hasil metrik mungkin sedikit berbeda karena perbedaan setup komputasi. Hasil yang dilaporkan dalam penelitian ini dihitung menggunakan TPU v5p-64.

bash
# Contoh perintah untuk verifikasi setup
# Pastikan JAX dan TPU sudah terinstal

Semua checkpoint yang sudah dilatih tersedia di HuggingFace di bawah lillian039/ELF. Checkpoint ini akan diunduh secara otomatis melalui .from_pretrained, jadi kamu tidak perlu mengunduhnya secara manual. Jika kamu ingin menggunakan checkpoint yang dilatih secara lokal, kamu cukup memasukkan jalur ke file checkpoint spesifik tersebut.

Konfigurasi dan Kustomisasi Eksperimen

ELF menggunakan sistem konfigurasi berlapis untuk memudahkan kustomisasi. Ada dua lapisan konfigurasi utama:

dataclass: Berisi semua hyperparameter default.
YAML: Berisi override khusus tugas yang dimuat oleh main.py.

Sistem akan menggabungkan keduanya, memungkinkan kamu menyesuaikan hanya parameter yang kamu butuhkan. Beberapa file konfigurasi sudah tersedia dan menunjuk ke data HuggingFace serta encoder T5, jadi kamu bisa menjalankannya langsung. Kamu mungkin ingin mengedit save_dir untuk menentukan lokasi penyimpanan sampel dan log, serta wandb_entity jika kamu menggunakan WandB.

Untuk membuat eksperimen kustom, kamu bisa membuat file konfigurasi baru di bawah configs/experiments/.

Melatih Model Bahasa ELF Kamu

ELF mendukung tiga pengaturan tugas utama:

Generasi tanpa syarat (unconditional generation) pada OpenWebText.
Terjemahan mesin (machine translation) pada WMT14 De-En.
Ringkasan abstrak (abstractive summarization) pada XSum.

Semua tugas ini menggunakan encoder T5 yang dibekukan untuk pemetaan teks ke embedding.

Data dan Dataset Kustom

Kami menyediakan split data yang sudah di-tokenisasi (dengan tokenizer T5) dan encoder JAX T5-small di HuggingFace di bawah lillian039/ELF-data. Ini akan dimuat langsung melalui .from_pretrained. Kamu juga bisa menggunakan salinan lokal dengan menunjuk dataset.path ke direktori yang disimpan dengan dataset.save_to_disk().

Untuk melatih model dengan dataset kustom, kamu perlu melakukan pra-tokenisasi menggunakan tokenizer T5 dan menyimpannya sebagai HuggingFace Dataset.

Generasi tanpa syarat: Setiap contoh hanya membutuhkan text_ids, yaitu ID token dari teks yang akan dihasilkan.
Generasi bersyarat: Setiap contoh membutuhkan target_ids (teks target/output) dan source_ids (teks sumber/input, misalnya kalimat bahasa Jerman atau artikel).

Collator akan secara otomatis menambahkan token eos dan membangun attention mask yang sesuai.

Perintah Pelatihan dan Konfigurasi yang Tersedia

Jalankan perintah berikut untuk memulai pelatihan:

bash
python -m elf.train --config <path_to_config_file>

Berikut adalah beberapa konfigurasi pelatihan yang tersedia:

Konfigurasi	Tugas
`unconditional_generation/elf_b` (default)	Generasi tanpa syarat pada OpenWebText (ELF-B)
`unconditional_generation/elf_m`	Generasi tanpa syarat pada OpenWebText (ELF-M)
`unconditional_generation/elf_l`	Generasi tanpa syarat pada OpenWebText (ELF-L)
`machine_translation/wmt14_de_en`	Terjemahan mesin WMT14 De-En
`abstractive_summarization/xsum`	Ringkasan abstrak XSum

Pelatihan default ELF-B menggunakan optimizer Muon dengan base learning rate (blr) 0.001. Learning rate efektif akan menjadi 0.002 pada batch size default 512. Model dilatih selama 5 epoch pada OpenWebText (~95K langkah) dengan TPU v5p-64, yang membutuhkan sekitar 1.5 jam per epoch.

Sampling dan Evaluasi

Proses sampling di ELF dipisahkan dari pelatihan. Ini dikendalikan oleh file YAML terpisah di configs/samplers/, yang direferensikan dari setiap konfigurasi pelatihan melalui sampler_config.

Ada beberapa jadwal sampling yang bisa kamu gunakan:

Generasi tanpa syarat: Dua jadwal SDE, yaitu 32-langkah (γ=1.5) dan 64-langkah (γ=1.0), keduanya dengan SC-CFG=3.
Generasi bersyarat (terjemahan / ringkasan): Satu jadwal ODE 64-langkah dengan CFG=2 dan SC-CFG=1.

Evaluator akan menghasilkan 1.000 sampel dan melaporkan metrik seperti Generative Perplexity (PPL) dan unigram entropy untuk generasi tanpa syarat. Untuk tugas bersyarat, akan dilaporkan metrik BLEU atau ROUGE.

Berikut adalah beberapa metrik yang diharapkan:

Tugas	Metrik	Nilai yang Diharapkan	Catatan
Generasi tanpa syarat (ELF-B, 32 SDE steps)	Gen. PPL (GPT-2 Large)	≈ 24
	Unigram Entropy	≈ 5.15
Terjemahan WMT14 De-En	BLEU	≈ 26.7	Pada set validasi. Angka di makalah menggunakan set tes.
Ringkasan XSum	ROUGE-1	≈ 36.3	Pada set validasi. Angka di makalah menggunakan set tes.
	ROUGE-2	≈ 12.5	Pada set validasi. Angka di makalah menggunakan set tes.
	ROUGE-L	≈ 28.1	Pada set validasi. Angka di makalah menggunakan set tes.

Untuk input JSONL khusus evaluasi (teks mentah, di-tokenisasi saat dimuat), kamu bisa melihat datasets.load_dataset('json', data_files=...).

Manajemen Checkpoint

Checkpoint disimpan pada akhir setiap epoch, atau pada interval fraksional jika save_interval_steps diatur. Sistem akan menyimpan hingga 10 checkpoint terbaru. Hanya proses 0 yang menulis ke disk. Jika hf_repo_id diatur dalam konfigurasi, seluruh direktori save_dir akan diunggah ke HuggingFace setelah setiap penyimpanan.

Jika checkpoint_path tidak ditentukan, pelatihan akan secara otomatis mendeteksi dan melanjutkan dari checkpoint terbaru di save_dir. checkpoint_path juga bisa menerima jalur lokal atau ID repo HuggingFace (misalnya, lillian039/ELF-B-OpenWebText). Untuk sebuah direktori, sistem akan menggunakan checkpoint terbaru di dalamnya.

Bobot encoder T5 (flax_params.msgpack) disimpan secara terpisah sebagai file .msgpack dan dimuat sekali saat startup. Bobot ini juga bisa ditentukan sebagai jalur HuggingFace (default: lillian039/t5-small-encoder).

Lisensi dan Penghargaan

Repo ini berada di bawah lisensi MIT. Kamu bisa melihat file LICENSE untuk detail lebih lanjut. Jika kamu merasa karya ini berguna dalam penelitianmu, pertimbangkan untuk mengutip makalahnya. Kami juga berterima kasih kepada Google TPU Research Cloud (TRC) atas pemberian akses TPU. Kami berharap karya ini akan menjadi sumber daya yang berguna bagi komunitas open-source.