Најсовремени техники кои ја забрзуваат вештачката интелигенција без дополнителни трошоци

Табела со врски

Апстракт и 1 Вовед

2 Државни вселенски модели

3 модели на простор за селективна состојба и 3.1 Мотивација: Избор како средство за компресија

3.2 Подобрување на SSM со селекција

3.3 Ефикасно спроведување на селективни ССМ

3.4 Поедноставена SSM архитектура

3.5 Својства на механизмите за селекција

3.6 Дополнителни детали за моделот

4 Емпириска евалуација и 4.1 Синтетички задачи

4.2 Јазично моделирање

4.3 Моделирање на ДНК

4.4 Аудио моделирање и генерирање

4.5 Репери за брзина и меморија

4.6 Модел аблации

5 Дискусија

6 Заклучок и референци

Дискусија: Механизам за селекција

Б Поврзана работа

C Механика на селективни SSMs

D Алгоритам за хардверски знаења за селективни SSMs

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

3.3 Ефикасно спроведување на селективни ССМ

Хардверските архитектури како што се конволуциите (Krizhevsky, Sutskever и Hinton 2012) и Transformers (Vaswani et al. 2017) уживаат широка примена. Овде имаме за цел да ги направиме селективните SSM ефикасни и на модерен хардвер (GPU). Механизмот на селекција е сосема природен, а претходните дела се обидоа да вклучат посебни случаи на селекција, како што е дозволувањето ∆ да варира со текот на времето во повторливи SSM (Гу, Дао, и сор. 2020). Сепак, како што беше претходно споменато, основно ограничување во користењето на SSM е нивната пресметковна ефикасност, поради што S4 и сите деривати користеа LTI (неселективни) модели, најчесто во форма на глобални конволуции.

3.3.1 Мотивација на претходните модели

Прво ја разгледуваме оваа мотивација и го прегледуваме нашиот пристап за надминување на ограничувањата на претходните методи.

• На високо ниво, повторливите модели како што се SSM секогаш го балансираат компромисот помеѓу експресивноста и брзината: како што беше дискутирано во Дел 3.1, моделите со поголема димензија на скриена состојба треба да бидат поефикасни, но побавни. Така сакаме да ја максимизираме димензијата на скриената состојба без да плаќаме трошоци за брзина и меморија.

• Забележете дека повторливиот режим е пофлексибилен од режимот на конволуција, бидејќи вториот (3) е изведен од проширување на првиот (2) (Gu, Goel и Ré 2022; Gu, Johnson, Goel, et al. 2021). Сепак, ова би барало пресметување и материјализирање на латентната состојба ℎ со облик (B, L, D, N), многу поголем (со фактор N, димензијата на состојбата на SSM) од влезот x и излезот y од обликот (B, L, D). Така беше воведен поефикасниот режим на конволуција кој може да го заобиколи пресметувањето на состојбата и материјализира конволуционо јадро (3а) од само (B, L, D).

• Претходните LTI SSM ги користат двојните рекурентно-конволуциони форми за да ја зголемат димензијата на ефективната состојба со фактор Nx (≈ 10 − 100), многу поголем од традиционалните RNN, без казни за ефикасност.

3.3.2 Преглед на селективно скенирање: проширување на состојбата со хардвер

Механизмот за избор е дизајниран да ги надмине ограничувањата на LTI моделите; во исто време, затоа треба повторно да го разгледаме пресметковниот проблем на ССМ. Ова го решаваме со три класични техники: фузија на јадрото, паралелно скенирање и повторно пресметување. Ние правиме две главни забелешки:

• Наивното повторливо пресметување користи O(BLDN) FLOP додека конволуционото пресметување користи O(BLD log(L)) FLOP, а првото има помал константен фактор. Така, за долги секвенци и не премногу голема димензија на состојбата N, рекурентниот режим всушност може да користи помалку FLOP.

• Двата предизвици се секвенцијалната природа на повторување и големата употреба на меморијата. За да се справиме со второто, исто како и конволуциониот режим, можеме да се обидеме да не ја материјализираме целосната состојба ℎ.

Главната идеја е да се искористат својствата на модерните акцелератори (GPU) за да се материјализира состојбата ℎ само во поефикасни нивоа на мемориската хиерархија. Особено, повеќето операции (освен множењето на матрицата) се ограничени со пропусниот опсег на меморијата (Дао, Фу, Ермон и сор. 2022; Иванов и сор. 2021; Вилијамс, Ватерман и Патерсон 2009). Ова ја вклучува нашата операција на скенирање и користиме фузија на јадрото за да го намалиме количеството мемориски IO, што доведува до значително забрзување во споредба со стандардната имплементација.

За да се избегне секвенцијалното повторување, забележуваме дека и покрај тоа што не е линеарен, сепак може да се паралелизира со алгоритам за паралелно скенирање ефикасен за работа (Blelloch 1990; Martin and Cundy 2018; Smith, Warrington и Linderman 2023).

Конечно, мораме да избегнеме и зачувување на средните состојби, кои се неопходни за назад пропагирање. Внимателно ја применуваме класичната техника на повторно пресметување за да ги намалиме барањата за меморија: средните состојби не се зачувуваат, туку се пресметуваат во назад премин кога влезовите се вчитуваат од HBM во SRAM. Како резултат на тоа, споениот слој за селективно скенирање ги има истите барања за меморија како оптимизираната имплементација на трансформаторот со FlashAttention.

Деталите за споеното јадро и повторното пресметување се во Додаток Г. Целосниот селективен SSM слој и алгоритам се илустрирани на Слика 1.