[Terjemahan] MoD (Mixture-of-Depths): Pendekatan untuk Optimasi Komputasi Model Bahasa Berbasis Transformer, dan MoDE (MoD+MoE)

(discuss.pytorch.kr)

7 poin oleh ninebow 2024-04-13 | Belum ada komentar. | Bagikan ke WhatsApp

PyTorchKR

Metode MoE (Mixture-of-Experts), yang memproses token input dengan memilih satu atau lebih di antara beberapa blok "expert" seperti Jamba, Qwen1.5-MoE, atau DBRX, belakangan ini semakin banyak digunakan. Pendekatan MoE ini menunjukkan performa yang lebih baik sambil mengalokasikan sumber daya komputasi secara efisien dengan mengatur sebagian dari LLM besar secara dinamis.
Jika MoE menyesuaikan "lebar (width) model" secara dinamis dengan memilih sebagian dari banyak expert, maka MoD (Mixture-of-Depth) yang diperkenalkan kali ini mengambil pendekatan untuk menyesuaikan "kedalaman (depth) model" secara dinamis. Artinya, alih-alih menjalankan komputasi pada seluruh layer model untuk token tertentu, metode ini menghitung secara dinamis dengan melewati sebagian layer. Mari kita lihat bersama.

MoD (Mixture-of-Depths): Pendekatan untuk Optimasi Komputasi Model Bahasa Berbasis Transformer (Mixture-of-Depths: Dynamically allocating compute in transformer-based language models)

Pengantar

Secara tradisional, model transformer mengalokasikan sumber daya komputasi yang sama ke semua bagian dari sekuens input. Namun, paper MoD (Mixture-of-Depth) Transformer yang diperkenalkan kali ini mencoba mengalokasikan sumber daya komputasi secara dinamis ke posisi sekuens tertentu dalam Transformer, dan mengubah alokasi tersebut di layer yang berbeda. Untuk menggunakan strategi Mixture-of-Depths (MoD) ini, penulis mengelola total beban komputasi dengan menggunakan mekanisme routing top-k untuk menentukan token mana yang akan diproses. Teknik ini menandingi atau melampaui performa model yang ada sambil membutuhkan FLOPs yang jauh lebih sedikit, serta memberikan peningkatan kecepatan hingga 50% saat sampling setelah pelatihan.

Arsitektur Transformer yang diperkenalkan dalam paper Attention is All You Need

Model transformer terbaru dan LLM (Large Language Model) terus bertambah dalam skala dan kompleksitas. Model-model ini menggunakan miliaran parameter untuk mendekati kemampuan pemahaman dan generasi bahasa setingkat manusia, dan sebagai hasilnya menunjukkan performa yang mengesankan. Namun, kemajuan ini disertai biaya komputasi dan konsumsi energi yang besar, dan peningkatan skala model secara drastis menambah waktu pelatihan serta inferensi.

Hal ini terjadi karena Transformer melakukan jumlah komputasi yang sama untuk semua token pada sekuens input. Akan tetapi, ide dasarnya berangkat dari kenyataan bahwa tidak semua token memberikan jumlah informasi yang setara: beberapa token memainkan peran kunci dalam memahami konteks, sementara yang lain relatif kurang penting.

MoD Transformer memungkinkan model memfokuskan sumber daya pada bagian yang benar-benar diperlukan dengan mengalokasikan komputasi secara dinamis untuk setiap token. Ini dapat mengurangi biaya komputasi keseluruhan model sekaligus meningkatkan efisiensi pemrosesan informasi dengan memberi perhatian lebih besar pada token penting. Khususnya pada model besar seperti LLM, pendekatan MoD dapat menjadi cara penting untuk menghemat waktu dan sumber daya dalam proses pelatihan maupun inferensi.

Pengenalan teknik MoD (Mixture-of-Depths)

Teknik MoD mengintegrasikan mekanisme routing top-k yang secara dinamis memilih token mana yang akan diproses pada tiap layer transformer. Dengan demikian, hanya sebagian token yang terpilih yang menjalani seluruh proses komputasi, sementara token lainnya melewati satu atau beberapa layer.

Gambar 1 / Kiri: Gambaran umum MoD (Mixture-of-Depths) Transformer: seperti struktur MoE (Mixture-of-Experts), jalur komputasi dipilih melalui router. Bedanya dengan MoE, di sini yang diputuskan adalah apakah blok standar (Self-Attention dan MLP) akan dihitung atau tidak. Jika token input ($X_i$) tidak mengaktifkan komputasi penuh sesuai keputusan router, token tersebut tidak akan melewati komputasi pada layer itu.
Gambar 1 / Kanan: Hasil keputusan routing untuk sekuens pendek yang diperkecil menjadi 64 token. (Sumbu X: sekuens, sumbu Y: layer) Warna ungu menunjukkan komputasi layer tersebut dijalankan penuh, sedangkan oranye menunjukkan komputasi layer tersebut dilewati.

Seperti dijelaskan pada gambar di atas, MoD mengimplementasikan router prediktif di dalam arsitektur transformer untuk menentukan apakah komputasi layer akan dijalankan untuk tiap token. Melalui cara ini, sumber daya komputasi dapat difokuskan pada token yang relatif "penting".

Pelatihan MoD Transformer

Proses pelatihan MoD Transformer pada dasarnya mirip dengan transformer standar. Sekuens input dimasukkan ke model, lalu diteruskan melalui setiap layer sambil menjalankan komputasi yang diperlukan. Namun, MoD Transformer mencakup langkah tambahan berupa router yang menentukan jalur komputasi mana yang akan diikuti tiap token.

Perhitungan bobot router: Router menghitung bobot untuk setiap token pada setiap layer model. Bobot ini menunjukkan jumlah komputasi yang harus diterima token tersebut, berdasarkan tingkat kepentingannya.
Pemilihan token teratas: Berdasarkan bobot yang dihitung, router memilih token-token teratas untuk setiap layer. Token-token ini diarahkan ke jalur yang menjalankan komputasi penuh.
Alokasi jalur komputasi: Token yang dipilih dialokasikan ke jalur yang menjalankan komputasi standar (Self-Attention dan MLP). Token lainnya diteruskan melalui residual connection dan berpindah ke layer berikutnya tanpa komputasi tambahan.
Perhitungan loss dan backpropagation: Output model dievaluasi melalui fungsi loss akhir, lalu gradient dipropagasikan balik melalui model untuk memperbarui bobot. Dalam proses ini, router juga dilatih sehingga seiring waktu dapat membuat keputusan routing yang lebih efisien.

Skema routing MoD Transformer (Routing Schemes)

Implementasi routing adalah bagian inti dari MoD Transformer. Penulis membandingkan tiga skema routing dan menemukan bahwa routing token dengan pendekatan Expert-Choice MoD menunjukkan performa terbaik. Mari kita bandingkan dengan skema routing lain yang diperkenalkan:

Routing pilihan token (Token-choice)

Pada routing pilihan token, setiap token memilih sendiri jalur yang akan dilewatinya. Model menghasilkan distribusi probabilitas untuk berbagai jalur komputasi bagi setiap token (misalnya layer atau blok pemrosesan yang berbeda), lalu mengarahkan token ke jalur tersebut sesuai distribusi itu.

Routing pilihan token memungkinkan setiap token memilih jalur komputasi yang paling sesuai, sehingga routing menjadi sangat fleksibel. Selain itu, model dapat mempertimbangkan tingkat kepentingan atau konteks tiap token dengan lebih presisi, sehingga optimasi per token menjadi memungkinkan.

Routing pilihan expert (Expert-choice)

Pada routing pilihan expert, model secara langsung memilih token yang akan diproses oleh tiap jalur (atau "expert"). Ini dilakukan dengan mekanisme top-k yang menyaring token-token teratas berdasarkan bobot router tiap token.

Dengan routing pilihan expert seperti ini, token dapat didistribusikan secara merata ke tiap jalur, sehingga beban komputasi tersebar lebih seimbang dan sumber daya komputasi model dapat dimanfaatkan lebih efisien. Selain itu, karena token teratas ditentukan dengan jelas, prediktabilitas dan stabilitas model juga meningkat.

Expert-choice MoD

Expert-choice MoD adalah penerapan konsep routing pilihan expert pada struktur MoD Transformer. Dalam pendekatan ini, model memilih token-token teratas untuk setiap layer atau tahap pemrosesan, dan hanya token yang dipilih itulah yang diproses melalui jalur yang menjalankan komputasi. Token sisanya mengikuti jalur berbiaya rendah melalui residual connection.

Karena hanya token terpilih yang menjalankan komputasi, total penggunaan FLOP model berkurang secara signifikan. Selain itu, pendekatan ini dapat mempertahankan atau meningkatkan performa model dengan menjamin komputasi terhadap seluruh layer untuk token penting, sambil meminimalkan komputasi yang tidak perlu. Terakhir, melalui routing dinamis oleh router, model memiliki keunggulan untuk mengalokasikan sumber daya komputasi secara fleksibel sesuai karakteristik sekuens input.

Cara implementasi routing

Implementasi routing dapat dilakukan melalui tiga langkah utama berikut:

Embedding token dan penentuan bobot router: Untuk setiap layer model, terdapat himpunan embedding token $X^i$ dari sekuens yang diberikan. Untuk setiap embedding token $x^l_i$, router menghitung bobot skalar $r^l_i$ melalui transformasi linear yang menunjukkan kemungkinan token tersebut ikut serta dalam komputasi.
Menentukan bobot teratas: Dari himpunan bobot router token $R^l$, dicari probabilitas persentil ke-$\beta$, yaitu $P_\beta(R^l)$. Di sini, $\beta$ menentukan rasio token yang akan diproses model pada tiap layer. Karena itu, token dengan bobot router lebih besar daripada $P_\beta(R^l)$ akan ikut dalam komputasi.
Pemilihan jalur komputasi: Setiap token dibandingkan dengan bobot router $r^l_i$; jika lebih besar atau sama dengan $P_\beta(R^l)$ maka token tersebut diproses pada komputasi layer terkait, dan jika lebih kecil maka langsung diteruskan ke layer berikutnya. Dengan cara ini, komputasi hanya dilakukan untuk token terpilih, sementara sisanya dapat diproses melalui jalur yang tidak menimbulkan biaya komputasi.

Dalam implementasi routing ini, hal-hal berikut perlu menjadi pertimbangan utama:

Keputusan dinamis berbasis bobot: Proses routing diputuskan secara dinamis berdasarkan bobot masing-masing token, yang mencerminkan tingkat kepentingannya. Token penting menerima lebih banyak komputasi untuk mengoptimalkan performa model, sementara token yang kurang penting menghemat biaya melalui jalur residual.
Efisiensi komputasi dan optimasi performa: Melalui mekanisme routing ini, MoD Transformer dapat menggunakan sumber daya komputasi secara efisien sambil mempertahankan atau meningkatkan performa model. Khususnya, dengan hanya memproses token penting pada layer yang mahal secara komputasi, total penggunaan FLOP dapat dikurangi.
Kesederhanaan dan efisiensi implementasi: Routing diimplementasikan melalui transformasi linear dan perhitungan persentil, sehingga relatif sederhana dan efisien. Ini memungkinkan alokasi sumber daya komputasi secara dinamis tanpa menambah beban komputasi yang besar pada proses pelatihan dan inferensi.

Contoh cara kerja MoD Transformer

Berdasarkan penjelasan sejauh ini, berikut contoh cara kerja MoD Transformer.

Misalnya, ada 100 token dalam sekuens input, dan kita asumsikan router menetapkan nilai menjadi 10. Dalam kasus ini, pada setiap layer router akan memilih 10 token yang paling penting, dan hanya token-token tersebut yang diproses melalui jalur komputasi standar.

Sisa 90 token akan melewati komputasi dan langsung diteruskan ke layer berikutnya melalui residual connection. Dengan proses ini, MoD Transformer menjamin komputasi yang diperlukan untuk token penting sambil mengurangi total penggunaan FLOP.

Perbandingan performa MoD Transformer

Performa pelatihan

Optimasi hyperparameter: Melalui eksperimen pelatihan varian MoD Transformer yang menempatkan blok routing berkapasitas 12.5% secara bergantian dengan blok attention penuh, ditemukan konfigurasi hyperparameter yang optimal. Konfigurasi ini memberikan performa yang lebih baik dibanding baseline optimal dalam analisis isoFLOP.

Efisiensi komputasi: Varian MoD tertentu menunjukkan kecepatan pemrosesan 66% lebih cepat sambil mencapai performa yang sama dengan model yang ada menggunakan FLOP yang lebih sedikit.

Analisis routing

Sparsitas keputusan routing: Dengan menggunakan routing berkapasitas 12.5%, sebagian besar token melewati blok, dan hanya sejumlah kecil token penting yang melaluinya. Hal ini juga dapat diamati pada distribusi bobot router, dan jaringan kadang merutekan dengan lebih menyukai token tertentu bergantung pada kedalaman.

Performa inferensi

Inferensi autoregresif: Saat beralih dari skema routing top-k non-causal saat pelatihan ke pendekatan berbasis prediksi causal saat inferensi, hampir tidak ada penurunan performa. Hal ini tampaknya karena masalah prediksi tersebut mudah dipelajari, dan akurasi lebih dari 97% telah dicapai sejak awal pelatihan.

Mixture-of-Depths-and-Experts (MoDE) = MoD + MoE

MoDE (Mixture-of-Depths-and-Experts) berarti integrasi teknik MoD (Mixture-of-Depths) dan teknik MoE (Mixture of Experts). Pendekatan ini bertujuan menggabungkan keunggulan kedua metode tersebut untuk sekaligus mengoptimalkan efisiensi komputasi dan performa model transformer. Jika MoD memilih agar komputasi untuk tiap token dilakukan lebih dalam atau lebih dangkal, maka MoE memproses token melalui "expert" yang paling sesuai di antara berbagai expert yang tersedia. MoDE menggabungkan kedua pendekatan ini sehingga token dapat melewati blok tertentu atau secara selektif menjalankan komputasi yang lebih dalam, sambil tetap memungkinkan pemilihan cara pemrosesan terbaik di antara berbagai expert.

Cara implementasi MoDE

MoDE memiliki dua pendekatan implementasi utama: Staged MoDE dan Integrated MoDE.

Staged MoDE: Pada Staged MoDE, teknik MoD diterapkan terlebih dahulu untuk menentukan jalur komputasi setiap token, lalu teknik MoE diterapkan untuk merutekan token yang akan diproses pada jalur terpilih ke masing-masing blok expert. Dengan cara ini, model terlebih dahulu menyesuaikan kedalaman komputasi, lalu memilih expert yang optimal pada setiap tahap.
Integrated MoDE: Pada Integrated MoDE, teknik MoD dan MoE diintegrasikan dalam satu tahap routing. Artinya, saat merutekan token, satu keputusan sekaligus memilih kedalaman jalur komputasi dan expert yang akan memprosesnya. Pendekatan ini bisa lebih sederhana dan efisien untuk diimplementasikan, serta memiliki keunggulan dapat menentukan jalur komputasi dan expert yang paling sesuai untuk setiap token secara bersamaan.

Keunggulan MoDE

Efisiensi komputasi: MoDE secara dinamis menentukan jumlah komputasi yang dibutuhkan tiap token dan expert yang akan memprosesnya, sehingga biaya komputasi keseluruhan model dapat dikelola secara efektif.
Optimasi performa: Dengan menempatkan setiap token pada jalur komputasi dan expert yang paling sesuai, MoDE memaksimalkan performa model. Ini bisa sangat berguna terutama saat menyelesaikan masalah kompleks atau memproses berbagai jenis data.
Fleksibilitas dan skalabilitas: MoDE menawarkan fleksibilitas untuk diterapkan pada berbagai struktur model dan task, serta menyediakan framework yang scalable untuk meningkatkan efisiensi komputasi dan performa model secara bersamaan.

MoDE menunjukkan arah perkembangan masa depan model transformer dan memberikan kontribusi penting bagi riset yang mengeksplorasi cara memaksimalkan performa model sambil meminimalkan biaya komputasi.

Baca lebih lanjut

https://arxiv.org/abs/2404.02258

Tulisan ini disusun berdasarkan ringkasan yang dibuat dengan model GPT, sehingga mungkin ada bagian yang diringkas berbeda dari isi atau maksud naskah asli. Jika topik ini menarik bagi Anda, silakan juga merujuk ke naskah aslinya. Jika saat membaca Anda menemukan bagian yang terasa janggal atau keliru, mohon beri tahu melalui komentar.

⚠️Iklan⚠️: Apakah tulisan yang dirangkum oleh Komunitas Pengguna PyTorch Korea ini bermanfaat? Jika Anda mendaftar sebagai anggota, kami akan mengirimkan tulisan-tulisan utama melalui email! (Default-nya Weekly, tetapi bisa diubah ke Daily.)