Mekanisme Perhatian Baru Mengurangkan KV Cache

Pengenalan

Penggunaan model bahasa besar (LLM) yang semakin meningkat dan kemunculan paradigma baru dalam inferensi telah membawa cabaran inferensi berskala besar yang efisien ke barisan hadapan. Kesesakan yang ketara adalah cache Key-Value (KV) dalam mekanisme perhatian tradisional, yang berkembang secara linear dengan saiz batch dan panjang urutan, menjadi ‘penghisap memori’ yang menghalang penskalaan dan pengembangan LLM.

Walaupun varian seperti MQA, GQA, dan MLA telah muncul untuk menangani masalah ini, mereka sering bergelut untuk mengekalkan prestasi di bawah kekangan memori yang ketat atau memperkenalkan kerumitan yang menimbulkan cabaran kejuruteraan dan isu keserasian.

Multi-matrix Factorization Attention (MFA)

Satu kertas kerja baru-baru ini oleh penyelidik dari Stepes, Universiti Tsinghua, dan institusi lain memperkenalkan seni bina mekanisme perhatian novel: Multi-matrix Factorization Attention (MFA) dan variannya MFA-Key-Reuse (MFA-KR). Mekanisme ini mengurangkan kos inferensi model bahasa dengan ketara sambil meningkatkan prestasi secara serentak.

MFA dan MFA-KR bukan sahaja mengatasi MLA dalam prestasi tetapi juga menyamai prestasi MHA tradisional sambil mengurangkan penggunaan KV Cache sehingga 93.7%.

MFA direka untuk kesederhanaan, mudah direproduksi, kepekaan rendah terhadap hiperparameter, dan keserasian dengan pelbagai kaedah Pos-embedding.

Pendekatan dan Analisis MFA

Pasukan penyelidik menganalisis reka bentuk umum dan kapasiti mekanisme perhatian, mengenal pasti dua dimensi kritikal yang berkaitan dengan kapasiti. Analisis ini membawa kepada pembangunan kaedah analisis dan prinsip reka bentuk baru.

Mereka memperkenalkan konsep Generalized Multi-Head Attention (GMHA) sebagai kerangka penyatuan untuk memahami varian MHA yang berbeza.

Pasukan itu juga meneroka pengiraan dan penyimpanan key-value dari perspektif inferensi dan mengkaji kapasiti model dari perspektif penguraian.

Fully Parameterized Bilinear Attention (FPBA) ditubuhkan sebagai had atas prestasi teori. Mereka mendapati bahawa MHA dan variannya adalah penguraian peringkat rendah FPBA.

Perbandingan dengan MQA dan MLA

Analisis tertumpu pada dua skim peningkatan perwakilan: Multi-Query Attention (MQA) dan Multi-Head Latent Attention (MLA).

MQA menggunakan strategi perkongsian parameter yang lebih agresif, di mana semua kepala perhatian berkongsi set parameter key-value yang sama. Ini mengurangkan penggunaan memori tetapi boleh menjejaskan kebolehan ekspresif model.

MLA memperkenalkan ruang laten yang dikongsi untuk pemampatan parameter, tetapi kuasa ekspresif sebenar dihadkan oleh dimensi terkecil, yang bermaksud bahawa peningkatan dimensi perantaraan tidak meningkatkan prestasi dengan ketara.

Inovasi Utama MFA

Pembangunan MFA didorong oleh matlamat untuk mencipta mekanisme perhatian yang meminimumkan penggunaan sumber sambil menghampiri had prestasi teori. Reka bentuk MFA menggabungkan tiga inovasi utama:

  • Meningkatkan bilangan dan dimensi kepala perhatian dengan ketara untuk memaksimumkan kapasiti model.
  • Menggunakan strategi penguraian peringkat rendah yang agresif untuk mengekalkan kecekapan parameter sambil mengembangkan bilangan dan dimensi kepala perhatian.
  • Menggunakan reka bentuk kepala key-value tunggal untuk memastikan penggunaan memori minimum, walaupun dengan peningkatan kerumitan model.

Pengukuran dan Perbandingan Kapasiti

Untuk menganalisis MFA dan mekanisme perhatian lain dengan lebih lanjut, pasukan memperkenalkan dua metrik utama:

  • Total Effective Rank (TER): Hasil daripada bilangan kepala perhatian dan peringkat Faktorisasi setiap kepala (FRH).
  • Shared Latent Subspace Dimension (SLSD): Dimensi ruang tersembunyi yang dikongsi oleh semua kepala perhatian.

MFA mencapai SLSD dan TER yang lebih tinggi berbanding dengan MQA.

Berbanding dengan MLA, MFA mencapai saiz cache KV yang lebih kecil dan TER yang lebih tinggi dengan belanjawan parameter yang serupa, sambil mengekalkan SLSD yang setanding.

Berbanding dengan MHA tradisional, MFA mempunyai TER yang lebih tinggi, walaupun SLSDnya lebih kecil.

Keputusan Eksperimen

Eksperimen yang meluas telah dijalankan untuk menilai prestasi seni bina baru pada skala yang lebih besar, menguji model dari 1B hingga 7B parameter dan data latihan dari 10B hingga 1T.

MFA menunjukkan keupayaan penskalaan yang setanding dengan MHA tradisional, mengekalkan prestasi yang sangat baik walaupun pada skala yang lebih besar.

Walaupun MFA-KR menunjukkan prestasi yang sedikit lebih rendah, trend penskalaannya sejajar dengan MHA. Kelebihan penjimatan memori MFA dan MFA-KR terus berkembang dengan saiz model, dengan MFA mencapai penjimatan memori 87.5% dan MFA-KR mengurangkan penggunaan memori kepada 6.25% pada skala terbesar.

Kajian Ablasi

Kajian ablasi mengesahkan keberkesanan MFA dan MFA-KR. Kelebihan prestasi mereka juga disahkan dalam pelbagai kaedah pengekodan posisi arus perdana.

Pandangan

MFA menawarkan peningkatan yang ketara dengan reka bentuk yang sederhana, dengan berkesan menangani kesesakan memori dalam inferensi LLM tanpa menambah kerumitan kejuruteraan tambahan. Ia berintegrasi dengan lancar ke dalam ekosistem Transformer sedia ada, mempercepatkan penggunaan LLM dalam pelbagai senario.