Kerja baharu daripada pengarang Mamba: Menyuling Llama3 menjadi RNN linear hibrid

Kunci kejayaan besar Transformer dalam bidang pembelajaran mendalam ialah mekanisme perhatian. Mekanisme perhatian membolehkan model berasaskan Transformer menumpukan pada bahagian yang berkaitan dengan jujukan input, mencapai pemahaman konteks yang lebih baik. Walau bagaimanapun, kelemahan mekanisme perhatian ialah overhed pengiraan adalah tinggi, yang meningkat secara kuadratik dengan saiz input, menjadikannya sukar bagi Transformer untuk mengendalikan teks yang sangat panjang.

Beberapa masa lalu, kemunculan Mamba memecahkan keadaan ini, yang boleh mencapai pengembangan linear apabila panjang konteks meningkat. Dengan keluaran Mamba, model angkasa lepas (SSM) ini sudah boleh memadankan atau bahkan mengatasi Transformer pada skala kecil hingga sederhana, sambil mengekalkan kebolehskalaan linear dengan panjang jujukan, yang memberikan ciri-ciri penggunaan yang menguntungkan Mamba.

Ringkasnya, Mamba mula-mula memperkenalkan mekanisme pemilihan yang mudah tetapi berkesan, yang boleh meparameterkan semula SSM mengikut input, membolehkan model mengekalkan maklumat yang diperlukan selama-lamanya sambil menapis maklumat yang tidak berkaitan dan data yang berkaitan.

Baru-baru ini, kertas kerja bertajuk "The Mamba in the Llama: Distilling and Accelerating Hybrid Models" membuktikan bahawa dengan menggunakan semula pemberat lapisan perhatian, transformer besar boleh disuling menjadi RNN linear hibrid besar, hanya pengiraan tambahan minimum sambil mengekalkan kebanyakan kualiti binaannya.

Model hibrid yang terhasil, yang mengandungi satu perempat daripada lapisan perhatian, mencapai prestasi yang setanding dengan Transformer asal dalam penanda aras sembang, dan mengatasi prestasi menggunakan data dalam penanda aras sembang dan penanda aras umum Model Mamba hibrid sumber terbuka dilatih dari awal dengan trilion token. Selain itu, kajian itu mencadangkan algoritma penyahkodan spekulatif peka perkakasan yang mempercepatkan inferens untuk model Mamba dan hibrid.

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

Model berprestasi terbaik kajian ini adalah daripada Llama3-8B-Instruct Distilled , ia mencapai kadar kemenangan terkawal panjang 29.61 pada AlpacaEval 2 berbanding GPT-4, dan kadar kemenangan 7.35 pada MT-Bench, mengatasi model RNN linear terlaras arahan terbaik.

Kaedah

Penyulingan Pengetahuan (KD) ialah teknik pemampatan model yang digunakan untuk memindahkan pengetahuan daripada model besar (model guru) kepada model yang lebih kecil (model pelajar) ), yang bertujuan untuk melatih rangkaian pelajar untuk meniru tingkah laku rangkaian guru. Penyelidikan bertujuan untuk menyaring Transformer supaya prestasinya setanding dengan model bahasa asal.

Kajian ini mencadangkan kaedah penyulingan berbilang peringkat yang menggabungkan penyulingan progresif, penalaan halus diselia dan pengoptimuman keutamaan arah. Berbanding dengan penyulingan biasa, kaedah ini boleh mencapai kebingungan yang lebih baik dan keputusan penilaian hiliran.

Kajian ini mengandaikan bahawa kebanyakan pengetahuan daripada Transformer dikekalkan dalam lapisan MLP yang dipindahkan daripada model asal, dan memfokuskan pada langkah penalaan halus dan penjajaran LLM suling. Semasa fasa ini, lapisan MLP kekal beku dan lapisan Mamba dilatih.

Kajian ini percaya bahawa terdapat beberapa hubungan semula jadi antara RNN linear dan mekanisme perhatian. Formula perhatian boleh dilinearkan dengan mengalih keluar softmax:

Tetapi perhatian linear akan membawa kepada kemerosotan keupayaan model. Untuk mereka bentuk RNN ​​linear suling yang cekap, kajian ini mendekati parameterisasi Transformer asal sedekat mungkin sambil memanjangkan kapasiti RNN linear dengan cara yang cekap. Kajian ini tidak cuba untuk memastikan model baharu menangkap fungsi perhatian asal yang tepat, sebaliknya menggunakan bentuk linear sebagai titik permulaan untuk penyulingan.

Seperti yang ditunjukkan dalam Algoritma 1, kajian ini menyuapkan kepala standard Q, K, V dari mekanisme perhatian terus ke pendiskretan Mamba dan kemudian menggunakan RNN linear yang terhasil. Ini boleh dianggap sebagai menggunakan perhatian linear untuk pemulaan kasar dan membolehkan model mempelajari interaksi yang lebih kaya melalui keadaan tersembunyi lanjutan.

Kajian ini secara langsung menggantikan kepala perhatian Transformer dengan lapisan RNN linear yang diperhalusi, memastikan lapisan MLP Transformer tidak berubah dan tidak melatihnya. Pendekatan ini juga perlu mengendalikan komponen lain, seperti perhatian pertanyaan berkumpulan yang berkongsi kunci dan nilai di seluruh kepala. Pasukan penyelidik menyatakan bahawa seni bina ini, tidak seperti yang digunakan dalam banyak sistem Mamba, membenarkan permulaan ini menggantikan sebarang blok perhatian dengan blok RNN ​​linear.

Penyelidikan ini juga mencadangkan algoritma baharu untuk penyahkodan spekulatif RNN linear menggunakan penjanaan berbilang langkah yang menyedari perkakasan.

Algoritma 2 dan Rajah 2 menunjukkan algoritma lengkap. Pendekatan ini hanya menyimpan keadaan tersembunyi RNN dalam cache untuk pengesahan dan malas memajukannya berdasarkan kejayaan kernel berbilang langkah. Memandangkan model penyulingan mengandungi lapisan pengubah, kajian ini juga memanjangkan penyahkodan spekulatif kepada seni bina hibrid Attention/RNN. Dalam persediaan ini, lapisan RNN melakukan pengesahan mengikut Algoritma 2, manakala lapisan Transformer hanya melakukan pengesahan selari.

Untuk mengesahkan keberkesanan kaedah ini, kajian menggunakan Mamba 7B dan Mamba 2.8B sebagai model sasaran untuk spekulasi. Keputusan ditunjukkan dalam Jadual 1.

Rajah 3 menunjukkan ciri prestasi kernel berbilang langkah itu sendiri.

Pecutan pada GPU H100. Algoritma yang dicadangkan dalam kajian ini menunjukkan prestasi yang kukuh pada GPU Ampere, seperti ditunjukkan dalam Jadual 1 di atas. Tetapi terdapat cabaran besar pada GPU H100. Ini terutamanya kerana operasi GEMM terlalu pantas, yang menjadikan overhed yang disebabkan oleh operasi caching dan pengiraan semula lebih ketara. Sesungguhnya, pelaksanaan mudah algoritma yang dikaji (menggunakan berbilang panggilan kernel berbeza) mencapai kelajuan yang besar pada 3090 GPU, tetapi tiada kelajuan langsung pada H100.

Eksperimen dan keputusan

Kajian ini menggunakan dua model sembang LLM untuk eksperimen: Zephyr-7B diperhalusi berdasarkan model Mistral 7B dan Llama- 3 Instruct 8B. Untuk model RNN linear, kajian ini menggunakan versi hibrid Mamba dan Mamba2 dengan lapisan perhatian masing-masing 50%, 25%, 12.5%, dan 0%, dan memanggil 0% model Mamba tulen. Mamba2 ialah varian seni bina Mamba yang direka khas untuk seni bina GPU terkini.

Penilaian pada Penanda Aras Sembang

Jadual 2 menunjukkan prestasi model pada Penanda Aras Sembang Model utama yang dibandingkan ialah model Transformer yang besar. Keputusan menunjukkan:

Model Mamba hibrid suling (50%) mencapai skor yang sama dengan model guru dalam penanda aras MT, dan lebih baik sedikit daripada model guru dalam penanda aras AlpacaEval dari segi kadar kemenangan LC dan kadar kemenangan keseluruhan.

Prestasi Mamba hibrid suling (25% dan 12.5%) adalah lebih buruk sedikit daripada model guru pada penanda aras MT, tetapi walaupun dengan lebih banyak parameter dalam AlpcaaEval ia masih mengatasi beberapa Transformer besar.

Ketepatan model Mamba tulen suling (0%) menurun dengan ketara.

Perlu diingat bahawa model hibrid suling berprestasi lebih baik daripada Falcon Mamba, yang dilatih dari awal menggunakan lebih daripada token 5T.

Penilaian tanda aras am

Penilaian sampel sifar. Jadual 3 menunjukkan prestasi pukulan sifar Mamba dan Mamba2 yang disuling daripada model guru yang berbeza pada penanda aras LM Eval. Model hibrid Mamba-Llama3 dan Mamba2-Llama3 yang disuling daripada Llama-3 Instruct 8B berprestasi lebih baik berbanding model TRI Mamba dan Nvidia Mamba sumber terbuka yang dilatih dari awal.

Penilaian penanda aras. Jadual 4 menunjukkan bahawa prestasi model hibrid suling sepadan dengan model RNN linear sumber terbuka terbaik pada Papan Pendahulu LLM Terbuka, sambil mengatasi prestasi model arahan sumber terbuka sepadan dalam GSM8K dan CRUX.

Penyahkodan Spekulatif Hibrid

Untuk model penyulingan 50% dan 25%, berbanding garis dasar bukan spekulatif, kajian ini Mencapai lebih 1.8x kelajuan pada Zephyr-Hybrid.

Eksperimen juga menunjukkan bahawa model draf 4 lapisan yang dilatih dalam kajian ini mencapai kadar penerimaan yang lebih tinggi, tetapi disebabkan peningkatan dalam saiz model draf, overhed tambahan juga menjadi lebih besar. Dalam kerja-kerja seterusnya, penyelidikan ini akan menumpukan pada mengecilkan model draf ini.

Perbandingan dengan kaedah penyulingan lain: Jadual 6 (kiri) membandingkan kebingungan varian model yang berbeza. Kajian itu melakukan penyulingan dalam satu zaman menggunakan Ultrachat sebagai gesaan benih dan membandingkan kebingungan. Ternyata mengalih keluar lebih banyak lapisan memburukkan keadaan. Kajian itu juga membandingkan kaedah penyulingan dengan garis dasar sebelumnya dan mendapati kaedah baharu menunjukkan kemerosotan yang lebih kecil, manakala model Distill Hyena dilatih pada dataset WikiText menggunakan model yang lebih kecil dan menunjukkan tahap kecelaruan yang lebih besar.

Jadual 6 (kanan) menunjukkan bahawa menggunakan SFT atau DPO sahaja tidak menghasilkan banyak peningkatan, manakala menggunakan SFT + DPO menghasilkan skor terbaik.

Jadual 7 membandingkan kajian ablasi untuk beberapa model berbeza. Jadual 7 (kiri) menunjukkan hasil penyulingan menggunakan pelbagai permulaan, dan Jadual 7 (kanan) menunjukkan keuntungan yang lebih kecil daripada penyulingan progresif dan lapisan perhatian interleaving dengan Mamba.

Jadual 8 membandingkan prestasi model hibrid menggunakan dua kaedah permulaan yang berbeza: keputusan mengesahkan bahawa permulaan berat perhatian adalah penting.

Jadual 9 membandingkan prestasi model dengan dan tanpa blok Mamba. Model dengan blok Mamba menunjukkan prestasi yang lebih baik daripada model tanpa blok Mamba. Ini mengesahkan bahawa penambahan lapisan Mamba adalah penting dan peningkatan prestasi bukan semata-mata disebabkan oleh mekanisme perhatian yang tinggal.

Pembaca yang berminat boleh membaca teks asal kertas kerja untuk mengetahui lebih lanjut tentang kandungan penyelidikan.

Atas ialah kandungan terperinci Kerja baharu daripada pengarang Mamba: Menyuling Llama3 menjadi RNN linear hibrid. Untuk maklumat lanjut, sila ikut artikel berkaitan lain di laman web China PHP!