Griffin
Griffin : Menggabungkan Pengulangan Linear Berpagar dengan Perhatian Lokal untuk Model Bahasa yang Efisien
arXiv
Arsitektur Model
Semua model kami berisi komponen berikut: (i) blok sisa, (ii) blok MLP, dan (iii) blok pencampuran sementara. Meskipun (i) dan (ii) sama di semua model, kami mempertimbangkan tiga blok pencampuran temporal: Multi-Query Attention (MQA) global, MQA lokal (jendela geser), dan blok berulang yang kami usulkan. Sebagai bagian dari blok berulang kami menggunakan Real-Gated Linear Recurrent Unit (RG-LRU) – lapisan berulang baru yang terinspirasi oleh Linear Recurrent Unit Orvieto dkk., 2023b.
Blok sisa, seperti yang ditunjukkan pada Gambar 2(a), mendefinisikan struktur global model kami dan terinspirasi oleh Transformers pra-norma (Xiong et al., 2020). Setelah menyematkan urutan masukan, kami meneruskannya $N$ blok seperti itu ( $N$ menunjukkan kedalaman model), lalu kami menerapkan RMSNorm Zhang dan Sennrich, 2019 untuk menghasilkan aktivasi akhir. Untuk menghitung probabilitas token, kami menerapkan lapisan linier akhir diikuti dengan softmax. Bobot lapisan ini dibagi dengan lapisan penyematan masukan.
Blok sisa
Gambar 2: a) Tulang punggung utama arsitektur mode kita adalah blok sisa, yang ditumpuk $N$ kali. b) Blok MLP berpagar yang kami gunakan. c) Blok berulang yang kami usulkan sebagai alternatif Multi Query Attention (MQA). Ia menggunakan lapisan RG-LRU yang kami usulkan, yang didefinisikan di Bagian 2.4.
Blok sisa berisi dua komponen, diterapkan secara berurutan. Komponen pertama mengambil status tersembunyi $chi$ dan menerapkan RMSNorm Zhang dan Sennrich, 2019, diikuti dengan blok pencampuran temporal. Kami kemudian menggabungkan output dengan koneksi lewati dari $chi$ melalui penambahan. Demikian pula, komponen kedua menerapkan RMSNorm, diikuti oleh blok MLP dan kemudian menggabungkan outputnya dengan koneksi lewati dari input RMSNorm. Blok ini diilustrasikan pada Gambar 2 (a).
blok MLP
Kami menggunakan blok MLP yang terjaga keamanannya Dauphin dkk., 2017 (diilustrasikan pada Gambar 2(b)), yang membuat dua cabang dari masukan dimensinya $D$ . Kami menerapkan lapisan linier dengan dimensi keluaran $MD$ di setiap cabang, di mana $M$ menunjukkan faktor ekspansi. Untuk mempermudah, kami menggunakan $M=3$ sepanjang pekerjaan ini. Kami menerapkan non-linier GeLU Hendrycks dan Gimpel, 2016 di salah satu cabang sebelum menggabungkannya dengan perkalian berdasarkan elemen, mirip dengan GeGeLU Shazeer, 2020. Namun, di blok MLP kami, kami menerapkan lapisan linier akhir dengan dimensi keluaran $D$ pada keluaran lapisan GeGeLU.
Blok pencampuran temporal
Blok pencampuran temporal adalah komponen model kami yang menggabungkan aktivasi lapisan tersembunyi di lokasi temporal berbeda dalam urutan. Kami mempertimbangkan tiga blok pencampuran temporal: MQA Shazeer global, 2019, MQA Beltagy et al., 2020 lokal, dan blok Berulang yang kami usulkan.
Perhatian multi-kueri global
Kecuali dinyatakan lain, kami menggunakan MQA, bukan MHA, untuk meningkatkan kecepatan inferensi garis dasar Transformer kami Shazeer, 2019. Kami menggunakan dimensi head tetap $D_{kepala}=128$ , dan kami memperbaiki jumlah kepala perhatian $H$ seperti yang $HD_{kepala}=D$ . Ini membutuhkan dimensi model $D$ menjadi kelipatan 128. Kami tidak menggunakan penyematan posisi absolut, namun kami menggunakan Penyematan Posisi Putar (RoPE) Su dkk., 2021 sebagai penyematan posisi relatif.
Perhatian jendela geser lokal
Salah satu kelemahan utama menggunakan perhatian global adalah kompleksitas komputasinya bertambah secara kuadrat dalam panjang urutan. Untuk mengatasi hal ini, beberapa karya telah mulai mengadopsi perhatian lokal Beltagy et al., 2020, yang juga dikenal sebagai perhatian jendela geser. Hal ini memungkinkan setiap posisi untuk hanya menangani sejumlah token di masa lalu. Hal ini tidak hanya mengurangi FLOP komputasi tetapi juga membatasi ukuran cache KV dengan ukuran jendela, membuatnya tidak lagi kuadrat dalam panjang urutan. Semua detail lainnya sama dengan MQA global.
Blok berulang
Blok berulang kami (Gambar 2(c)) mirip dengan blok GSS Mehta dkk., 2022 dan blok yang digunakan oleh Mamba Gu dan Dao, 2023. Kami mengambil masukan dimensi $D$ dan terapkan dua lapisan linier dengan dimensi keluaran $D_{RNN}$ secara paralel, menciptakan dua cabang. Pada cabang pertama, kami menerapkan lapisan Conv1D kecil yang dapat dipisahkan, terinspirasi oleh Shift-SSM di H3 Dao et al., 2022b, dengan dimensi filter temporal 4. Perhatikan bahwa lapisan Conv1D ini sangat kecil, hanya dengan $4D$ parameter. Kami mengikuti lapisan Conv1D dengan lapisan RG-LRU yang kami usulkan (didefinisikan di bawah.) Pada cabang kedua kami menerapkan nonlinier GeLU dan kemudian menggabungkan cabang-cabang tersebut dengan perkalian berdasarkan elemen. Kami kemudian menerapkan lapisan linier akhir dengan dimensi keluaran $D$ .
Unit Berulang Linier Berpagar Nyata (RG-LRU)
Lapisan RG-LRU yang kami usulkan memiliki perulangan sederhana yang terinspirasi oleh Linear Recurrent Unit (LRU) Orvieto dkk., 2023b, tetapi menggabungkan mekanisme gerbang yang dimotivasi oleh literatur tentang RNN non-linier, khususnya LSTM Hochreiter dan Schmidhuber, 1997 dan GRUs Chung et al., 2014. Persamaan yang menggambarkan lapisan tersebut adalah sebagai berikut:
$$begin{align} r_t &= sigma(W_{a} x_t + b_a), & text{gerbang pengulangan} \ i_t &= sigma(W_{x} x_t + b_x), & text{ gerbang masukan} \ a_t &= a^{cr_t}, & text{} \ h_t &= a_t odot h_{t-1} + sqrt{1 - a_t^2} odot (i_t odot x_t). & teks{} end{align}$$
Output dari lapisan tersebut adalah $y_t=h_t$ , dan non-linearitas $sigma$ dalam persamaan adalah fungsi sigmoid. Berat badan yang berulang $a$ pada Persamaan (4) adalah diagonal. Oleh karena itu, semua operasi berdasarkan elemen. Kami membuat parameter $a$ pada Persamaan (3) sebagai $a=sigma(Lambda)$ , Di mana $Lambda$ adalah parameter yang bisa dipelajari. Ini menjamin hal itu $0 <= dan <= 1$ , memastikan kekambuhan stabil. Variabel $c$ adalah konstanta bernilai skalar yang disetel ke 8. Untuk stabilitas numerik, dalam praktiknya kita menghitung $a^{cr_t}$ di ruang log (lihat Lampiran A). Lapisan tersebut memiliki gerbang pada kedua masukan $x$ dan berat badan yang berulang $a$ . Namun, tidak ada gerbang yang bergantung pada keadaan berulang $h_{t-1}$ , yang memastikan bahwa komputasi dapat dijalankan secara efisien di perangkat. Kami menginisialisasi keduanya $W_{a}$ Dan $W_{b}$ menggunakan LeCun init LeCun et al., 2002. Kami menginisialisasi $Lambda$ seperti yang $a^c$ didistribusikan secara merata di antara $0,9$ Dan $0,999$ di awal pelatihan, serupa dengan (Orvieto et al., 2023b.). Tidak seperti banyak karya terbaru dalam literatur SSM, RG-LRU tidak menggunakan inisialisasi yang terinspirasi oleh teori polinomial ortogonal Gu dkk., 2020, dan juga tidak didefinisikan sebagai diskritisasi sistem kontinu yang mendasari Gu dkk., 2021a. Berbeda dengan lapisan LRU asli, kami tidak menggunakan aljabar kompleks dalam perulangannya. Meskipun penggunaan perulangan kompleks akan menghasilkan lapisan yang lebih ekspresif Orvieto dkk., 2023a kami menemukan bahwa perulangan kompleks tidak bermanfaat untuk pemodelan bahasa dalam praktiknya, seperti yang juga diamati oleh Gu dan Dao, 2023. (lihat Lampiran B)
Perilaku gerbang
Gerbang masukan $i_t$ mirip dengan yang ada di LSTM, yang dapat memfilter (atau memperkecil) masukan $x_t$ . Namun, sepengetahuan kami, gerbang pengulangan kami $r_t$ berbeda dari mekanisme gerbang lain dalam literatur. Misalnya, mekanisme seleksi yang diusulkan di Mamba Gu dan Dao, 2023 sebanding dengan gerbang pembaruan GRU yang melakukan interpolasi $x_t$ . Efeknya pada keadaan tersembunyi memungkinkannya untuk mengatur ulang keadaannya dan melupakan informasi apa pun yang disimpannya dari masa lalu, mirip dengan gerbang lupa di LSTM. Sebaliknya, gerbang perulangan kami kira-kira dapat melakukan interpolasi antara pembaruan LRU standar dari Orvieto et al., 2023a dan keadaan tersembunyi sebelumnya, yang memungkinkannya membuang masukan secara efektif dan menyimpan semua informasi dari riwayat sebelumnya (lihat Lampiran A untuk rincian lebih lanjut ). Kami yakin peran utama gerbang ini adalah memungkinkan model mencapai memori super eksponensial dengan mengurangi pengaruh masukan yang tidak informatif.
