Моделі та платформи ШІ

Mamba: Перевизначення моделей послідовності та перевершення архітектури трансформерів

mm
Додайте Unite.AI до бажаних джерел у Google

У цій статті про Mamba ми дослідимо, як ця інноваційна модель простору стану (SSM) революціонізує моделі послідовності. Розроблена Альбертом Гу та Трі Дао, Mamba відрізняється своєю ефективністю при обробці складних послідовностей у галузях, таких як обробка мови, геноміка та аналіз аудіо. Її лінійна обробка послідовностей з вибірковими просторами стану забезпечує виняткову продуктивність у цих різних модальностях.

Ми зануримося у можливість Mamba подолати обчислювальні виклики, з якими зіштовхуються традиційні трансформери, особливо при роботі з довгими послідовностями. Її вибірковий підхід у моделях простору стану дозволяє здійснювати швидшу інференцію та лінійне масштабування з довжиною послідовності, суттєво покращуючи пропускну здатність.

Унікальність Mamba полягає в її швидкій можливості обробки, вибірковому шарі SSM та апаратно-орієнтованому дизайні, натхненному FlashAttention. Ці особливості дозволяють Mamba перевершити багато існуючих моделей, включаючи ті, що базуються на підході трансформерів, роблячи її помітним досягненням у машинному навчанні.

Трансформери проти Mamba

Трансформери, як GPT-4, встановили стандарти в обробці природної мови. Однак їхня ефективність знижується при роботі з довгими послідовностями. Саме тут Mamba виходить вперед, завдяки своїй здатності обробляти довгі послідовності більш ефективно та своїй унікальній архітектурі, яка спрощує весь процес.

Трансформери вміють обробляти послідовності даних, такі як текст для мовних моделей. На відміну від попередніх моделей, які обробляли дані послідовно, трансформери обробляють всю послідовність одночасно, дозволяючи їм захоплювати складні відносини всередині даних.

Вони використовують механізм уваги, який дозволяє моделі зосередитися на різних частинах послідовності при прогнозуванні.

Ця увага обчислюється за допомогою трьох наборів ваг: запитів, ключів і значень, отриманих з вхідних даних. Кожен елемент у послідовності порівнюється з кожним іншим елементом, надаючи вагу, яка позначає важливість, або “увагу”, яку кожний елемент повинен отримати при прогнозуванні наступного елемента у послідовності.

Трансформери підтримують два основні блоки: кодувач, який обробляє вхідні дані, і декодувач, який генерує вихідні дані. Кодувач складається з декількох шарів, кожен з яких містить два підшари: механізм багатократної уваги та просту, позиційну повністю зв’язану нейронну мережу. Нормалізація та залишкові з’єднання використовуються на кожному підшарі для допомоги у навчанні глибоких мереж.

Декодувач також має шари з двома підшарами, подібними до кодувача, але додає третій підшар, який здійснює багатократну увагу над вихідними даними кодувача. Послідовна природа декодувача забезпечує, що прогнози для позиції можуть розглядатися лише раніше позицій, зберігаючи автoregressивну властивість.

На відміну від трансформерів, модель Mamba використовує інший підхід. Хоча трансформери займаються проблемою довгих послідовностей за допомогою більш складних механізмів уваги, Mamba використовує вибіркові простори стану, надаючи більш ефективний спосіб обробки.

Ось високорівневий огляд того, як працює трансформер:

  1. Обробка вхідних даних: Трансформери спочатку кодують вхідні дані у формат, який модель може зрозуміти, часто використовуючи вкладення, які також включають позицію кожного елемента у послідовності.
  2. Механізм уваги: У своєму ядрі механізм уваги обчислює оцінку, яка представляє, наскільки потрібно зосередитися на інших частинах вхідної послідовності при розумінні поточного елемента.
  3. Архітектура кодувача-декодувача: Модель трансформера складається з кодувача для обробки вхідних даних і декодувача для генерації вихідних даних. Кожен складається з декількох шарів, які уточнюють розуміння моделі вхідних даних.
  4. Багатократна увага: У кодувачі та декодувачі багатократна увага дозволяє моделі одночасно зосередитися на різних частинах послідовності з різних представницьких просторів, покращуючи її здатність навчатися з різних контекстів.
  5. Позиційно-незалежні повністю зв’язані нейронні мережі: Після уваги проста нейронна мережа обробляє вихід кожного положення окремо та однаково. Це поєднується з вхідними даними через залишкове з’єднання та слідує за нормалізацією шару.
  6. Генерація вихідних даних: Декодувач потім прогнозує вихідну послідовність, під впливом контексту кодувача та того, що було згенеровано до цього.

Спроможність трансформера обробляти послідовності паралельно та його потужний механізм уваги роблять його потужним для завдань, таких як переклад та генерація тексту.

На відміну від цього, модель Mamba працює інакше, використовуючи вибіркові простори стану для обробки послідовностей. Цей підхід вирішує обчислювальну неефективність трансформерів при роботі з довгими послідовностями. Дизайн Mamba дозволяє здійснювати швидшу інференцію та лінійне масштабування з довжиною послідовності, встановлюючи новий парадигму для моделей послідовностей, який може бути більш ефективним, особливо при роботі з довгими послідовностями.

Mamba

Що робить Mamba真正 унікальним, це її відхід від традиційної уваги та блоків MLP. Ця спрощення веде до легшої, швидшої моделі, яка масштабується лінійно з довжиною послідовності – досягнення, яке не має аналогів серед її попередників.

Ключові особливості Mamba включають:

  1. Вибіркові SSM: Ці особливості дозволяють Mamba фільтрувати неважливу інформацію та зосередитися на важливій інформації, покращуючи її обробку послідовностей. Ця вибірковість є важливою для ефективної контекстно-залежної логіки.
  2. Апаратно-орієнтований алгоритм: Mamba використовує паралельний алгоритм, оптимізований для сучасного апаратного забезпечення, особливо GPU. Цей дизайн дозволяє здійснювати швидшу обчислення та зменшувати вимоги до пам’яті порівняно з традиційними моделями.
  3. Спрощена архітектура: Інтегруючи вибіркові SSM та ліквідуючи блоки уваги та MLP, Mamba пропонує спрощену, більш однорідну структуру. Це веде до кращої масштабованості та продуктивності.

Mamba продемонструвала вищу продуктивність у різних галузях, включаючи мову, аудіо та геноміку, перевершуючи у попередньому навчанні та галузевих завданнях. Наприклад, у мовній моделі Mamba дорівнює або перевершує продуктивність більших моделей трансформерів.

Код та попередньо навчені моделі Mamba відкрито доступні для спільноти на GitHub.

Стандартні завдання копіювання прості для лінійних моделей. Вибіркове копіювання та індуктивні голови вимагають динамічної, контекстно-залежної пам'яті для LLM.

Стандартні завдання копіювання прості для лінійних моделей. Вибіркове копіювання та індуктивні голови вимагають динамічної, контекстно-залежної пам’яті для LLM.

Структуровані моделі простору стану (S4) недавно з’явилися як перспективний клас моделей послідовностей, що поєднують риси РНН, СНН та класичних моделей простору стану. Моделі S4 походять від безперервних систем, зокрема типу системи, яка відображає одновимірні функції чи послідовності через неявний латентний стан. У контексті глибокого навчання вони представляють суттєву інновацію, надаючи новий метод для розробки моделей послідовностей, які є ефективними та високоадаптивними.

Динаміка моделей S4

SSM (S4) Це базова структурована модель простору стану. Вона приймає послідовність x та генерує вихід y за допомогою вивчених параметрів A, B, C та затримки Δ. Трансформація включає дискретизацію параметрів (перетворення безперервних функцій у дискретні) та застосування операції SSM, яка є часово-інваріантною – тобто вона не змінюється протягом різних часових кроків.

Значення дискретизації

Дискретизація є ключовим процесом, який перетворює безперервні параметри у дискретні за допомогою фіксованих формул, дозволяючи моделям S4 підтримувати зв’язок з безперервними системами часу. Це наділяє моделі додатковими властивостями, такими як інваріантність роздільності, та забезпечує належну нормалізацію, підвищуючи стабільність та продуктивність моделі. Дискретизація також проводить паралелі з механізмами затворів, знайденими в РНН, які є критичними для управління потоком інформації через мережу.

Лінійна часова інваріантність (LTI)

Ключовою особливістю моделей S4 є їх лінійна часова інваріантність. Це властивість означає, що динаміка моделі залишається сталою протягом часу, з фіксованими параметрами для всіх часових кроків. LTI є каменем у фундаменті повторюваності та конволюцій, надаючи спрощену, але потужну основу для побудови моделей послідовностей.

Подолання фундаментальних обмежень

Фреймворк S4 традиційно обмежувався своєю LTI-природою, яка створює виклики при моделюванні даних, які вимагають адаптивної динаміки. Недавня дослідницька робота представляє підхід, який подолує ці обмеження, вводячи часові параметри, тим самим усуваючи обмеження LTI. Це дозволяє моделям S4 обробляти більш різноманітний набір послідовностей та завдань, суттєво розширюючи їх застосуваність.

Термін “модель простору стану” загально охоплює будь-який повторюваний процес, що включає латентний стан, і був використаний для опису різних концепцій у різних галузях. У контексті глибокого навчання моделі S4, або структуровані SSM, відносяться до певного класу моделей, які були оптимізовані для ефективних обчислень, зберігаючи можливість моделювати складні послідовності.

Моделі S4 можуть бути інтегровані в архітектури нейронних мереж, функціонуючи як самостійні трансформації послідовностей. Їх можна розглядати як аналогічні до шарів зворотного зв’язку у СНН, надаючи основу для моделювання послідовностей у різних архітектурах нейронних мереж.

SSM проти SSM + Вибір

SSM проти SSM + Вибір

Мотивація для вибірковості у моделюванні послідовностей

Структуровані SSM

Структуровані SSM

Стаття стверджує, що фундаментальним аспектом моделювання послідовностей є стиснення контексту у керованому стані. Моделі, які можуть вибірково зосередитися на або фільтрувати вхідні дані, забезпечують більш ефективний засіб підтримання цього стисненого стану, що призводить до більш ефективних та потужних моделей послідовностей. Ця вибірковість є важливою для того, щоб моделі могли адаптивно контролювати, як інформація тече уздовж послідовності, суттєва здатність для обробки складних завдань у мовній моделі та за її межами.

Вибіркові SSM покращують традиційні SSM, дозволяючи їхнім параметрам бути залежними від вхідних даних, що вводить ступінь адаптивності, раніше недосяжної з часово-інваріантними моделями. Це призводить до часових SSM, які вже не можуть використовувати конволюції для ефективних обчислень, а замість цього покладаються на лінійний механізм повторення, суттєву відмінність від традиційних моделей.

SSM + Вибір (S6) Цей варіант включає механізм вибору, додаючи залежність від вхідних даних до параметрів B та C, та затримки Δ. Це дозволяє моделі вибірково зосередитися на певних частинах вхідної послідовності x. Параметри дискретизуються з урахуванням вибору, а операція SSM застосовується у часово-змінному порядку за допомогою операції сканування, яка обробляє елементи послідовно, динамічно регулюючи фокус протягом часу.

Видатні результати Mamba

Mamba є найкращою у своєму класі за кожним окремим результатом оцінки

Mamba є найкращою у своєму класі за кожним окремим результатом оцінки

У плані продуктивності Mamba виділяється як за швидкістю інференції, так і за точністю. Її дизайн дозволяє краще використовувати довші контексти, що демонструється як у моделюванні ДНК, так і в аудіо, перевершуючи попередні моделі на складних завданнях, які вимагають довгочасних залежностей. Її універсальність також підкреслюється у нульових оцінках по різних завданнях, встановлюючи новий стандарт для таких моделей за ефективністю та масштабованістю.

Початок роботи з Mamba

Для тих, хто цікавиться використанням Mamba, технічні вимоги включають Linux OS, NVIDIA GPU, PyTorch 1.12+ та CUDA 11.6+. Встановлення включає прості команди pip для встановлення необхідних пакетів з репозиторію Mamba. Якщо виникають питання сумісності з версіями PyTorch, використання прапора –no-build-isolation з pip може допомогти. Ці моделі, навчені на великих наборах даних, таких як Pile та SlimPajama, призначені для задоволення різних обчислювальних потреб та продуктивних стандартів.

Mamba пропонує різні рівні інтерфейсів, від вибіркового шару SSM до блоку Mamba та повних структур мовних моделей. Блок Mamba, який є основним модулем архітектури, використовує каузальний шар Conv1d та може бути легко інтегрований у дизайн нейронної мережі. Приклад використання у Python демонструє створення екземпляру моделі Mamba та обробку даних через неї, підкреслюючи простоту та гнучкість системи.

Попередньо навчені моделі Mamba доступні на Hugging Face, з розмірами від 130M до 2,8B параметрів, навчених на великому наборі даних Pile та SlimPajama. Ці моделі призначені для задоволення різних обчислювальних та продуктивних вимог, дотримуючись стандартів розмірів GPT-3. Користувачі можуть очікувати високої пропускної здатності та точності від цих моделей, роблячи Mamba конкурентоспроможним вибором для різних застосунків, включаючи, але не обмежуючись, мовну модель.

Вплив Mamba

Mamba представляє собою крок вперед у моделюванні послідовностей, пропонуючи потужну альтернативу архітектурам трансформерів для обробки даних, насичених інформацією. Її дизайн відповідає вимогам сучасного апаратного забезпечення, оптимізуючи як використання пам’яті, так і паралельні обчислювальні можливості. Відкрита доступність кодової бази Mamba та попередньо навчених моделей робить її доступною та надійною інструментом для дослідників та розробників у сфері штучного інтелекту та глибокого навчання.

Я провів останні п'ять років, занурючись у захопливий світ машинного навчання та глибокого навчання. Моя пристрасть та експертиза привели мене до внеску у понад 50 різних проектів програмної інженерії, з особливим акцентом на AI/ML. Моя тривала цікавість також привела мене до природної обробки мови, галузі, яку я бажаю дослідити далі.