Нов механизъм за внимание намалява KV кеша

Въведение

Нарастващото използване на големи езикови модели (LLM) и появата на нови парадигми в извода изведоха на преден план предизвикателството за ефективен мащабен извод. Значително ограничение е кешът Key-Value (KV) в традиционните механизми за внимание, който се разширява линейно с размера на партидата и дължината на последователността, превръщайки се в 'поглъщащ паметта', който възпрепятства мащабирането и разширяването на LLM.

Докато варианти като MQA, GQA и MLA се появиха, за да се справят с това, те често се борят да поддържат производителност при строги ограничения на паметта или въвеждат сложности, които поставят инженерни предизвикателства и проблеми със съвместимостта.

Multi-matrix Factorization Attention (MFA)

Скорошна статия от изследователи от Stepes, университета Tsinghua и други институции представя нова архитектура на механизма за внимание: Multi-matrix Factorization Attention (MFA) и неговия вариант MFA-Key-Reuse (MFA-KR). Този механизъм значително намалява разходите за извод на езикови модели, като същевременно подобрява производителността.

MFA и MFA-KR не само надминават MLA по производителност, но и съответстват на традиционната производителност на MHA, като същевременно намаляват използването на KV кеша до 93.7%.

MFA е проектиран за простота, лесно възпроизвеждане, ниска чувствителност към хиперпараметри и съвместимост с различни методи за позиционно вграждане.

Подход и анализ на MFA

Изследователският екип анализира общия дизайн и капацитет на механизмите за внимание, идентифицирайки две критични измерения, свързани с капацитета. Този анализ доведе до разработването на нови аналитични методи и принципи на проектиране.

Те въведоха концепцията за Generalized Multi-Head Attention (GMHA) като обединяваща рамка за разбиране на различните варианти на MHA.

Екипът също така изследва изчислението и съхранението на key-values от гледна точка на извода и изследва капацитета на модела от гледна точка на декомпозиция.

Fully Parameterized Bilinear Attention (FPBA) беше установен като теоретичната горна граница на производителността. Те установиха, че MHA и неговите варианти са нискорангови декомпозиции на FPBA.

Сравнение с MQA и MLA

Анализът се фокусира върху две представителни схеми за подобрение: Multi-Query Attention (MQA) и Multi-Head Latent Attention (MLA).

MQA използва по-агресивна стратегия за споделяне на параметри, при която всички глави за внимание споделят един и същ набор от key-value параметри. Това намалява използването на паметта, но може да повлияе на изразителността на модела.

MLA въвежда споделено латентно пространство за компресиране на параметрите, но действителната изразителна сила е ограничена от най-малкото измерение, което означава, че увеличаването на междинните измерения не подобрява значително производителността.

Ключови иновации на MFA

Разработването на MFA беше водено от целта да се създаде механизъм за внимание, който да минимизира консумацията на ресурси, като същевременно се доближава до теоретичните граници на производителност. Дизайнът на MFA включва три ключови иновации:

• Значително увеличаване на броя и измерението на главите за внимание, за да се максимизира капацитетът на модела.
• Използване на агресивна нискорангова стратегия за декомпозиция, за да се поддържа ефективност на параметрите, като същевременно се разширява броят и измеренията на главите за внимание.
• Използване на дизайн с единична key-value глава, за да се поддържа минимална консумация на памет, дори при повишена сложност на модела.

Измерване и сравнение на капацитета

За по-нататъшен анализ на MFA и други механизми за внимание, екипът въведе два ключови показателя:

• Total Effective Rank (TER): Произведението на броя на главите за внимание и ранга на факторизация на глава (FRH).
• Shared Latent Subspace Dimension (SLSD): Измерението на скритото пространство, споделено от всички глави за внимание.

MFA постига по-високи SLSD и TER в сравнение с MQA.

В сравнение с MLA, MFA постига по-малък размер на KV кеша и по-висок TER при подобни бюджети на параметри, като същевременно поддържа сравним SLSD.

В сравнение с традиционния MHA, MFA има по-висок TER, въпреки че неговият SLSD е по-малък.

Експериментални резултати

Проведени бяха обширни експерименти за оценка на производителността на новата архитектура в по-големи мащаби, тествайки модели, вариращи от 1B до 7B параметри, и данни за обучение от 10B до 1T.

MFA демонстрира мащабируемост, сравнима с традиционния MHA, като поддържа отлична производителност дори при по-големи мащаби.

Докато MFA-KR показа малко по-ниска производителност, неговата тенденция на мащабиране се изравни с MHA. Предимствата на спестяване на памет на MFA и MFA-KR продължиха да се разширяват с размера на модела, като MFA постигна 87.5% спестяване на памет, а MFA-KR намали използването на памет до 6.25% в най-голям мащаб.

Аблационни проучвания

Аблационните проучвания потвърдиха ефективността на MFA и MFA-KR. Техните предимства в производителността бяха потвърдени и при различни основни методи за позиционно кодиране.

Перспективи

MFA предлага значителни подобрения с прост дизайн, като ефективно адресира ограничението на паметта при извод на LLM, без да добавя допълнителна инженерна сложност. Той се интегрира безпроблемно в съществуващата екосистема на Transformer, ускорявайки прилагането на LLM в различни сценарии.