MoE dan Mamba bekerjasama untuk menskalakan model ruang negeri kepada berbilion parameter

Model Angkasa Negeri (SSM) adalah teknologi yang telah menarik perhatian ramai dan dianggap sebagai alternatif kepada Transformer. Berbanding dengan Transformer, SSM boleh mencapai penaakulan masa linear apabila memproses tugas konteks yang panjang, dan mempunyai latihan selari dan prestasi cemerlang. Khususnya, Mamba, yang berasaskan SSM terpilih dan reka bentuk peka perkakasan, telah menunjukkan prestasi cemerlang dan telah menjadi salah satu alternatif yang berkuasa kepada seni bina Transformer berasaskan perhatian.

Baru-baru ini, penyelidik juga meneroka menggabungkan SSM dan Mamba dengan kaedah lain untuk mencipta seni bina yang lebih berkuasa. Contohnya, Machine Heart pernah melaporkan bahawa "Mamba boleh menggantikan Transformer, tetapi ia juga boleh digunakan dalam kombinasi."

Baru-baru ini, pasukan penyelidik Poland mendapati bahawa jika SSM digabungkan dengan sistem pakar hibrid (MoE/Campuran Pakar), SSM boleh dijangka mencapai pengembangan berskala besar. MoE ialah teknologi yang biasa digunakan untuk memanjangkan Transformer Sebagai contoh, model Mixtral baru-baru ini menggunakan teknologi ini. Hasil penyelidikan yang diberikan oleh pasukan penyelidik Poland ini ialah MoE-Mamba, model yang menggabungkan Mamba dan lapisan pakar hibrid.

Alamat kertas: https://arxiv.org/pdf/2401.04081.pdf

MoE-Mamba boleh meningkatkan kecekapan SSM dan KPM pada masa yang sama. Dan pasukan juga mendapati bahawa MoE-Mamba berkelakuan boleh diramal apabila bilangan pakar berbeza-beza.

Pasukan menjalankan demonstrasi eksperimen dan keputusan menunjukkan bahawa berbanding dengan Mamba, MoE-Mamba memerlukan 2.2 kali lebih sedikit langkah latihan dengan keperluan prestasi yang sama, menunjukkan prestasi kaedah baharu berbanding potensi kelebihan Transformer dan Transformer-MoE. Keputusan awal ini juga menunjukkan arah penyelidikan yang menjanjikan: SSM mungkin boleh berskala kepada berpuluh bilion parameter.

Penyelidikan Berkaitan

Model Angkasa Negeri

Model Angkasa Negeri (SSM) ialah sejenis seni bina yang digunakan untuk pemodelan jujukan. Idea untuk model ini berasal dari bidang sibernetik dan boleh dilihat sebagai gabungan RNN dan CNN. Walaupun mereka mempunyai kelebihan yang besar, mereka juga mempunyai beberapa masalah yang menghalang mereka daripada menjadi seni bina yang dominan untuk tugas pemodelan bahasa. Walau bagaimanapun, penemuan penyelidikan baru-baru ini telah membolehkan SSM mendalam menskalakan kepada berbilion-bilion parameter sambil mengekalkan kecekapan pengiraan dan prestasi yang kukuh.

Mamba

Mamba ialah model yang dibina pada SSM, yang boleh mencapai kelajuan inferens masa linear (untuk panjang konteks), dan ia juga mencapai proses latihan yang cekap melalui reka bentuk yang menyedari perkakasan. Mamba menggunakan pendekatan imbasan selari cekap kerja yang mengurangkan kesan jujukan gelung, manakala operasi GPU yang digabungkan menghapuskan keperluan untuk melaksanakan keadaan lanjutan. Keadaan pertengahan yang diperlukan untuk perambatan belakang tidak disimpan tetapi dikira semula semasa hantaran ke belakang, dengan itu mengurangkan keperluan memori. Kelebihan Mamba berbanding mekanisme perhatian amat ketara dalam peringkat inferens kerana ia bukan sahaja mengurangkan kerumitan pengiraan, tetapi juga penggunaan memori tidak bergantung pada panjang konteks.

Mamba boleh menyelesaikan pertukaran asas antara kecekapan dan keberkesanan model jujukan, yang menyerlahkan kepentingan pemampatan keadaan. Model yang cekap mesti memerlukan keadaan kecil, dan keadaan yang diperlukan oleh model yang berkesan harus mengandungi semua maklumat utama konteks. Tidak seperti SSM lain yang memerlukan invarian temporal dan input, Mamba memperkenalkan mekanisme pemilihan yang mengawal cara maklumat disebarkan sepanjang dimensi jujukan. Pilihan reka bentuk ini diilhamkan oleh pemahaman intuitif tentang tugas sintesis kelas pertama seperti replikasi terpilih dan aruhan, yang membolehkan model membezakan dan mengekalkan maklumat kritikal sambil menapis maklumat yang tidak berkaitan.

Penyelidikan mendapati bahawa Mamba mempunyai keupayaan untuk menggunakan konteks yang lebih panjang dengan cekap (sehingga token 1M), dan apabila panjang konteks bertambah, kebingungan pra-latihan juga akan bertambah baik. Model Mamba terdiri daripada blok Mamba bertindan dan telah mencapai hasil yang sangat baik dalam pelbagai bidang seperti NLP, genomik, audio, dll. Prestasinya setanding dan mengatasi model Transformer sedia ada. Oleh itu, Mamba telah menjadi model calon yang kukuh untuk model tulang belakang pemodelan jujukan umum Sila rujuk "Lima kali pemprosesan, prestasi menyeluruh mengelilingi Transformer: Seni bina baharu Mamba meletupkan bulatan AI".

Pakar Campuran

Teknik Pakar Campuran (KPM) boleh meningkatkan bilangan parameter model tanpa menjejaskan FLOP yang diperlukan untuk inferens dan latihan model. MoE pertama kali dicadangkan oleh Jacobs et al pada tahun 1991 dan mula digunakan untuk tugasan NLP pada tahun 2017 oleh Shazeer et al.

MoE mempunyai kelebihan: pengaktifan sangat jarang - untuk setiap token yang diproses, hanya sebahagian kecil daripada parameter model digunakan. Disebabkan keperluan pengiraannya, lapisan hadapan dalam Transformer telah menjadi sasaran standard untuk beberapa teknik KPM.

Komuniti penyelidik telah mencadangkan pelbagai kaedah untuk menyelesaikan masalah teras KPM, iaitu proses pemberian token kepada pakar, juga dikenali sebagai proses routing. Pada masa ini terdapat dua algoritma penghalaan asas: Pilihan Token dan Pilihan Pakar. Yang pertama adalah untuk menghalakan setiap token ke bilangan (K) pakar tertentu, manakala yang terakhir adalah untuk menghalakan setiap token ke bilangan pakar yang tetap.

Fedus et al mencadangkan Suis dalam kertas 2022 "Tukar transformer: Menskalakan kepada trilion model parameter dengan kesederhanaan yang mudah dan cekap", yang merupakan seni bina Pilihan Token yang mengarahkan setiap token kepada pakar tunggal ( K=1), dan mereka menggunakan kaedah ini untuk berjaya mengembangkan saiz parameter Transformer kepada 1.6 trilion. Pasukan di Poland ini juga menggunakan reka bentuk MoE ini dalam eksperimen mereka.

Baru-baru ini, KPM juga telah mula memasuki komuniti sumber terbuka, seperti OpenMoE.

Alamat projek: https://github.com/XueFuzhao/OpenMoE

Adalah wajar dinyatakan bahawa sumber terbuka Mistral Mixtral 8×7B mempunyai prestasi yang setanding dengan LLaMa 2 70B, sementara memerlukan inferens pengiraan hanya kira-kira satu perenam daripada yang terakhir.

Seni Bina Model

Walaupun mekanisme asas utama Mamba agak berbeza daripada mekanisme perhatian yang digunakan dalam Transformer, Mamba mengekalkan struktur peringkat tinggi berasaskan modul model Transformer. Menggunakan paradigma ini, satu atau lebih lapisan modul yang sama disusun di atas satu sama lain, dan output setiap lapisan ditambah kepada aliran baki, lihat Rajah 2. Nilai akhir aliran baki ini kemudiannya digunakan untuk meramalkan token seterusnya untuk tugas pemodelan bahasa.

MoE-Mamba mengambil kesempatan daripada keserasian kedua-dua seni bina ini. Seperti yang ditunjukkan dalam Rajah 2, dalam MoE-Mamba, setiap lapisan Mamba selang digantikan oleh lapisan suapan MoE berasaskan Suis.

Walau bagaimanapun, pasukan juga menyedari bahawa reka bentuk ini agak serupa dengan reka bentuk "Mamba: Pemodelan jujukan masa linear dengan ruang keadaan terpilih" secara bergantian menyusun lapisan Mamba dan lapisan suapan, tetapi hasilnya Model ini sedikit lebih rendah daripada Mamba yang ringkas. Reka bentuk ini dilambangkan sebagai Mamba-MLP dalam Rajah 1.

MoE-Mamba memisahkan pemprosesan tanpa syarat bagi setiap token yang dilakukan oleh lapisan Mamba dan pemprosesan bersyarat yang dilakukan oleh lapisan MoE dengan cekap boleh menyepadukan keseluruhan konteks jujukan ke dalam perwakilan dalaman, manakala pemprosesan bersyarat; Pemprosesan boleh menggunakan pakar yang paling berkaitan untuk setiap token. Idea pemprosesan bersyarat dan tanpa syarat berselang-seli ini telah digunakan dalam beberapa model berasaskan MoE, tetapi ia biasanya bergantian lapisan asas dan suapan MoE.

Keputusan utama

Persediaan latihan

Pasukan membandingkan 5 persediaan berbeza: Transformer Asas, Mamba, Mamba-MLP, KPM dan MoE-Mamba.

Dalam kebanyakan Transformer, lapisan suapan hadapan mengandungi parameter 8dm², manakala kertas Mamba menjadikan Mamba lebih kecil (kira-kira 6dm²), supaya bilangan parameter dua lapisan Mamba adalah sama dengan satu lapisan suapan dan satu lapisan perhatian Ia kelihatan hampir yang sama. Untuk mendapatkan kira-kira bilangan parameter aktif yang sama bagi setiap token dalam Mamba dan model baharu, pasukan mengurangkan saiz setiap lapisan hadapan pakar kepada 6dm². Kecuali untuk lapisan benam dan menyahbenam, semua model menggunakan kira-kira 26 juta parameter bagi setiap token. Proses latihan menggunakan 6.5 bilion token dan bilangan langkah latihan ialah 100k.

Set data yang digunakan untuk latihan ialah set data Bahasa Inggeris C4, dan tugasnya adalah untuk meramalkan token seterusnya. Teks ditandakan menggunakan tokenizer GPT2. Jadual 3 memberikan senarai lengkap hiperparameter.

Keputusan

Jadual 1 memberikan keputusan latihan. MoE-Mamba menunjukkan prestasi yang lebih baik daripada model Mamba biasa.

Adalah diperhatikan bahawa MoE-Mamba mencapai tahap keputusan yang sama seperti Mamba biasa dalam hanya 46% daripada langkah latihan. Memandangkan kadar pembelajaran diselaraskan untuk Mamba biasa, boleh dijangkakan bahawa jika proses latihan dioptimumkan untuk MoE-Mamba, MoE-Mamba akan berprestasi lebih baik.

Kajian Ablasi

Untuk menilai sama ada skala Mamba serta bilangan pakar bertambah, penyelidik membandingkan model menggunakan bilangan pakar yang berbeza.

Rajah 3 menunjukkan langkah latihan semasa menggunakan bilangan pakar yang berbeza.

Jadual 2 memberikan keputusan selepas 100k langkah.

Keputusan ini menunjukkan bahawa kaedah yang baru dicadangkan berskala baik dengan bilangan pakar. Jika bilangan pakar adalah 8 atau lebih, prestasi akhir model baharu adalah lebih baik daripada Mamba biasa. Memandangkan Mamba-MLP lebih teruk daripada Mamba biasa, boleh dijangkakan MoE-Mamba yang menggunakan sebilangan kecil pakar akan berprestasi lebih teruk daripada Mamba. Kaedah baharu itu memberikan hasil terbaik apabila bilangan pakar ialah 32 orang.

Atas ialah kandungan terperinci MoE dan Mamba bekerjasama untuk menskalakan model ruang negeri kepada berbilion parameter. Untuk maklumat lanjut, sila ikut artikel berkaitan lain di laman web China PHP!