RWKV (RNN with Transformer attention): как комбинирует RNN и attention?

Краткий тезис

RWKV (Receptance Weighted Key Value) — это архитектура нейронной сети, которая объединяет рекуррентные вычисления (как в RNN) с механизмом внимания (как в Transformer), но без квадратичной сложности. Вместо полноценного self-attention RWKV использует линейную рекуррентную формулу с экспоненциально затухающими весами, что даёт сложность O(n) по длине последовательности и позволяет обрабатывать контекст до 16k–32k токенов. Основная инновация — WKV (Weighted Key Value) — аналог attention, реализованный через рекуррентное обновление скрытого состояния. Это делает RWKV эффективным для инференса (как RNN) и при этом способным моделировать долгосрочные зависимости (как Transformer).

1. Введение: зачем комбинировать RNN и attention?

RNN (рекуррентные нейронные сети) обрабатывают последовательности шаг за шагом, поддерживая state|скрытое состояние. Их инференс дешёв (O(n)), но они страдают от исчезающего градиента и плохо улавливают дальние зависимости. Transformer с self-attention решает эту проблему, но требует O(n²) памяти и времени, что ограничивает длину контекста.

RWKV предлагает компромисс: сохранить линейную сложность RNN, но добавить способность внимания к глобальному контексту. Это достигается заменой softmax-внимания на рекуррентное взвешивание ключей и значений.

2. Основная идея: рекуррентность + attention

В RWKV каждый слой состоит из нескольких компонентов, но ключевой — WKV (Weighted Key Value). Он вычисляет выход как взвешенную сумму значений V, где веса определяются рекуррентно через ключи K и предыдущие веса. Формально это напоминает attention|линейное внимание (linear attention), но с экспоненциальным затуханием.

Архитектура RWKV (упрощённо):

• Embedding токенов.
• Несколько RWKV-блоков, каждый содержит:
  • Time-mix (рекуррентный слой с WKV).
  • Channel-mix (FFN с активацией).
• LayerNorm и residual connections.

В отличие от Transformer, здесь нет multi-head attention с softmax — вместо этого один «головной» WKV, работающий рекуррентно.

3. Механизм WKV (Weighted Key Value)

WKV — это сердце RWKV. Он принимает три вектора на каждом шаге t:

• R (Receptance) — «восприимчивость», управляет, сколько информации пропустить.
• K (Key) — ключ, определяет вес текущего токена.
• V (Value) — значение, информация, которую нужно передать.

Выход WKV для шага t:

wkv_t = (sum_{i=1}^{t-1} exp(-(t-1-i) * w + K_i) * V_i + exp(K_t) * V_t) / (sum_{i=1}^{t-1} exp(-(t-1-i) * w + K_i) + exp(K_t))

где w — обучаемый параметр затухания (decay rate). На практике используется рекуррентная форма, чтобы не суммировать всё каждый раз.

Рекуррентная форма (из черновика):

wkv_t = (w_{t-1} * wkv_{t-1} + exp(K_t) * V_t) / (w_{t-1} + exp(K_t))

Здесь w_{t-1} — это накопленный «вес» предыдущих шагов (экспоненциально затухающая сумма exp(K_i)). Эта формула позволяет обновлять wkv за O(1) на шаг.

Интерпретация: wkv_t — это взвешенное среднее всех V, где вес каждого V_i пропорционален exp(K_i) и затухает экспоненциально с расстоянием (t-i). Это аналог attention, но без softmax и с линейной сложностью.

4. Формула рекуррентности и её интерпретация

Разберём рекуррентное обновление подробнее. Пусть на шаге t мы имеем:

• a_t — числитель (weighted sum of V)
• b_t — знаменатель (sum of weights)

Инициализация: a_0 = 0, b_0 = 0.

Для каждого t:

a_t = a_{t-1} * exp(-w) + exp(K_t) * V_t
b_t = b_{t-1} * exp(-w) + exp(K_t)
wkv_t = a_t / b_t

Здесь exp(-w) — константа затухания (0 < exp(-w) < 1). Это эквивалентно экспоненциальному скользящему среднему с обучаемым коэффициентом.

Связь с attention: если бы мы использовали softmax, веса были бы exp(K_i) / sum(exp(K_j)). В RWKV знаменатель b_t — это сумма всех exp(K_i) с затуханием, а числитель — взвешенная сумма V. Разница в том, что веса не нормализуются на каждом шаге заново, а накапливаются рекуррентно. Это даёт линейную сложность, но теряет свойство «внимания к каждому токену в равной степени» — старые токены имеют меньший вес.

5. Сравнение RWKV с Transformer и традиционными RNN

Вывод: RWKV занимает промежуточное положение — быстрее Transformer на длинных последовательностях, но сложнее в обучении из-за рекуррентности.

6. Преимущества RWKV

• Линейная сложность инференса: генерация каждого нового токена требует O(1) времени и памяти, независимо от длины контекста. Это делает RWKV идеальным для AI-агентов и чат-ботов с длинными диалогами.
• Эффективное использование памяти: не нужно хранить KV cache для всех предыдущих токенов (как в Transformer), только два скаляра a и b на слой.
• Хорошая производительность на длинных контекстах: до 32k токенов без потери качества (в версии RWKV-5).
• Простота реализации: рекуррентная формула легко реализуется на любом фреймворке.

7. Недостатки и ограничения

• Сложность обучения: из-за рекуррентности обучение требует BPTT (Backpropagation Through Time), что медленнее, чем параллельное обучение Transformer. Для длинных последовательностей градиенты могут затухать.
• Ограниченная длина контекста: хотя RWKV поддерживает до 32k, это меньше, чем у современных Transformer (например, GPT-4 — 128k). Увеличение контекста требует изменения параметра затухания w.
• Меньшая точность на коротких контекстах: из-за экспоненциального затухания RWKV может хуже улавливать точные позиционные зависимости (например, в задачах счётного характера).
• Меньшая экосистема: по сравнению с Transformer, меньше предобученных моделей и инструментов.

8. Варианты RWKV: RWKV-4, RWKV-5 (Eagle) и RWKV-6

• RWKV-4 (2023): базовая версия с WKV и рекуррентностью. Контекст до 8k.
• RWKV-5 (Eagle) (2024): улучшенная версия с многоголовым WKV (multi-head), что повышает качество. Контекст до 32k.
• RWKV-6 (2024): добавляет динамическое затухание (per-head decay) и gating (аналогично LSTM), что улучшает контроль над памятью.

Эти версии постепенно приближаются к качеству Transformer, сохраняя линейную сложность.

9. Применение RWKV в RAG и AI-агентах

В контексте Agentic RAG RWKV может служить эффективным LLM backbone для агентов, которым нужно обрабатывать длинные истории диалогов, документы и инструменты. Преимущества:

• Долгая память: агент может «помнить» весь контекст сессии без потери производительности.
• Быстрая генерация: каждый шаг агента (вызов инструмента, анализ результата) требует минимальной задержки.
• Низкие требования к GPU: RWKV можно запускать на CPU с приемлемой скоростью.

Однако для RAG, где требуется точный retrieval, RWKV может уступать Transformer в задачах, чувствительных к позиции (например, извлечение фактов из середины документа). Тем не менее, гибридные подходы (RWKV + retrieval) активно исследуются.

10. Пет-проект для закрепления

Задача: Реализовать упрощённый RWKV-слой с нуля на PyTorch и сравнить его скорость инференса с обычным Transformer-слоем при разной длине последовательности.

Инструменты: Python, PyTorch, NumPy, matplotlib (для графиков).

Шаги:

1. Реализовать класс RWKVLayer с параметрами C (размерность), decay (w). Внутри: линейные проекции для R, K, V, рекуррентное обновление a, b, вычисление wkv.
2. Реализовать класс TransformerLayer с multi-head attention (1 голова для простоты) и softmax.
3. Сгенерировать случайные последовательности длины L = 100, 500, 1000, 2000, 5000.
4. Замерить время инференса (forward pass) для каждого слоя на одинаковых данных (batch_size=1).
5. Построить график зависимости времени от L. Ожидать, что Transformer растёт квадратично, RWKV — линейно.
6. Дополнительно: проверить, что выход RWKV приближается к attention при увеличении decay (w -> 0).

Ожидаемый результат: график, наглядно показывающий преимущество RWKV при длинных последовательностях. Код можно выложить на GitHub.

11. Связь с другими вопросами

12. Навигация

• Предыдущий: 713
• Следующий: 715
• Индекс: 00. Индекс разборов

Навигация

• Предыдущий: 713
• Следующий: 715
• Индекс: 00. Индекс разборов