gonzo-обзоры ML статей

RWKV: Reinventing RNNs for the Transformer Era
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Xiangru Tang, Bolun Wang, Johan S. Wind, Stansilaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Jian Zhu, Rui-Jie Zhu
Статья: https://arxiv.org/abs/2305.13048
Код: https://github.com/BlinkDL/RWKV-LM
Модели: https://huggingface.co/BlinkDL
Посты с описанием: https://johanwind.github.io/2023/03/23/rwkv_overview.html, https://johanwind.github.io/2023/03/23/rwkv_details.html

Давно хотелось написать про эту работу, которая долгое время существовала исключительно практически как народный проект в виде репы на гитхабе. И вот в мае 2023 года авторы наконец засабмитили статью на архив.

Работа продолжает традицию развития рекуррентных или шире нетрансформерных архитектур, про которые мы тут любим писать, хоть и делаем это не так часто как хотелось бы. Среди последних интересных архитектур были S4 (https://www.group-telegram.com/us/gonzo_ML.com/1424) и LEM (https://www.group-telegram.com/us/gonzo_ML.com/857). А на очереди ещё много всего интересного.

С трансформерами в целом всё хорошо кроме пресловутого квадратичного (от длины последовательности) внимания. Было много многообещающих работ про оптимизированные механизмы вплоть до линейных, но почему-то в массы они так и не пошли, что интересно. У рекуррентных сетей обычно сложность линейная, но хуже перформанс, потому что параллелятся и скейлятся они плохо (а ещё сложности с затухающими градиентами). Есть ещё шустрые свёрточные, но они больше нацелены на поиск локальных паттернов, а на больших длинах им сложно. Текущая работа предлагает новую архитектуру под названием Receptance Weighted Key Value (RWKV), обладающую достоинствами и хорошей параллелизации трансформеров, и эффективного инференса рекуррентных сетей.

Как этого добились? Рациональный дизайн!

По сути, авторы предлагают новый механизм внимания, вдохновлённый Attention Free Transformer (AFT, https://arxiv.org/abs/2105.14103, работа команды из Apple). Традиционное трансформерное QKV (query, key и value) заменяется на RWKV, где K и V -- те же самые, R -- Receptance vector отвечающий за принятие прошлой информации, а W -- это обучаемый вес (Weight) затухающий в зависимости от позиции.

Как и обычный трансформер, RWKV состоит из состыкованных блоков c residual connection, внутри которых time-mixing и channel-mixing подблоки с рекуррентностями.

В time-mixing блоке входной сигнал (эмбеддинги) x преобразуется в набор r, k и v через линейную интерполяцию текущего входа и предыдущего, и проекцию через соответствующие матрицы W_r, W_k, W_v. Например, для k формула выглядит так:

k_t = W_k * (µ_k*x_t + (1 − µ_k)*x_{t−1}), для r и v аналогично.

Такая интерполяция текущего и предыдущего входов называется time-shift mixing в time-mixing блоке и token shift в channel-mixing блоке.

Дальше вычисляется WKV, аналог QKV внимания в традиционном трансформере. В QKV веса внимания для v считались как q_i*k_j/sqrt(d), а здесь как -(t-1-i)*w+k_i. W — это обучаемый channel-wise вектор, умноженный на относительную позицию токена. Есть ещё дополнительный вектор U, введённый чтобы отдельно обращать внимание на текущий токен t, это сделано для борьбы с потенциальной дегенерацией W. От всего этого берётся softmax. При увеличении длины последовательности (времени t) в получающуюся взвешенную сумму входит всё больше элементов. Квадратичной сложности от умножения здесь нет. Затем полученный wkv поэлементно умножается на гейтирующую функцию (сигмоиду) от r (receptance). Для wkv реализовали свой кастомный CUDA кернел.

Блок channel-mixing попроще. Там аналогичным образом считаются r_t и k_t. От k_t берётся квадрат ReLU и делается проекция через матрицу W_v (которая как бы для v, value, а само v при этом не используется — вот этот момент я недопонял, почему именно так, видимо это и есть channel mixing?). И дальше так же сигмоида от r.

GitHub

GitHub - BlinkDL/RWKV-LM: RWKV (pronounced RwaKuv) is an RNN with great LLM performance, which can also be directly trained like…

RWKV (pronounced RwaKuv) is an RNN with great LLM performance, which can also be directly trained like a GPT transformer (parallelizable). We are at RWKV-7 "Goose". So it'...

www.group-telegram.com/us/gonzo_ML.com/1647

4.7K viewsJun 21, 2023 at 21:27

RWKV: Reinventing RNNs for the Transformer Era
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Xiangru Tang, Bolun Wang, Johan S. Wind, Stansilaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Jian Zhu, Rui-Jie Zhu
Статья: https://arxiv.org/abs/2305.13048
Код: https://github.com/BlinkDL/RWKV-LM
Модели: https://huggingface.co/BlinkDL
Посты с описанием: https://johanwind.github.io/2023/03/23/rwkv_overview.html, https://johanwind.github.io/2023/03/23/rwkv_details.html

Давно хотелось написать про эту работу, которая долгое время существовала исключительно практически как народный проект в виде репы на гитхабе. И вот в мае 2023 года авторы наконец засабмитили статью на архив.

Работа продолжает традицию развития рекуррентных или шире нетрансформерных архитектур, про которые мы тут любим писать, хоть и делаем это не так часто как хотелось бы. Среди последних интересных архитектур были S4 (https://www.group-telegram.com/us/gonzo_ML.com/1424) и LEM (https://www.group-telegram.com/us/gonzo_ML.com/857). А на очереди ещё много всего интересного.

С трансформерами в целом всё хорошо кроме пресловутого квадратичного (от длины последовательности) внимания. Было много многообещающих работ про оптимизированные механизмы вплоть до линейных, но почему-то в массы они так и не пошли, что интересно. У рекуррентных сетей обычно сложность линейная, но хуже перформанс, потому что параллелятся и скейлятся они плохо (а ещё сложности с затухающими градиентами). Есть ещё шустрые свёрточные, но они больше нацелены на поиск локальных паттернов, а на больших длинах им сложно. Текущая работа предлагает новую архитектуру под названием Receptance Weighted Key Value (RWKV), обладающую достоинствами и хорошей параллелизации трансформеров, и эффективного инференса рекуррентных сетей.

Как этого добились? Рациональный дизайн!

По сути, авторы предлагают новый механизм внимания, вдохновлённый Attention Free Transformer (AFT, https://arxiv.org/abs/2105.14103, работа команды из Apple). Традиционное трансформерное QKV (query, key и value) заменяется на RWKV, где K и V -- те же самые, R -- Receptance vector отвечающий за принятие прошлой информации, а W -- это обучаемый вес (Weight) затухающий в зависимости от позиции.

Как и обычный трансформер, RWKV состоит из состыкованных блоков c residual connection, внутри которых time-mixing и channel-mixing подблоки с рекуррентностями.

В time-mixing блоке входной сигнал (эмбеддинги) x преобразуется в набор r, k и v через линейную интерполяцию текущего входа и предыдущего, и проекцию через соответствующие матрицы W_r, W_k, W_v. Например, для k формула выглядит так:

k_t = W_k * (µ_k*x_t + (1 − µ_k)*x_{t−1}), для r и v аналогично.

Такая интерполяция текущего и предыдущего входов называется time-shift mixing в time-mixing блоке и token shift в channel-mixing блоке.

Дальше вычисляется WKV, аналог QKV внимания в традиционном трансформере. В QKV веса внимания для v считались как q_i*k_j/sqrt(d), а здесь как -(t-1-i)*w+k_i. W — это обучаемый channel-wise вектор, умноженный на относительную позицию токена. Есть ещё дополнительный вектор U, введённый чтобы отдельно обращать внимание на текущий токен t, это сделано для борьбы с потенциальной дегенерацией W. От всего этого берётся softmax. При увеличении длины последовательности (времени t) в получающуюся взвешенную сумму входит всё больше элементов. Квадратичной сложности от умножения здесь нет. Затем полученный wkv поэлементно умножается на гейтирующую функцию (сигмоиду) от r (receptance). Для wkv реализовали свой кастомный CUDA кернел.

Блок channel-mixing попроще. Там аналогичным образом считаются r_t и k_t. От k_t берётся квадрат ReLU и делается проекция через матрицу W_v (которая как бы для v, value, а само v при этом не используется — вот этот момент я недопонял, почему именно так, видимо это и есть channel mixing?). И дальше так же сигмоида от r.

BY gonzo-обзоры ML статей