RWKV-X: Новая архитектура для эффективного моделирования языка с длинным контекстом
Постоянно растущий спрос на обработку более длинных и сложных последовательностей расширил границы больших языковых моделей (LLMs). Традиционные архитектуры на основе Transformer, хотя и мощные, сталкиваются со значительными проблемами масштабирования из-за их квадратичной сложности по отношению к длине последовательности. Это ограничение становится особенно очевидным при работе с расширенными контекстными входными данными, препятствуя их способности эффективно захватывать и использовать информацию из отдаленных частей последовательности. В ответ на эту проблему возникла волна инновационных подходов, направленных на достижение линейной сложности при обработке длинных последовательностей.
Эти методы включают модели линейного внимания, модели пространства состояний (такие как Mamba), линейные RNN (такие как DeltaNet) и RWKV. Каждая из этих архитектур предлагает уникальное решение проблемы квадратичной сложности, обеспечивая более эффективную обработку длинных последовательностей. Однако эти линейные архитектуры часто сталкиваются с трудностями в полном понимании и использовании информации из длинного контекста.
Например, RWKV-7 (модель с 2,9 млрд параметров) демонстрирует высокую точность в задачах извлечения секретных ключей до 28 тыс. токенов. Однако его производительность быстро ухудшается после этого порога. Даже при непрерывном предварительном обучении с использованием данных длиной 128 тыс., ограничения длинного контекста сохраняются. Эта проблема не является уникальной для RWKV; она распространяется и на другие архитектуры, такие как Mamba, представляя собой фундаментальную проблему для этого класса моделей. Трудности с поддержанием производительности в расширенных контекстах подчеркивают важную область для улучшения в языковых моделях с линейной сложностью.
Ландшафт языковых моделей с линейной сложностью
Языковые модели с линейной сложностью появились в качестве заманчивых альтернатив архитектурам на основе Transformer, обходя квадратичное вычислительное бремя, присущее обработке длинных последовательностей. Семейство моделей RWKV, выделяющееся в этой области, искусно сочетает параллелизируемость Transformer во время обучения с рекуррентным представлением состояния, подобным RNN.
Эволюция RWKV охватывает несколько итераций, начиная с базовой RWKV-4, переходя к RWKV-5, RWKV-6 и завершаясь RWKV-7. Каждая итерация принесла усовершенствования и улучшения, повышающие возможности модели и устраняющие ограничения. Кроме того, гибридные языковые модели, такие как Jamba, Zamba и MiniMax, оставили свой след, представив уникальные гибридные конструкции, еще больше обогатив ландшафт моделей с линейной сложностью.
Стремление к эффективной обработке длинного контекста также привело к разработке инновационных механизмов внимания. Native Sparse Attention, например, организует токены во временные блоки, используя три различных пути внимания: сжатые крупнозернистые токены для глобального контекста, выборочно сохраненные мелкозернистые токены для локальных деталей и скользящие окна для захвата локальной контекстной информации. Другие известные механизмы внимания включают SeerAttention и Block Attention (MoBA), каждый из которых предлагает уникальные стратегии для обращения к релевантной информации в длинных последовательностях.
RWKV-X: Гибридная архитектура для улучшения моделирования контекста дальнего действия
Исследователи из Guangdong Laboratory of Artificial Intelligence and Digital Economy (SZ), Shenzhen, Hohai University, Nanjing, Shenzhen University, и Qinghai University, Xining, представили новую гибридную архитектуру под названием RWKV-X. Эта архитектура гениально сочетает эффективность RWKV в моделировании зависимостей ближнего действия с разреженным механизмом внимания, специально разработанным для захвата контекста дальнего действия.
В отличие от предыдущих гибридных подходов, RWKV-X достигает линейной сложности по времени во время обучения и постоянной сложности по времени во время вывода декодирования. Это делает его исключительно эффективным для обработки длинных последовательностей. Модель демонстрирует почти идеальную точность в тесте извлечения секретных ключей 64K при предварительном обучении на последовательностях длиной 64K токенов непрерывно. Он неизменно превосходит предыдущие модели RWKV-7 в тестах с длинным контекстом, сохраняя при этом высокую производительность в задачах с коротким контекстом.
Инновации в RWKV-X представляют собой значительный шаг вперед в решении проблем моделирования языка с длинным контекстом. Сочетая сильные стороны рекуррентных моделей и разреженных механизмов внимания, RWKV-X достигает баланса между эффективностью и точностью, открывая путь для более эффективной обработки расширенных последовательностей.
RWKV-X: Архитектура и обучение
RWKV-X воплощает гибридную архитектуру, интегрируя блоки RWKV-7 с разреженными блоками внимания, чтобы использовать сильные стороны обоих подходов. Вместо обучения с нуля, RWKV-X строится на основе существующих моделей с использованием чередующегося подхода к расширению блоков и механизма нулевой инициализации, вдохновленного LLaMA Pro.
Процесс обучения состоит из двух этапов, тщательно разработанных для оптимизации производительности модели как в коротком, так и в длинном контекстах:
- Предварительное обучение на коротком контексте: Первоначально модель обучается на коротких контекстах длиной 1024 токена, извлеченных из набора данных MiniPile. На этом этапе все параметры, кроме тех, что находятся в новых добавленных блоках, замораживаются, гарантируя, что предварительно полученные знания из базовой модели RWKV-7 будут сохранены. Это позволяет новым добавленным блокам адаптироваться к существующей архитектуре, не нарушая предварительно обученные представления.
- Непрерывное предварительное обучение на длинном контексте: Второй этап включает в себя непрерывное предварительное обучение на длинном контексте с использованием набора данных ProLong-64K и длины контекста 64 тыс. токенов, обрабатывая примерно 1 миллиард токенов в общей сложности. На этом этапе все параметры размораживаются и совместно оптимизируются, что позволяет модели точно настраивать свои представления и изучать зависимости дальнего действия. Обучение использует Long-context Cross-Entropy (LongCE) loss, который динамически взвешивает токены на основе их важности. Эта функция потерь помогает модели сосредоточиться на наиболее важных частях последовательности, улучшая ее способность захватывать отношения дальнего действия.
Двухэтапный процесс обучения позволяет RWKV-X эффективно сочетать эффективность RWKV-7 для моделирования ближнего действия с осознанием контекста дальнего действия разреженного механизма внимания. Сначала предварительно обучившись на коротких контекстах, а затем точно настроившись на длинных контекстах, модель учится эффективно интегрировать информацию из разных частей последовательности.
RWKV-X: Оценка и производительность
Оценка короткого контекста показывает, что RWKV-X поддерживает конкурентоспособную производительность в стандартных тестах, демонстрируя свою способность эффективно обрабатывать более короткие последовательности. Меньший RWKV-X (0,22 млрд) достигает среднего балла 51,0, что сопоставимо с 51,8 у RWKV-7. В большем масштабе RWKV-X (3,6 млрд) достигает 71,9, что близко к RWKV-7 (2,9 млрд, 72,8) и Qwen2.5-3B (71,4), превосходя LLaMA3.2-3B (69,7). Эти результаты подтверждают эффективность RWKV-X в качестве универсальной основы LLM без ущерба для производительности в более коротких контекстах.
Более того, анализ эффективности демонстрирует превосходные характеристики масштабирования RWKV-X для длинных последовательностей. При 128 тыс. токенах RWKV-X достигает ускорения в 1,37 раза по сравнению с Flash-Attention v3, и это преимущество расширяется по мере увеличения длины контекста. Это указывает на то, что RWKV-X становится все более эффективным по сравнению с другими механизмами внимания по мере увеличения длины последовательности.
Высокая производительность RWKV-X как в коротком, так и в длинном контекстах подчеркивает его универсальность и эффективность в качестве языковой модели. Его способность поддерживать конкурентоспособную производительность в более коротких последовательностях при достижении значительного ускорения в более длинных последовательностях делает его перспективной архитектурой для широкого спектра приложений.
RWKV-X: Ограничения и будущие направления
RWKV-X возникает как гибридная языковая модель, которая успешно сочетает эффективность RWKV для моделирования зависимостей ближнего действия с новым разреженным механизмом внимания, разработанным специально для моделирования контекста дальнего действия. Хотя RWKV-X демонстрирует высокую производительность и эффективность в моделировании языка с длинным контекстом, остается несколько ограничений.
Во-первых, его разреженный механизм внимания, который основан на выборе топ-k чанков, использует эвристический подход, который может игнорировать семантически релевантные зависимости. Стратегия выбора топ-k не всегда может захватывать наиболее важную информацию в последовательности, что потенциально может привести к субоптимальной производительности.
Во-вторых, текущая реализация показывает, что декодирование с разреженным вниманием работает медленнее, чем vanilla RWKV, что указывает на необходимость дальнейших инженерных усилий для оптимизации производительности. Хотя RWKV-X достигает значительного ускорения по сравнению с другими механизмами внимания в длинных последовательностях, его декодирование с разреженным вниманием все еще медленнее, чем vanilla RWKV, что предполагает, что есть возможности для улучшения его реализации.
Будущие исследования могут быть сосредоточены на устранении этих ограничений путем изучения более сложных разреженных механизмов внимания, оптимизации реализации декодирования с разреженным вниманием и изучения альтернативных стратегий обучения. Преодолев эти проблемы, RWKV-X может стать еще более мощной и эффективной языковой моделью для приложений с длинным контекстом. Кроме того, следует рассмотреть возможность использования квантования для дополнительной оптимизации и повышения эффективности модели, особенно при работе на устройствах с ограниченными ресурсами. Также, интересным направлением является исследование адаптивных техник разрежения, которые динамически настраивают структуру разреженного внимания на основе характеристик входной последовательности, потенциально улучшая как точность, так и эффективность. Важно также отметить, что масштабируемость RWKV-X с еще большими объемами данных и параметрами остается областью для дальнейшего изучения, чтобы понять его пределы и потенциал в будущем. Наконец, оценка устойчивости RWKV-X к состязательным атакам и шуму в данных имеет решающее значение для обеспечения надежности и безопасности модели в реальных сценариях.