Неужели наконец AMD'шные GPU можно будет нормально использовать для deep learning'а? Даже не верится.

https://huggingface.co/blog/huggingface-and-amd

In the meantime, a discussion on AGI between Joscha Bach and Connor Leahy

Video: https://www.youtube.com/watch?v=Z02Obj8j6FQ
Podcast: https://podcasters.spotify.com/pod/show/machinelearningstreettalk/episodes/Joscha-Bach-and-Connor-Leahy-on-AI-risk-e25ukc5

Очередные слухи про устройство GPT-4:

Mixture model с 220B параметров x 8 экспертов, обученных на разные распределения задач и данных. И ещё почему-то делают 16 инференсов.

Вероятно к нынешнему моменту всё отдистиллировано.

https://twitter.com/soumithchintala/status/1671267150101721090?t=lKs8fcz_Nzlda4kiEQvDsw&s=19

Вспоминается эта безумная картинка, широко ходившая по сети в недалёком прошлом

:)

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://news.1rj.ru/str/gonzo_ML/1424) и LEM (https://news.1rj.ru/str/gonzo_ML/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.

Тема с микшированием одного и другого, кажется, инспирирована MLP-Mixer (https://news.1rj.ru/str/gonzo_ML/776, там был token-mixing и channel-mixing), а история с гейтированием — gMLP (https://news.1rj.ru/str/gonzo_ML/787).

Обучение RWKV можно параллелить в time-parallel mode. Матричные умножения параллелятся как в обычных трансформерах. Для инференса RWKV можно сформулировать рекурсивно, так чтобы текущий токен зависел только от последнего состояния.

RWKV -- штука универсальная, но текущая имплементация сфокусирована на языковом моделировании (RWKV-LM). На входе слой эмбеддингов, на выходе линейная проекция на логиты выходных токенов.

Проблем с градиентами при обучении не возникает, авторы уповают на софтмакс, а также на layernorm, которые там перед подблоками стоят. Было интересное работающее решение с инициализацией матрицы эмбеддингов и прогоном её через дополнительный layernorm. Это стабилизировало и ускоряло процесс обучения. Все параметры инициализировались максимально близко к identity transformation, bias термы в линейных слоях не использовались.

Обучили модели размеров 169M, 430M, 1.5B, 3B, 7B, 14B, от 12 до 40 слоёв.

Авторы задались несколькими вопросами:

RQ1: является ли RWKV конкурентом обычным квадратичным трансформерам с сопоставимым числом параметров и токенов в обучении?
RQ2: при увеличении числа параметров остаётся ли RWKV конкурентоспособным?
RQ3: улучшает ли лосс увеличение числа параметров на тех длинах контекста, где обычные трансформеры уже не алё?

На RQ1 и RQ2 RWKV выглядит достойно на шести бенчмарках (Winogrande, PIQA, ARC-C, ARC-E, LAMBADA, SciQ) в сравнении с Pythia, OPT и BLOOM. На четырёх задачах даже побил Pythia и GPT-Neo.

По RQ3 показали, что увеличение контекста ведёт к более низкому лоссу на The Pile, то есть модель использует более далёкую информацию.

Померили скорость инференса на CPU (x86) и GPU (NVIDIA A100 80GB) в fp32. В подсчёт параметров включали всё, с эмбеддингами. Время генерации растёт существенно медленнее, чем у трансформеров, и явно линейно, а не квадратично.

Все актуальные модели выложены на HuggingFace: https://huggingface.co/BlinkDL. Текущая версия RWKV-4 (не понял, это просто так оказалось или чтобы лучше соответствовать GPT-4…).

Лучшая модель RWKV-4-World генерирует текст, код и чаты на 100+ языках и у неё лучшие способности zero-shot и in-context learning на английском.

Также есть Raven 7B/14B, файнтюн на Alpaca+ShareGPT и т.п. Есть аналог ChatGPT на RWKV (https://github.com/BlinkDL/ChatRWKV).