Pengembangan Large Language Models (LLM) yang dibangun dari decoder-only transformer models telah memainkan peran penting dalam mengubah domain Natural Language Processing (NLP), serta memajukan berbagai aplikasi deep learning termasuk reinforcement learning, analisis time-series, pemrosesan gambar, dan banyak lagi. Namun, meskipun skalabilitas dan kinerja yang kuat, LLM yang dibangun dari decoder-only transformer models masih menghadapi kelemahan signifikan. Meskipun ekspresif, mekanisme perhatian dalam transformer-derived LLM memerlukan sumber daya komputasi yang tinggi selama inferensi dan pelatihan, memerlukan memori yang besar untuk panjang urutan dan operasi floating-point (FLOPs) yang meningkat secara kuadrat. Kebutuhan komputasi yang tinggi ini membatasi panjang konteks model, meningkatkan biaya tugas generasi autoregresif seiring dengan skala model, dan menghambat kemampuan model untuk belajar dari aliran data kontinu dan memproses urutan dengan panjang tak terbatas secara efisien.

Belakangan ini, State Space Models (SSM) telah menunjukkan kemampuan dan kinerja yang luar biasa, bersaing dengan model arsitektur transformer dalam benchmark pemodelan skala besar sambil mencapai kompleksitas memori sebagai fungsi dari panjang urutan dan waktu linier. Selain itu, Mamba, sebuah State Space Model yang baru dirilis, telah menunjukkan kinerja yang luar biasa dalam berbagai tugas pemodelan bahasa dan pemrosesan urutan panjang. Bersamaan dengan itu, Mixture of Expert (MoE) models juga telah menunjukkan kinerja yang mengesankan sambil mengurangi latency dan biaya komputasi inferensi, meskipun dengan biaya kaki yang lebih besar. Dengan membangun pada Mamba dan MoE models, artikel ini akan membahas BlackMamba, sebuah arsitektur baru yang menggabungkan Mamba State Space Model dengan MoE models untuk memanfaatkan kelebihan yang ditawarkan oleh kedua kerangka kerja ini. Eksperimen pada BlackMamba telah menunjukkan kemampuan untuk mengungguli kerangka kerja Mamba yang ada dan baseline transformer dalam pelatihan FLOPs dan inferensi. Kinerja luar biasa dari kerangka kerja BlackMamba menunjukkan bahwa dapat menggabungkan kemampuan Mamba dan MoE dengan sangat baik, menawarkan inferensi yang cepat dan efisien dari MoE dengan generasi linier-kompleks dari Mamba.

Artikel ini bertujuan untuk membahas kerangka kerja BlackMamba secara mendalam. Kami menjelajahi mekanisme, metodologi, dan arsitektur kerangka kerja ini, bersama dengan perbandingannya dengan kerangka kerja generasi gambar dan video yang ada. Mari kita mulai.

BlackMamba : Pengenalan MoE untuk State Space Models

Perkembangan Large Language Models (LLM), khususnya yang berbasis arsitektur transformer decoder-only, telah mempengaruhi bidang Natural Language Processing (NLP) dan memperluas ke berbagai aplikasi deep learning, termasuk reinforcement learning, analisis time-series, pemrosesan gambar, dan lain-lain. Namun, meskipun skalabilitas dan kinerja yang kuat, LLM yang dibangun dari decoder-only transformer models masih menghadapi tantangan signifikan. Mekanisme perhatian, yang merupakan fitur kunci dari transformer-based LLM, memerlukan sumber daya komputasi yang luas untuk inferensi dan pelatihan. Ini melibatkan kebutuhan memori yang tumbuh dengan panjang urutan dan operasi floating-point (FLOPs) yang meningkat secara kuadrat. Kebutuhan komputasi yang tinggi ini membatasi panjang konteks model, meningkatkan biaya tugas generasi autoregresif seiring dengan skala model, dan menghambat kemampuan model untuk belajar dari aliran data kontinu dan memproses urutan dengan panjang tak terbatas secara efisien.

Upaya signifikan telah dilakukan dalam beberapa tahun terakhir untuk mengatasi keterbatasan ini, dan perhatian telah beralih ke pengembangan alternatif arsitektur untuk model transformer yang padat, dengan State Space Models dan MoE models menjadi arsitektur kandidat paling menjanjikan. Kelebihan utama yang diperoleh dengan memilih State Space Models daripada model arsitektur transformer adalah kompleksitas komputasi linier terhadap panjang urutan input yang ditawarkan oleh SSM, bukan kompleksitas kuadrat yang ditawarkan oleh transformer. Secara teoretis, kompleksitas komputasi linier terhadap panjang urutan input memungkinkan State Space Models untuk memproses urutan yang lebih panjang daripada model arsitektur transformer untuk anggaran FLOPS atau Floating-point operations per second yang sama, dan untuk membuat generasi autoregresif konstan dalam komputasi tanpa cache KV. State Space Models yang baru dikembangkan, termasuk Mamba, RetNet, dan beberapa lainnya, telah menunjukkan inferensi dan pelatihan urutan panjang yang efisien, serta kinerja tugas pemodelan bahasa yang kompetitif dengan transformer dengan sifat penskalaan yang sama. Di sisi lain, arsitektur Mixture of Expert telah menjadi populer sebagai alternatif untuk transformer yang padat, karena memungkinkan pengurangan signifikan dalam FLOPs inferensi dan pelatihan yang diperlukan untuk mencapai kualitas yang setara dengan model yang padat. MoE (Mixture of Experts) models bekerja dengan mengaktifkan hanya seleksi parsial dari total parameter selama satu pass maju. Mereka menggunakan fungsi routing untuk menentukan ‘ahli’ mana yang dipanggil berdasarkan konteks yang diberikan. Pendekatan ini menciptakan pemisahan antara biaya komputasi inferensi dan jumlah total parameter, memungkinkan kinerja yang ditingkatkan dalam anggaran inferensi yang tetap, meskipun dengan kaki yang lebih besar dan kebutuhan memori yang lebih besar.

Kemajuan ini menawarkan kelebihan signifikan dibandingkan dengan transformer tradisional dan mewakili arah yang menarik untuk pengembangan lebih lanjut. Kami berpendapat bahwa mengintegrasikan peningkatan ini ke dalam model Mamba-MoE yang tergabung dapat secara signifikan mempercepat kemampuan pemodelan bahasa dan efisiensi di luar model transformer standar. Kelebihan yang diharapkan dari arsitektur Mamba-MoE dibandingkan dengan model transformer yang padat termasuk:

Mamba: Mencapai kompleksitas komputasi linier relatif terhadap panjang urutan input untuk fase pelatihan dan inferensi. Ini memungkinkan generasi autoregresif terjadi dalam waktu konstan dan dengan penggunaan memori yang konstan.

MoE: Menawarkan kecepatan inferensi dan efisiensi komputasi pelatihan yang setara dengan model dasar yang lebih kecil, sambil mempertahankan tingkat kualitas model yang setara dengan model yang memiliki jumlah parameter yang sama dengan versi yang lebih padat.

Dengan demikian, perlu dinyatakan bahwa model arsitektur transformer masih merupakan state of the art, dan telah menunjukkan kinerja yang konsisten dan luar biasa pada tugas pemodelan bahasa dan pemrosesan urutan. Pada intinya, arsitektur transformer menggunakan perhatian diri yang melakukan perbandingan semua-ke-semua dari kesamaan produk dot antara embedding token yang berbeda dalam sebuah urutan, dan melakukan pemetaan linier ke vektor output. Model transformer terdiri dari blok perhatian diri yang disusun di antara blok MLP atau Multilayer Perceptron yang lebih lanjut terdiri dari dua lapis MLP dengan fungsi aktivasi yang diberikan.

BlackMamba : Arsitektur dan Metodologi

State Space Models

State Space Models termasuk dalam kelompok model urutan dengan kompleksitas linier terhadap panjang urutan input. Arsitektur State Space Models lebih sejalan dengan Jaringan Saraf Berulang dan Jaringan Saraf Konvolusional daripada arsitektur berbasis perhatian, dan terinspirasi dari sistem dinamis kontinu yang memetakan fungsi satu dimensi melalui ruang laten implisit. Sistem dinamis linier membuat komputasi paralel efisien menggunakan scan asosiatif atau konvolusional. Dalam skenario praktis, sifat berulang dari State Space Models telah menjadi alasan mengapa masih belum diadopsi pada perangkat keras AI yang sangat paralel seperti GPU. Namun, munculnya SSM seperti RWKV dan Mamba telah menggunakan kernel scan paralel untuk memetakan operasi berulang dengan efisiensi yang setara dengan model transformer, sehingga memungkinkan pelatihan arsitektur baru dengan efisiensi yang setara dengan model transformer.

Keterbatasan kuadrat yang melekat dalam transformer terkait dengan panjang urutan adalah keterbatasan yang diketahui dengan baik yang menghambat penalaran dan pemahaman atas konteks yang sangat panjang. Inovasi baru-baru ini telah memperkenalkan gagasan untuk memperpanjang panjang konteks, memungkinkan transformer untuk dilatih pada skala yang layak sebelum diterapkan pada konteks yang jauh lebih panjang selama inferensi. Meskipun kemajuan ini, proses inferensi masih memerlukan sejumlah besar sumber daya komputasi dan memori, terutama untuk mempertahankan cache Key-Value (KV), membuatnya menjadi upaya yang intensif sumber daya. Upaya penelitian baru-baru ini telah fokus pada meningkatkan kemampuan ekspresif dari model ruang keadaan dengan mengintegrasikan mekanisme gerbang yang bergantung pada input, serupa dengan matriks QKV yang ditemukan dalam mekanisme perhatian.

Upaya ini bertujuan untuk mempertahankan progresi linier yang melekat dalam rekursi ruang keadaan, memungkinkan eksekusi yang efisien melalui konvolusi atau proses scan selektif. Pendekatan ini secara signifikan menyempitkan kesenjangan kinerja dengan transformer dalam aplikasi praktis. Di antara kemajuan ini, Mamba menonjol sebagai model ruang keadaan yang mencerminkan tujuan penelitian sebelumnya, menunjukkan tingkat kinerja yang mengesankan yang setara dengan transformer pada skala hingga 2,8 miliar parameter. Ini dicapai dengan menerapkan gerbang yang bergantung pada input ke input rekursi model ruang keadaan (SSM), sambil memastikan komputasi yang efisien melalui penggunaan kernel scan selektif khusus.

Mixture of Expert Models

Mixture of Expert (MoE) models mencapai pemisahan antara biaya inferensi dan jumlah total parameter dengan mengaktifkan parameter secara selektif selama pass maju. Alih-alih menggunakan semua parameter, model ini mengarahkan token ke ahli Multilayer Perceptron (MLP) tertentu. Idealnya, setiap ahli dirancang untuk memproses jenis input tertentu, dengan mekanisme routing, yang pada dasarnya adalah jaringan neural kompak, menentukan ahli yang paling sesuai untuk setiap token. Pendekatan ini bertujuan untuk mempertahankan kekuatan ekspresif model yang setara dengan jumlah parameter yang sama dalam konfigurasi yang lebih padat, tetapi dengan biaya komputasi yang jauh lebih rendah. Biasanya, router adalah pemetaan lapisan linier dari token ke indeks ahli dengan setiap ahli menjadi Multilayer Perceptron standar transformer. Namun, pengembang masih belum menemukan metode pelatihan router yang optimal karena masalah penugasan ahli yang tidak dapat dibedakan, dan model Mixture of Expert sering mengalami kesulitan dengan keseimbangan beban dan stabilitas pelatihan antara ahli yang berbeda untuk efisiensi perangkat keras.

Arsitektur

Pada intinya, BlackMamba menggunakan model transformer standar yang terdiri dari blok MLP dan blok perhatian yang disusun secara berurutan di sepanjang aliran residual. Sekarang, sebagian besar model Mixture of Expert hanya menggantikan blok multilayer perceptron dengan lapisan ahli yang diarahkan. Di sisi lain, kerangka kerja BlackMamba tidak hanya menggantikan blok multilayer perceptron dalam transformer dengan lapisan ahli yang diarahkan, tetapi juga menggantikan lapisan perhatian dengan lapisan Model Ruang Keadaan Mamba. Arsitektur kerangka kerja BlackMamba ditunjukkan dalam gambar berikut.

Pelatihan dan Dataset

Model BlackMamba dilatih pada lebih dari 300 miliar token pada dataset khusus, dan menggunakan fungsi aktivasi SwiGLU untuk multilayer perceptron ahli. Kerangka kerja ini dilatih dengan 8 ahli, sebuah jumlah yang ditemukan oleh pengembang sebagai keseimbangan yang tepat antara jejak memori dan biaya inferensi model. Dataset khusus yang digunakan untuk melatih kerangka kerja BlackMamba terdiri dari campuran dataset sumber terbuka yang ada, termasuk Starcoder, SlimPajama, Pile, dan lain-lain. Tabel berikut menunjukkan bobot dari setiap dataset yang digunakan untuk melatih kerangka kerja BlackMamba. Secara keseluruhan, ada 1,8 triliun token dalam dataset.

BlackMamba : Hasil

Untuk memastikan perbandingan yang adil antara Mamba dan BlackMamba, pengembang telah melatih kedua model dengan parameter pelatihan yang sama pada data pelatihan yang sama. Kerangka kerja BlackMamba dapat mengungguli kedua model Mamba dan transformer untuk ukuran model pass maju yang identik pada waktu inferensi serta pelatihan Floating-point operations per second. Gambar berikut menunjukkan waktu yang dibutuhkan untuk menghasilkan urutan dengan panjang tertentu secara autoregresif dari prompt awal satu token sebagai fungsi dari panjang urutan.

Selain itu, kelebihan latency dari model Mixture of Expert dan Mamba digabungkan dalam kerangka kerja BlackMamba, menghasilkan waktu inferensi yang jauh lebih cepat dibandingkan dengan model transformer, model Mamba murni, dan model MoE. Selain itu, kelebihan inferensi dari kerangka kerja BlackMamba secara langsung sebanding dengan panjang urutan, membuat BlackMamba sangat efektif untuk generasi urutan panjang. Berikutnya, gambar berikut menunjukkan jumlah token yang ditugaskan ke model BlackMamba dengan 340 juta dan 640 juta parameter, masing-masing. Seperti yang dapat dilihat, sebagian besar lapisan menunjukkan keseimbangan ahli yang tinggi sebagai hasil dari algoritma Sinkhorn yang ditingkatkan yang diimplementasikan oleh model BlackMamba.

Tabel berikut mencakup skor evaluasi kerangka kerja BlackMamba dibandingkan dengan berbagai model bahasa pra-dilatih sumber terbuka. Seperti yang dapat dilihat, kerangka kerja BlackMamba dapat bersaing dan mengungguli sebagian besar kerangka kerja di semua baseline. Selain itu, perlu diperhatikan bahwa model yang mengungguli BlackMamba memiliki jumlah parameter yang jauh lebih besar, dan kesenjangan kinerja minimal, menunjukkan kemampuan kerangka kerja BlackMamba dengan parameter yang lebih sedikit.

Pemikiran Akhir

Dalam artikel ini, kita telah membahas BlackMamba, sebuah arsitektur baru yang menggabungkan Mamba State Space Model dengan MoE models untuk memanfaatkan kelebihan yang ditawarkan oleh kedua kerangka kerja ini. Eksperimen pada BlackMamba telah menunjukkan kemampuan untuk mengungguli kerangka kerja Mamba yang ada dan baseline transformer dalam pelatihan FLOPs dan inferensi. Kinerja luar biasa dari kerangka kerja BlackMamba menunjukkan bahwa dapat menggabungkan kemampuan Mamba dan MoE dengan sangat baik, menawarkan inferensi yang cepat dan efisien dari MoE dengan generasi linier-kompleks dari Mamba. Kita telah membahas bagaimana arsitektur kerangka kerja BlackMamba dapat mengungguli model bahasa besar yang dilatih dengan kuat, kerangka kerja Mamba yang ada, dan model Mixture of Expert dalam hal pelatihan FLOPs dan biaya inferensi. Selain itu, kerangka kerja BlackMamba juga mewarisi generasi FLOPs dan pelatihan yang berkurang dari model Mixture of Expert dan Mamba secara bersamaan.