Принстън и CMU разширяват границите на AI с модела Mamba Sequence

Таблица с връзки

Резюме и 1 въведение

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

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

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

3.3 Ефективно прилагане на селективни SSM

3.4 Опростена SSM архитектура

3.5 Свойства на механизмите за подбор

3.6 Допълнителни подробности за модела

4 Емпирична оценка и 4.1 Синтетични задачи

4.2 Езиково моделиране

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

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

4.5 Показатели за скорост и памет

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

5 Дискусия

6 Заключение и литература

Дискусия: Механизъм за подбор

B Свързана работа

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

D Хардуерно съобразен алгоритъм за селективни SSM

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

Резюме

Основните модели, които сега захранват повечето от вълнуващите приложения в дълбокото обучение, са почти универсално базирани на архитектурата на Transformer и нейния основен модул за внимание. Много архитектури на субквадратично време като линейно внимание, затворена конволюция и повтарящи се модели и модели на структурирано пространство на състоянието (SSM) са разработени за справяне с изчислителната неефективност на Transformers при дълги последователности, но те не са се представили толкова добре, колкото вниманието върху важни модалности като като език. Ние идентифицираме, че ключова слабост на такива модели е тяхната неспособност да извършват разсъждения, базирани на съдържание, и правим няколко подобрения. Първо, простото оставяне на параметрите на SSM да бъдат функции на входа адресира тяхната слабост с дискретни модалности, позволявайки на модела селективно да разпространява или забравя информация по дължината на измерението на последователността в зависимост от текущия токен. Второ, въпреки че тази промяна предотвратява използването на ефективни навивки, ние проектираме паралелен алгоритъм, съобразен с хардуера, в повтарящ се режим. Ние интегрираме тези селективни SSMs в опростена архитектура на невронна мрежа от край до край без внимание или дори без MLP блокове (Mamba). Mamba се радва на бърз извод (5 пъти по-висока пропускателна способност от Transformers) и линейно мащабиране на дължината на последователността, а нейната производителност се подобрява върху реални данни до последователности с дължина милион. Като основен модел на обща последователност, Mamba постига най-съвременна производителност в няколко модалности като език, аудио и геномика. Що се отнася до моделирането на езика, нашият модел Mamba-3B превъзхожда Transformers със същия размер и съвпада с Transformers два пъти по-големия си, както при предварително обучение, така и при оценка надолу по веригата.

1 Въведение

Основните модели (FM) или големи модели, предварително обучени върху масивни данни, след това адаптирани за задачи надолу по веригата, се очертаха като ефективна парадигма в съвременното машинно обучение. Гръбнакът на тези FM често са модели на последователност, работещи върху произволни последователности от входове от голямо разнообразие от области като език, изображения, реч, аудио, времеви серии и геномика (Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz et al., 2016 г.; Vinyals и Quoc V Le 2014). Въпреки че тази концепция е агностична за конкретен избор на архитектура на модела, съвременните FM се основават предимно на един тип модел на последователност: Трансформаторът (Vaswani et al. 2017) и неговият основен слой на вниманието (Bahdanau, Cho и Bengio 2015) ефикасността на самовниманието се приписва на способността му да насочва информацията плътно в контекстен прозорец, което му позволява да моделира сложни данни. Това свойство обаче носи основни недостатъци: невъзможност да се моделира нещо извън краен прозорец и квадратично мащабиране по отношение на дължината на прозореца. Появиха се огромни изследвания върху по-ефективни варианти на внимание за преодоляване на тези недостатъци (Tay, Dehghani, Bahri, et al. 2022), но често за сметка на самите свойства, които го правят ефективен. Досега нито един от тези варианти не е доказано емпирично ефективен в мащаб в различни области.

Наскоро структурирани модели на последователност в пространството на състоянието (SSM) (Gu, Goel и Ré 2022; Gu, Johnson, Goel и др. 2021) се появиха като обещаващ клас архитектури за моделиране на последователност. Тези модели могат да се интерпретират като комбинация от повтарящи се невронни мрежи (RNNs) и конволюционни невронни мрежи (CNNs), с вдъхновение от класическите модели на пространството на състоянието (Kalman 1960). Този клас модели може да бъде изчислен много ефективно или като повторение, или като конволюция, с линейно или почти линейно мащабиране в дължината на последователността. Освен това те имат принципни механизми за моделиране на дългосрочни зависимости (Gu, Dao, et al. 2020) в определени модалности на данни и са доминирали в бенчмаркове като Long Range Arena (Tay, Dehghani, Abnar, et al. 2021). Много разновидности на SSM (Gu, Goel и Ré 2022; Gu, Gupta и др. 2022; Gupta, Gu и Berant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; Smith, Warrington и Linderman 2023) са успешни в области, включващи непрекъснати сигнални данни като аудио и визия (Goel et al. 2022; Nguyen, Goel, et al. 2022; Saon, Gupta и Cui 2023). Въпреки това, те са били по-малко ефективни при моделиране на дискретни и плътни данни, като текст.

Предлагаме нов клас селективни модели на пространството на състоянията, които подобряват предишната работа по няколко оси, за да постигнат мощността на моделиране на Transformers, докато се мащабират линейно в дължината на последователността.

Механизъм за подбор. Първо, идентифицираме ключово ограничение на предходните модели: способността за ефикасен избор на данни по начин, зависим от входа (т.е. фокусиране върху или игнориране на конкретни входове). Въз основа на интуицията, базирана на важни синтетични задачи като селективни копиращи и индукционни глави, ние проектираме прост механизъм за избор чрез параметризиране на SSM параметрите въз основа на входа. Това позволява на модела да филтрира неуместната информация и да помни подходящата информация за неопределено време.

Хардуерно съобразен алгоритъм. Тази проста промяна представлява техническо предизвикателство за изчисляването на модела; всъщност всички предишни модели на SSMs трябва да бъдат инвариантни във времето и входа, за да бъдат изчислително ефективни. Ние преодоляваме това с алгоритъм, съобразен с хардуера, който изчислява модела периодично със сканиране вместо конволюция, но не материализира разширеното състояние, за да избегне IO достъп между различните нива на йерархията на паметта на GPU. Полученото внедряване е по-бързо от предишните методи както на теория (линейно мащабиране в дължината на последователността, в сравнение с псевдолинейно за всички SSM, базирани на конволюция), така и на модерен хардуер (до 3 пъти по-бързо на A100 GPU).

Архитектура . Ние опростяваме предишни архитектури на модели с дълбока последователност, като комбинираме дизайна на предишни SSM архитектури (Dao, Fu, Saab, et al. 2023) с MLP блока на Transformers в един блок, което води до прост и хомогенен архитектурен дизайн (Mamba), включващ селективни пространства на състояния.

Селективните SSM и като разширение архитектурата Mamba са напълно повтарящи се модели с ключови свойства, които ги правят подходящи като гръбнак на общите базови модели, работещи върху последователности. (i) Високо качество: селективността носи силна производителност при плътни модалности като език и геномика. (ii) Бързо обучение и извод: изчислението и паметта се мащабират линейно в дължината на последователността по време на обучението, а разгръщането на модела авторегресивно по време на извод изисква само постоянно време на стъпка, тъй като не изисква кеш на предишни елементи. (iii) Дълъг контекст: качеството и ефективността заедно дават подобрения на производителността на реални данни до дължина на последователност 1M.

Ние емпирично валидираме потенциала на Mamba като основна последователност на FM гръбнака, както в качеството на предварителна тренировка, така и в изпълнението на специфични за домейна задачи, на няколко вида модалности и настройки:

• Синтетика. При важни синтетични задачи като копиране и индукционни глави, които са предложени като ключови за големи езикови модели, Mamba не само ги решава лесно, но може да екстраполира решения за неопределено време (>1M токени).

• Аудио и геномика. Mamba превъзхожда предишни най-съвременни модели като SaShiMi, Hyena и Transformers при моделиране на аудио вълнови форми и ДНК последователности, както в качеството преди обучението, така и в показателите надолу по веригата (напр. намаляване на FID при предизвикателен набор от данни за генериране на реч с повече от половината ). И в двете настройки неговата производителност се подобрява с по-дълъг контекст до поредици с дължина милион.

• Езиково моделиране. Mamba е първият модел на линейна времева последователност, който наистина постига производителност с качество на Transformer, както при объркване преди тренировка, така и при оценки надолу по веригата. Със закони за мащабиране до 1B параметри, ние показваме, че Mamba надхвърля производителността на голям набор от базови линии, включително много силни съвременни рецепти за обучение на Transformer, базирани на LLaMa (Touvron et al. 2023). Нашият езиков модел Mamba има 5 пъти по-голяма производителност в сравнение с Transformers с подобен размер и качеството на Mamba-3B съвпада с това на Transformers, два пъти по-големи от него (напр. 4 пункта по-високо средно при здрави разумни разсъждения в сравнение с Pythia-3B и дори надвишава Pythia-7B ).