UltraMem: Сверхразреженная память для языковых моделей

В мире современных языковых моделей, где размер и сложность моделей растут экспоненциально, возникает проблема эффективности и скорости вывода. Одним из подходов к решению этой проблемы является использование Mixture of Experts (MoE), но даже этот метод не лишен недостатков. В данной статье мы рассмотрим UltraMem — инновационную архитектуру, которая обещает значительное улучшение производительности и скорости вывода, используя сверхразреженные слои памяти.

С развитием больших языковых моделей (LLM), таких как Transformer, выяснилось, что увеличение количества параметров напрямую влияет на производительность модели, но также и на вычислительную сложность. Модели MoE попытались разрешить этот конфликт, разделяя параметры на множество "экспертов", каждый из которых специализируется на определённых типах задач, но и они сталкиваются с проблемами при выводе из-за высоких затрат на доступ к памяти.

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

Основные концепции UltraMem

Сверхразреженная память

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

Преимущества перед MoE

• Скорость вывода: UltraMem значительно уменьшает время вывода, поскольку активирует только необходимые параметры, в отличие от MoE, где все эксперты должны быть доступны для вывода.
• Масштабируемость: UltraMem демонстрирует лучшие масштабные свойства, позволяя увеличивать количество параметров без пропорционального увеличения времени вывода.

Структура UltraMem

Архитектура UltraMem включает в себя:

1. Модули памяти: Слои памяти, которые могут содержать до 20 миллионов слотов памяти, активируются только для соответствующих токенов.
2. Ключи и значения: Вместо использования всех ключей и значений, как в традиционных моделях, UltraMem использует только те, которые имеют наивысший "скор", что уменьшает количество вычислений.
3. Асинхронное выполнение: UltraMem позволяет параллельно выполнять операции с памятью и плотными слоями, что улучшает общую производительность.

Эксперименты и результаты

В ходе экспериментов были протестированы модели с различным количеством параметров, включая модели с 151M, 680M, 1.6B и 6.5B параметрами. Результаты показали:

• Улучшение производительности: UltraMem превосходит MoE по всем ключевым показателям, включая точность на различных бенчмарках.
• Скорость вывода: При одинаковом количестве параметров и вычислительных затрат, UltraMem может быть до 6 раз быстрее MoE при малых размерах батча.
• Масштабируемость: UltraMem демонстрирует сильные масштабные способности, позволяя увеличивать количество параметров без значительного ухудшения времени вывода.

Валидация и масштабирование

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

Заключение

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

Таким образом, UltraMem открывает новые возможности для создания ещё более крупных и сложных моделей без тех же затрат на вывод, что и у предыдущих архитектур. Это не только улучшает производительность текущих моделей, но и задаёт направление для будущих исследований в области LLM и AI.