Published on

Mekanisme Perhatian Baru Step Mengurangi KV Cache untuk LLM

Penulis
  • avatar
    Nama
    Ajax
    Twitter

Pengantar

Peningkatan penggunaan Large Language Models (LLM) dan munculnya paradigma baru dalam inferensi telah menyoroti tantangan inferensi skala besar yang efisien. Kendala signifikan adalah Key-Value (KV) cache dalam mekanisme perhatian tradisional, yang berkembang secara linear dengan ukuran batch dan panjang sequence, menjadi "pemakan memori" yang menghambat penskalaan dan perluasan LLM.

Meskipun varian seperti MQA, GQA, dan MLA telah muncul untuk mengatasi hal ini, mereka sering kesulitan mempertahankan kinerja di bawah batasan memori yang ketat atau memperkenalkan kompleksitas yang menimbulkan tantangan rekayasa dan masalah kompatibilitas.

Multi-matrix Factorization Attention (MFA)

Sebuah makalah baru-baru ini oleh para peneliti dari Stepes, Universitas Tsinghua, dan lembaga lain memperkenalkan arsitektur mekanisme perhatian baru: Multi-matrix Factorization Attention (MFA) dan variannya MFA-Key-Reuse (MFA-KR). Mekanisme ini secara signifikan mengurangi biaya inferensi model bahasa sambil secara bersamaan meningkatkan kinerja.

MFA dan MFA-KR tidak hanya melampaui MLA dalam kinerja tetapi juga menyamai kinerja MHA tradisional sambil mengurangi penggunaan KV Cache hingga 93,7%.

MFA dirancang untuk kesederhanaan, reproduksi mudah, sensitivitas rendah terhadap hyperparameter, dan kompatibilitas dengan berbagai metode Pos-embedding.

Pendekatan dan Analisis MFA

Tim peneliti menganalisis desain umum dan kapasitas mekanisme perhatian, mengidentifikasi dua dimensi penting yang terkait dengan kapasitas. Analisis ini mengarah pada pengembangan metode analisis dan prinsip desain baru.

Mereka memperkenalkan konsep Generalized Multi-Head Attention (GMHA) sebagai kerangka kerja pemersatu untuk memahami berbagai varian MHA.

Tim juga mengeksplorasi komputasi dan penyimpanan key-value dari perspektif inferensi dan memeriksa kapasitas model dari perspektif dekomposisi.

Fully Parameterized Bilinear Attention (FPBA) ditetapkan sebagai batas atas teoritis kinerja. Mereka menemukan bahwa MHA dan variannya adalah dekomposisi low-rank dari FPBA.

Perbandingan dengan MQA dan MLA

Analisis ini berfokus pada dua skema peningkatan representatif: Multi-Query Attention (MQA) dan Multi-Head Latent Attention (MLA).

MQA menggunakan strategi berbagi parameter yang lebih agresif, di mana semua attention head berbagi set parameter key-value yang sama. Ini mengurangi penggunaan memori tetapi dapat memengaruhi daya ekspresif model.

MLA memperkenalkan ruang laten bersama untuk kompresi parameter, tetapi daya ekspresif aktual dibatasi oleh dimensi terkecil, yang berarti bahwa peningkatan dimensi perantara tidak secara signifikan meningkatkan kinerja.

Inovasi Kunci MFA

Pengembangan MFA didorong oleh tujuan untuk menciptakan mekanisme perhatian yang meminimalkan konsumsi sumber daya sambil mendekati batas kinerja teoretis. Desain MFA menggabungkan tiga inovasi utama:

  • Secara signifikan meningkatkan jumlah dan dimensi attention head untuk memaksimalkan kapasitas model.
  • Menggunakan strategi dekomposisi low-rank yang agresif untuk mempertahankan efisiensi parameter sambil memperluas jumlah dan dimensi attention head.
  • Memanfaatkan desain single key-value head untuk menjaga konsumsi memori minimal, bahkan dengan peningkatan kompleksitas model.

Pengukuran dan Perbandingan Kapasitas

Untuk lebih menganalisis MFA dan mekanisme perhatian lainnya, tim memperkenalkan dua metrik utama:

  • Total Effective Rank (TER): Hasil perkalian dari jumlah attention head dan Factorization rank per head (FRH).
  • Shared Latent Subspace Dimension (SLSD): Dimensi ruang tersembunyi yang dibagi oleh semua attention head.

MFA mencapai SLSD dan TER yang lebih tinggi dibandingkan dengan MQA. Dibandingkan dengan MLA, MFA mencapai ukuran KV cache yang lebih kecil dan TER yang lebih tinggi dengan anggaran parameter yang serupa, sambil mempertahankan SLSD yang sebanding. Dibandingkan dengan MHA tradisional, MFA memiliki TER yang lebih tinggi, meskipun SLSD-nya lebih kecil.

Hasil Eksperimen

Eksperimen ekstensif dilakukan untuk mengevaluasi kinerja arsitektur baru pada skala yang lebih besar, menguji model mulai dari 1B hingga 7B parameter dan data pelatihan dari 10B hingga 1T.

MFA menunjukkan kemampuan penskalaan yang sebanding dengan MHA tradisional, mempertahankan kinerja yang sangat baik bahkan pada skala yang lebih besar. Meskipun MFA-KR menunjukkan kinerja yang sedikit lebih rendah, tren penskalaannya sejajar dengan MHA. Keunggulan penghematan memori dari MFA dan MFA-KR terus meningkat seiring dengan ukuran model, dengan MFA mencapai penghematan memori 87,5% dan MFA-KR mengurangi penggunaan memori menjadi 6,25% pada skala terbesar.

Studi Ablasi

Studi ablasi memvalidasi efektivitas MFA dan MFA-KR. Keunggulan kinerja mereka juga dikonfirmasi di berbagai metode pengkodean posisi utama.

Prospek

MFA menawarkan peningkatan signifikan dengan desain sederhana, secara efektif mengatasi hambatan memori dalam inferensi LLM tanpa menambahkan kompleksitas rekayasa tambahan. Ini terintegrasi dengan mulus ke dalam ekosistem Transformer yang ada, mempercepat penerapan LLM di berbagai skenario.