МАМБА с нуля: Нейронные сети лучше и быстрее трансформеров

Mamba представляет собой прорывную архитектуру нейронной сети, которая превосходит transformers за счет резкого снижения вычислительных требований и масштабирования с использованием O(n log(n)), а не O(n2). Его многообещающая производительность, продемонстрированная на моделях с несколькими миллиардами параметров, указывает на потенциал для работы с гораздо большими объемами контекста. Дизайн использует информацию, полученную из моделей пространства состояний, и в то же время тесно связан с простотой рекуррентных нейронных сетей, предлагая доступный путь к продвинутому языковому моделированию.

Рекуррентные нейронные сети преодолевают ограничения сверточных уровней, передавая предыдущий результат вместе с текущим входным, позволяя каждому результату инкапсулировать информацию из всех предыдущих входных данных. Этот подход позволяет интегрировать долгосрочные зависимости без дополнительных вычислительных затрат. Однако последовательный характер вычислений препятствует эффективному распараллеливанию на современном оборудовании, а проблемы с обучением ограничивают возможности сети обрабатывать только несколько десятков входных данных одновременно. Эти ограничения привели к доминированию сверточных нейронных сетей и преобразователей в практических приложениях.

Линейные рекуррентные нейронные сети открывают новые возможности в обработке длинных последовательностей, заменяя сложные нелинейные рекурренции простыми линейными функциями. Этот инновационный подход дополняет линейный уровень нейронными сетями с разбивкой по элементам, точно отражая стратегию, используемую в transformers для восстановления нелинейной обработки. Позволяя выполнять параллельные вычисления за O(log(n)) время, этот метод эффективно преодолевает традиционные проблемы RNN и устанавливает новый уровень эффективности моделирования последовательностей.

Ускорение линейных рецидивов с помощью параллельного сканирования Линейные повторения интерпретируются как операции сканирования, в которых каждый результат строится на основе предыдущего состояния с использованием ассоциативного оператора, который объединяет матричные произведения и векторные дополнения. Метод вычисляет кумулятивные операции параллельно, итеративно удваивая размер шага, тем самым сокращая последовательный процесс до логарифмических (n) итераций. Этот подход заменяет последовательное сложение параллельными сканированиями, сокращая общее время обработки с O(n) до O(log n).

Диагонализация матрицы для эффективного вычисления Простое параллельное сканирование приводит к дорогостоящему умножению матриц, увеличивая время вычислений с O(d ^ 2) до O (d ^ 3) при каждом обновлении. При разложении рекуррентной весовой матрицы на обратимую матрицу, диагональную матрицу и ее инверсию сложные вычисления заменяются эффективным поэлементным умножением. Такая диагонализация преобразует повторение в операции над сложными векторами, значительно ускоряя обработку и повышая выразительность модели за счет сложной арифметики без полной инверсии матриц.

Рекуррентные нейронные сети сталкиваются с неизбежной проблемой, когда один и тот же вес применяется многократно, что приводит к исчезновению или резкому увеличению градиентов. Это экспоненциальное изменение происходит потому, что умножение веса на множество шагов приводит к быстрому уменьшению, если значения меньше единицы, или к неконтролируемому росту, если они превышают единицу. Такая нестабильность серьезно ограничивает способность сети эффективно использовать длинный контекст - проблема, которая в стандартных нейронных сетях решается с помощью распределенных весов, зависящих от уровня.

Новая стратегия инициализации устанавливает веса в сложной полярной форме, где значения жестко контролируются вблизи единицы, а входные данные поддерживаются на уровне, близком к нулю. Этот подход ограничивает значения веса с помощью функции, которая обеспечивает стабильность градиентов, устраняя проблемы с исчезновением и взрывом. При инициализации весов, близких к стабильному значению, модель успешно изучает контекст на протяжении десятков тысяч шагов. Оценки в задачах долгосрочного анализа, таких как классификация плоских изображений в задаче PathX, показывают, что линейная схема RNN превосходит трансформаторы.

Модели пространства состояний, полученные на основе теории управления для дискретизации непрерывных систем, по сути, представляют собой линейные RNN с уникальной схемой инициализации, которая устанавливает повторяющиеся веса, близкие к единице, для обеспечения стабильности. В моделях используется параметризация, при которой веса инициализируются малыми значениями, что обеспечивает производительность, аналогичную линейным RNN, при выполнении долгосрочных задач. Однако их неспособность выборочно забывать предыдущую информацию означает, что они либо отбрасывают, либо бесконечно накапливают входные данные, что приводит к снижению производительности в языковом моделировании по сравнению с преобразователями.

Сеть может выборочно забывать или сохранять информацию, динамически вычисляя весовые коэффициенты для каждого входного вектора вместо использования фиксированных весовых коэффициентов. Линейная функция преобразует каждый входной сигнал в уникальный весовой вектор, позволяя некоторым входным данным получать значения, близкие к 0, для удаления сохраненных данных, а другим - значения, близкие к 1, для их сохранения. Затем при повторном сканировании эти пользовательские значения интегрируются для определения обновлений памяти. Этот подход также позволяет избежать исчезновения или резкого увеличения градиентов, обеспечивая различные стабильные конфигурации сети.

Mamba увеличивает размер выходного вектора в 16 раз, что позволяет собирать больше исторической информации, прежде чем возвращать вектору его первоначальное измерение. Конструкция основана на том факте, что передача данных из основной памяти в высокопроизводительную является узким местом в вычислениях, поэтому она передает только векторы стандартного размера, в то время как вычисления больших блоков выполняются исключительно в высокопроизводительной памяти. Этот метод сокращает время передачи данных, сокращая при этом дополнительные вычислительные затраты, что позволяет эффективно использовать обогащенные векторы с минимальными затратами. Следовательно, Mamba превосходит transformers в языковом моделировании при меньшей сложности, работая при O(n log(n)) по сравнению с O(n2).

Статья Mamba, в которой была продемонстрирована превосходная производительность языкового моделирования при меньших вычислительных затратах по сравнению с transformers, была отклонена ICLR из-за противоречивых отзывов. Критики обвинили его в том, что он не учитывает несвязанный долгосрочный тест, и подвергли сомнению его показатели оценки, несмотря на то, что в стандартных последующих задачах результаты были получены с нулевым результатом. Неверные толкования использования памяти еще больше подогрели споры, подчеркнув обеспокоенность по поводу целостности академической экспертной оценки в машинном обучении.