RWKV-X: Революция в моделировании длинных контекстов с помощью разреженного внимания и рекуррентной памяти

Проблемы масштабирования моделей на базе трансформеров

Большие языковые модели (LLM), основанные на архитектуре трансформеров, сталкиваются с серьезными проблемами при обработке длинных контекстов из-за квадратичной сложности вычислений по длине последовательности. Для решения этих задач разработаны альтернативы, такие как модели с линейным вниманием, модели с пространством состояний (например, Mamba), линейные RNN (DeltaNet) и RWKV. Однако такие линейные архитектуры испытывают трудности с пониманием очень длинных контекстов. Например, модель RWKV-7 (2,9 млрд параметров) показывает высокую точность при извлечении ключей до 28K токенов, но быстро теряет эффективность за этим пределом, даже при непрерывном дообучении на данных длиной 128K токенов.

Возникновение моделей с линейной сложностью

Модели с линейной сложностью становятся перспективной альтернативой трансформерам, которые требуют квадратичных вычислительных ресурсов для длинных последовательностей. Серия моделей RWKV совмещает параллельность трансформеров во время обучения с представлением состояния, подобным RNN. Модель эволюционировала от RWKV-4 до RWKV-7, улучшая эффективность и производительность. Другие гибридные модели, такие как Jamba, Zamba и MiniMax, также предлагают уникальные конструкции. Механизмы разреженного внимания (Native Sparse Attention) организуют токены в временные блоки с тремя путями внимания: сжатые токены крупного масштаба, выборочно сохраняемые мелкомасштабные токены и скользящие окна для локального контекста. Другие механизмы внимания включают SeerAttention и Block Attention (MoBA).

Представляем RWKV-X: новую гибридную архитектуру

Группа исследователей из Гуандунской лаборатории искусственного интеллекта и цифровой экономики, Хохайского университета, Шэньчжэньского университета и Цинхайского университета предложила RWKV-X — архитектуру, объединяющую эффективность RWKV для коротких зависимостей с разреженным вниманием для длинных контекстов. RWKV-X достигает линейной сложности при обучении и постоянной сложности при выводе, что является значительным прорывом.

Стратегия обучения и результаты

RWKV-X сочетает блоки RWKV-7 с блоками разреженного внимания, используя подход с перемежающимся расширением блоков и инициализацией с нуля по образцу LLaMA Pro. Обучение проходит в два этапа:

Первый этап: обучение на коротких контекстах длиной 1024 токена из набора MiniPile, при этом заморожены все параметры, кроме добавленных блоков.
Второй этап: непрерывное дообучение на длинных контекстах с использованием ProLong-64K (64K токенов), обработка около 1 миллиарда токенов с оптимизацией всех параметров.

Для обучения используется функция потерь Long-context Cross-Entropy (LongCE), которая динамически взвешивает токены по их значимости.

Оценка производительности

В задачах с коротким контекстом RWKV-X показывает конкурентоспособные результаты. Модель меньшего размера (0,22 млрд параметров) достигает 51.0, близко к RWKV-7 (51.8). Большая модель (3,6 млрд) показывает 71.9, почти сравнима с RWKV-7 (2,9 млрд) и Qwen2.5-3B, превосходя LLaMA3.2-3B. Анализ эффективности демонстрирует превосходство RWKV-X при работе с длинными последовательностями: при 128K токенах скорость выше Flash-Attention v3 в 1,37 раза, и преимущество растет с увеличением длины контекста.

Ограничения и перспективы

У RWKV-X есть ограничения. Разреженный механизм внимания основан на эвристическом выборе топ-k чанков, что может приводить к пропуску семантически важных связей. Кроме того, декодирование с разреженным вниманием работает медленнее, чем у стандартного RWKV, что требует дальнейших инженерных улучшений.

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

Подробности доступны в оригинальной статье, а новости публикуются в Twitter.