Apa itu perhatian kueri yang dikelompokkan (GQA)?

Ainslie dkk menyusun perhatian kueri yang dikelompokkan sebagai pengoptimalan perhatian multi-kepala (MHA), algoritma perhatian-diri inovatif yang diperkenalkan dalam makalah “Attention is All You Need” tahun 2017 yang membentuk neural networks. Lebih khusus lagi, ini diusulkan sebagai generalisasi dan aplikasi yang lebih terkendali dari multi-query attention (MQA), sebuah pengoptimalan MHA sebelumnya.

Meskipun perhatian multi-kepala standar mengkatalisasi lompatan evolusioner dalam machine learning, pemrosesan bahasa alami (NLP), dan AI generatif, hal ini sangat menuntut sumber daya komputasi dan bandwidth memori. Karena LLM tumbuh lebih besar dan lebih canggih, persyaratan penggunaan memori ini menjadi hambatan dalam kemajuan, terutama untuk LLM khusus dekoder autoregresif yang digunakan untuk pembuatan teks, peringkasan, dan tugas AI generatif lainnya.

Penelitian selanjutnya berfokus pada teknik untuk meningkatkan atau merampingkan perhatian multi-kepala. Beberapa, seperti perhatian kilat dan perhatian cincin, meningkatkan cara GPU yang digunakan untuk melatih dan menjalankan model menangani perhitungan dan penyimpanan memori. Lainnya, seperti GQA dan MQA, mengeksplorasi perubahan pada cara arsitektur transformator memproses token.

Group Query Attention (GQA) bertujuan untuk menyeimbangkan keunggulan antara perhatian multi kepala dan multi kueri. Hal tersebut meningkatkan akurasi tetapi memerlukan lebih banyak bandwidth memori dan mengurangi kecepatan, sementara perhatian multi kueri lebih cepat dan efisien tetapi dengan akurasi yang lebih rendah.

Untuk memahami bagaimana perhatian kueri yang dikelompokkan mengoptimalkan model transformator, penting untuk terlebih dahulu memahami cara kerja perhatian multi-kepala secara umum. Baik GQA dan MQA hanya menyempurnakan, bukan menggantikan, metodologi inti MHA.

Kekuatan pendorong di balik LLM dan model lain yang menggunakan arsitektur transformator adalah perhatian-diri, suatu kerangka kerja untuk memahami hubungan antara setiap token yang berbeda dalam suatu urutan. Perhatian diri memungkinkan LLM untuk menafsirkan data teks tidak hanya melalui definisi dasar statis, tetapi juga konteks yang disediakan oleh kata dan frasa lain.

Dalam LLM autoregresif yang digunakan untuk pembuatan teks, mekanisme perhatian membantu model memprediksi token berikutnya dalam sebuah urutan dengan menentukan token sebelumnya yang paling layak “diperhatikan” pada saat itu. Informasi dari token yang dianggap paling relevan diberi bobot perhatian yang lebih besar, sementara informasi dari token yang dianggap tidak relevan diberi bobot perhatian mendekati 0.

Mekanisme perhatian multi-kepala dalam model transformator memperkaya informasi kontekstual dengan menghitung perhatian diri secara paralel. Ini dicapai dengan membagi lapisan perhatian menjadi beberapa kepala perhatian.

Penulis "Attention is All You Need" mengartikulasikan mekanisme perhatiannya dengan menggunakan terminologi basis data relasional: kueri, kunci, dan nilai. Basis data relasional dirancang untuk menyederhanakan penyimpanan dan pengambilan data yang relevan: basis data ini memberikan pengenal unik (“kunci”) untuk setiap bagian data, dan setiap kunci dikaitkan dengan nilai yang sesuai. Tujuan dari database relasional adalah mencocokkan setiap kueri dengan kunci yang sesuai.

Untuk setiap token dalam urutan, perhatian multi-kepala membutuhkan pembuatan 3 vektor.

• Vektor kueri, Q, menunjukkan informasi yang “dicari“ oleh token. Misalnya, vektor kueri untuk kata benda mungkin menunjukkan pencarian kata sifat yang menggambarkannya.
  
  
• Vektor kunci, K, menunjukkan informasi yang terkandung dalam token. Skor penyelarasan yang menunjukkan relevansi dari vektor kunci setiap token terhadap vektor kueri setiap token lainnya digunakan untuk menghitung bobot perhatian.
  
  
• Sebuah vektor nilai, V, yang mewakili informasi kontekstual yang akan diperbarui oleh kontribusi berbobot perhatian dari vektor kunci token lainnya.

Interaksi matematis antara ketiga vektor ini, yang dimediasi oleh mekanisme perhatian, adalah bagaimana sebuah model menyesuaikan pemahaman konteks spesifiknya terhadap setiap token.
 

Untuk menghasilkan masing-masing dari 3 vektor ini untuk token yang diberikan, model dimulai dengan penyematan vektor asli token tersebut :pengkodean numerik di mana setiap dimensi vektor sesuai dengan beberapa elemen abstrak dari makna semantik token. Jumlah dimensi dalam vektor ini adalah hyperparameter yang telah ditentukan.

Vektor Q, K, dan V untuk setiap token dihasilkan dengan melewatkan penyematan token asli melalui lapisan linier yang mendahului lapisan perhatian pertama. Lapisan linier ini dipartisi menjadi 3 matriks unik bobot model: W_Q, W_K dan W_V. Nilai bobot spesifik di dalamnya dipelajari melalui pra- pelatihan yang diawasi sendiri pada kumpulan data besar contoh teks.

Mengalikan penyematan vektor asli token dengan_WQ, _WK dan_WV menghasilkan vektor kueri, vektor kunci dan vektor nilai yang sesuai. Jumlah dimensi d yang dikandung setiap vektor ditentukan oleh ukuran setiap matriks bobot. Q dan K akan memiliki jumlah dimensi yang sama, d_k.

Ketiga vektor ini kemudian diteruskan ke lapisan perhatian.

Pada lapisan perhatian, vektor Q, K dan V digunakan untuk menghitung skor keselarasan antara setiap token pada setiap posisi dalam sebuah urutan. Skor penyelarasan tersebut kemudian dinormalisasi menjadi bobot perhatian menggunakan fungsi softmax.

Untuk setiap token x dalam sebuah urutan, skor perataan dihitung dengan menghitung dot product dari vektor kueri token tersebut, _Qx, dengan vektor kunci K dari setiap token lainnya: dengan kata lain, dengan mengalikan keduanya. Jika hubungan yang berarti antara 2 token tercermin dalam kemiripan antara vektor masing-masing, mengalikannya bersama-sama akan menghasilkan nilai yang besar. Jika kedua vektor tidak sejajar, mengalikannya akan menghasilkan nilai yang kecil atau negatif. Sebagian besar model transformator menggunakan varian yang disebut perhatian produk titik berskala, di mana QK diskalakan-yaitu, dikalikan dengan $\frac{1}{\sqrt{d_k}}$  untuk meningkatkan stabilitas pelatihan.

Skor penyelarasan kunci kueri ini kemudian diketik ke dalam fungsi softmax. Softmax menormalkan semua input ke nilai antara 0 dan 1 sehingga semuanya berjumlah 1. Hasil dari fungsi softmax adalah bobot perhatian, masing-masing mewakili bagian (dari 1) perhatian tokenyang harus dibayarkan ke masing-masing token lainnya. Jika bobot perhatian token mendekati 0, maka akan diabaikan. Bobot perhatian 1 berarti token menerima seluruh perhatian xdan yang lainnya akan diabaikan.

Akhirnya, vektor nilai untuk setiap token dikalikan dengan bobot perhatiannya. Kontribusi berbobot perhatian dari setiap token sebelumnya dirata-ratakan dan ditambahkan ke penyematan vektor asli untuk token x. Dengan ini, penyematan token xsekarang diperbarui untuk mencerminkan konteks yang disediakan oleh token lain dalam urutan yang relevan dengannya.

Vektor penyematan yang diperbarui kemudian dikirim ke lapisan linier berikutnya dengan matriks bobotnya sendiri, W_Z, vektor konteks dinormalisasi agar memiliki jumlah dimensi yang konsisten sebelum diteruskan ke lapisan perhatian selanjutnya. Setiap lapisan perhatian berikutnya menangkap nuansa kontekstual yang lebih mendalam.
 

Menggunakan rata-rata kontribusi tertimbang dari token lain, dibandingkan mempertimbangkan setiap bagian konteks secara individual, lebih efisien secara matematis. Namun, pendekatan ini dapat menyebabkan hilangnya detail penting.

Untuk mengimbanginya, jaringan transformator membagi penyematan token input asli menjadi h bagian yang berukuran sama. Mereka juga membagi _WQ, _WK dan_WV menjadi h himpunan bagian yang masing-masing disebut kepala kueri, kepala kunci dan kepala nilai. Setiap kepala kueri, kepala kunci, dan kepala nilai menerima bagian dari penyematan token asli. Vektor yang dihasilkan oleh masing-masing triplet paralel dari kepala kueri, kepala kunci, dan kepala nilai dimasukkan ke dalam kepala perhatian yang sesuai.Pada akhirnya, hasil dari rangkaian paralel h ini disatukan kembali untuk memperbarui penyematan token secara penuh.

Selama pelatihan, setiap sirkuit dalam model mempelajari bobot yang berbeda untuk menangkap aspek makna semantik yang unik. Proses ini memungkinkan model memahami berbagai makna sebuah kata berdasarkan konteks kata-kata lain di sekitarnya.

Kelemahan utama dari perhatian multi-kepala standar bukan terletak pada kekurangan mendasar, tetapi pada kurangnya pengoptimalan. Sebagai algoritma pertama dalam kategorinya, MHA merupakan implementasi paling kompleks dari mekanisme perhatian dalam komputasi.

Sebagian besar ketidakefisienan MHA berasal dari banyaknya perhitungan dan parameter model. Pada MHA standar, setiap kepala kueri, kepala kunci, dan kepala nilai di setiap blok perhatian memiliki matriks bobotnya sendiri. Jadi, sebagai contoh, sebuah model dengan 8 attention head di setiap lapisan perhatian—jauh lebih sedikit daripada kebanyakan LLM modern— akan membutuhkan 24 matriks bobot unik untuk lapisan Q, K, dan V head saja. Hal ini memerlukan banyak sekali perhitungan perantara di setiap lapisan.

Salah satu konsekuensi dari konfigurasi ini adalah bahwa hal ini mahal secara komputasi. Kebutuhan komputasi untuk skala MHA secara kuadratik sehubungan dengan panjang urutan: menggandakan jumlah token dalam urutan input membutuhkan empat kali lipat kompleksitas. Hal ini memberikan batasan praktis pada ukuran jendela konteks.

MHA juga menempatkan tekanan besar pada memori sistem. GPU tidak memiliki banyak memori on-board untuk menyimpan hasil dari sejumlah besar perhitungan menengah yang harus diingat pada setiap langkah pemrosesan berikutnya. Hasil antara ini disimpan dalam memori bandwidth tinggi (HBM), yang tidak terletak pada chip GPU itu sendiri. Solusi ini memerlukan sedikit latensi setiap kali kunci dan nilai harus dibaca dari memori. Ketika model transformator mulai untuk menskalakan hingga miliaran parameter, waktu dan komputasi yang diperlukan untuk melatih dan menjalankan inferensi menjadi hambatan pada kinerja model.

Kemajuan lebih lanjut memerlukan metode untuk mengurangi jumlah langkah komputasi tanpa mengorbankan kemampuan transformer dalam mempelajari dan mereproduksi pola linguistik yang kompleks. Dalam konteks ini, MQA diperkenalkan, diikuti oleh GQA sebagai pengembangannya.

Perhatian multi-kueri (MQA) adalah mekanisme yang lebih efisien secara komputasi karena menyederhanakan struktur perhatian multi-kepala, sehingga mengurangi penggunaan memori dan perhitungan tambahan. Alih-alih melatih kepala kunci dan kepala nilai yang berbeda untuk setiap kepala perhatian, MQA menggunakan satu kepala kunci dan satu kepala nilai per lapisan. Dengan demikian, vektor kunci dan nilai dihitung hanya sekali, kemudian satu set vektor tersebut digunakan oleh semua h kepala perhatian.

Penyederhanaan ini sangat mengurangi jumlah proyeksi linier yang harus dihitung dan disimpan oleh model dalam memori bandwidth tinggi. Menurut makalah tahun 2019 yang memperkenalkan MQA, MQA memungkinkan penyimpanan pasangan nilai-kunci (atau cache KV) yang 10-100 kali lebih kecil dan inferensi dekoder 12 kali lebih cepat. Pengurangan penggunaan memori MQA juga secara signifikan mempercepat pelatihan dengan mengaktifkan ukuran batch yang lebih besar.

Terlepas dari manfaat, MQA hadir dengan beberapa kelemahan yang tak terhindarkan.

• Degradasi kinerja: Tidak mengherankan, mengurangi jumlah parameter model yang unik dan dapat dilatih akan mengurangi kapasitas model untuk pengetahuan dan nuansa. MQA menyebabkan penurunan akurasi yang cukup berarti dibandingkan dengan MHA standar, sehingga tidak cocok untuk situasi dan contoh penggunaan tertentu.
  
  
• Harus dilatih dari awal: Sebuah model yang dilatih dengan MHA standar tidak dapat begitu saja diadaptasi ke MQA, tetapi harus dilatih dengan MQA dari awal. Hal ini berarti MQA tidak dapat digunakan untuk mengoptimalkan model yang sudah ada dan memerlukan biaya peluang yang cukup besar ketika bereksperimen dengan MQA untuk model baru.
  
  
• Redundansi dalam paralelisme tensor: Salah satu manfaat utama melatih model transformator pada GPU adalah kemampuan untuk melakukan beberapa operasi tensor kompleks secara paralel. Nilai K dan V harus ada pada setiap node cluster GPU yang melakukan operasi ini, yang berarti bahwa dalam praktiknya mereka harus direplikasi untuk setiap node. Ini bukan penggunaan sumber daya komputasi yang optimal meskipun masih lebih efisien daripada MHA standar.

Perhatian kueri yang dikelompokkan adalah formulasi yang lebih umum dan fleksibel dari perhatian multi-kueri yang mempartisi kepala kueri ke dalam beberapa kelompok yang masing-masing berbagi satu set kunci dan nilai, daripada berbagi satu set kunci dan nilai di semua kepala kueri.

Setelah publikasi makalah "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints" pada Mei 2023, banyak model bahasa besar (LLM) dengan cepat mengadopsi pendekatan GQA. Misalnya, Meta mulai menerapkan GQA pada model Llama 2 pada Juli 2023 dan terus menggunakannya pada model Llama 3 yang dirilis pada tahun 2024. Begitu pula, Mistral AI mengaplikasikan GQA pada model Mistral 7B yang diluncurkan pada September 2023. Dan model-model Granite 3.0 dari IBM juga menggunakan GQA untuk mencapai inferensi yang lebih cepat.

Secara teori, GQA dapat dilihat sebagai bentuk generalisasi antara MHA standar dan MQA penuh. Jika jumlah kelompok kepala nilai-kunci dalam GQA sama dengan jumlah kepala perhatian, maka GQA setara dengan MHA standar. Sementara itu, jika hanya ada satu kelompok kepala, maka GQA setara dengan MQA.

Dalam praktiknya, GQA hampir selalu menyiratkan beberapa pendekatan perantara, yang mana jumlah kelompok itu sendiri merupakan hyperparameter penting.

Perhatian kueri yang dikelompokkan menawarkan beberapa keuntungan yang telah menyebabkan adopsi yang relatif luas untuk LLM terkemuka.

• Penggunaan GPU yang efisien: Distribusi pasangan nilai-kunci GQA mengambil keuntungan dari paralelisme tensor, sehingga mengurangi jumlah komputasi yang "terbuang" dengan mereplikasi nilai yang berlebihan.
  
  
• Keseimbangan yang optimal: GQA memberikan keseimbangan ideal antara kecepatan inferensi decoder dan akurasi kinerja, dengan tingkat akurasi mendekati MHA dan kecepatan hampir setara dengan MQA.
  
  
• Mengurangi overhead bandwidth memori: Seperti MQA, GQA secara signifikan mengurangi jumlah kalkulasi perantara yang harus dikomputasi, disimpan, dan diambil pada waktu inferensi.
  
  
• Pelatihan yang fleksibel: Berbeda dengan MQA, group query attention (GQA) tidak mengharuskan model dilatih dari awal dengan pendekatan ini. Model yang sebelumnya dilatih menggunakan MHA standar dapat disesuaikan untuk menggunakan GQA melalui proses penyempurnaan yang dikenal sebagai "pelatihan ulang."