RWKV: Reinventing RNNs for the Transformer Era

уже упоминали эту интересную статью, своего рода тоже база

можно ли попробовать пересобрать рнн в контексте трансформера и получить преимущества в виде параллелизации обучения (трансформер) и рекурсивного инференса, где в памяти можно хранить только один объект скрытой репрезентации (рнн)? - да, получится свертка

so, авторы отталкиваются от идеи attention-free трансформеров, где outer product (per-token attention) между запросами и ключами заменяют на сумму между обучаемой матрицей W и ключами K - якобы предполагают, что в этой матрице и так может содержаться вся релевантная инфа, которую просто надо сложить с ключами (что уже не так конечно, ибо произведение в линале куда мощнее сложения)

сначала они добавляют относительное позиционное кодирование в вычислении “attention-free аттеншна” к этой обучаемой матрице W, а затем добавляют новые понятия, аналогичные ключам K и значениям V + receptance вектор R как приемник информации прошлого (не путать с хидден стейтом, тут он комплексный и зависит от всех этих терминов) → метод и назвали по первым буквам этих модулей - RWKV

разделяют на два подэтапа - где time mixing & channel mixing (хотя и там и там миксуют каналы но не суть), не что иное, как FFN, в которые скармливаем токены на данном и предыдущем таймстепах + разные нелинейности применяем. а в качестве хидден стейта работает WKV operator, который очень хитровыделанно оперирует над предыдущими таймстепами и сохраняет идею attention-free трансформеров

что в итоге? получаем, что последовательность обрабатывается в линию относительно ее длины, мы можем это дело параллелизовать во время обучения (time-parallel mode, имеет аналогии со свертками если не полное сходство) и инферить реккурентно (отголосок ссм, которые тоже во время обучения параллелятся сверткой) + реализовали кстати куда кернелы для пущей эффективности

графики перформанса можете посмотреть в статье, вроде бы неплохо, но доля скепсиса у меня всегда остается (хоть я и не нашел как ее применить)

также есть интересные графики в аппендиксе, показывающие насколько значима информация из прошлого для разных слоев модели + есть хитмапа, которая демонстрирует интуитивно “продвижение хидден стейта по последовательности для генерации последнего токена” (в виде упоминания Эйфеля и генерации слова Париж)

👀LINK