Ура! Моя книга “Deep Learning with JAX” (в девичестве "JAX in Action") вышла в печать! Я только что получил свои бумажные копии 🙂

https://www.manning.com/books/deep-learning-with-jax

Для тех, кто не следил, JAX -- это питоновская библиотека для высокопроизводительных вычислений и large-scale ML, с отличной поддержкой ускорителей, в частности TPU.

На данный момент JAX является вполне реальной альтернативой TensorFlow и PyTorch (torch.func, в юности functorch, до сих пор пытается угнаться и всё ещё beta), и многие компании, в частности Google DeepMind, Cohere, xAI и прочие, перешли на него. На JAX созданы такие известные модели как AlphaFold, GraphCast, Gemini, Gemma, Grok, и я уже молчу сколько разного рисёча.

JAX -- это больше, чем библиотека для ML, это библиотека для очень разных высокопроизводительных, параллельных и распределённых вычислений. Не просто так его называют “NumPy на стероидах”. За пределами ML/DL, например, JAX активно используется для физических симуляций, и на GitHub есть уже огромное количество производных библиотек.

Сейчас отличное время, чтобы застолбить себе немного будущего :)

Отдельная радость должна быть для любителей функционального программирования, ибо JAX -- это первый фреймворк с большим охватом, работающий в этой парадигме. Очень прикольно использовать функции для трансформации других функций. Написали функцию для обработки одного элемента -- трансформировали в функцию для обработки батча. Написали сложную математическую функцию -- трансформировали в функцию, вычисляющую её производную. Аналогично с компиляцией и распараллеливанием. Никаких hidden state и side-effects, код чист, красив и понятен. А также БЫСТР! (см. https://x.com/fchollet/status/1735420737744507374)

Книга состоит из трёх частей на 370+ страницах.

Part 1: First steps.
Верхнеуровневое введение в JAX для менеджеров и вообще всех, рассказывающее, где и почему стоит использовать JAX. Плюс отдельная глава для тех, кто любит видеть код, где показан полный цикл реализации простой нейросети с использованием большинства фишек JAX.

Part 2: Core JAX.
Основная часть книги, где покрыты все основы JAX, шаг за шагом. От работы с массивами (тензорами), autodiff, компиляция, векторизация, параллелизация и шардирование, случайные числа (в функциональном программировании старые приёмы из NumPy не работают эффективно, зато теперь всё наглядно и воспроизводимо!) и pytrees.

Part 3: Ecosystem.
Большая глава с практическим знакомством с экосистемой высокоуровневых библиотек для DL (Flax, Optax, Orbax, CLU, …), а также примеры использования HuggingFace Transformers/Diffusers, которые давно уже добавили поддержку JAX. Также есть отдельная глава с очень верхнеуровневым и широким обзором того, что есть в JAX и вокруг за пределами нейросетевого мейнстрима.

Много крутых и умных людей читало и ревьюило мою книгу, спасибо куче GDE и не только. И отдельное спасибо Франсуа Шолле за добрые слова 🙂

“A comprehensive guide to mastering JAX, whether you’re a seasoned deep learning practitioner or just venturing into the realm of differentiable programming and large-scale numerical simulations.”
-- François Chollet, Software Engineer, Google

В общем это был прикольный опыт, я доволен результатом, надеюсь, вам тоже понравится.

Ещё отдельное спасибо всем, кто поддерживал GonzoML на Патреоне (https://www.patreon.com/GonzoML). Всем действующим платным членам нашей тесной группы я отправил коды для получения книги бесплатно (проверьте сообщения!) -- у вас будет постоянно обновляемая версия (a JAX очевидно будет меняться!) в онлайн доступе.

Не забыли ещё про KAN'ы? А тут уже KAT'ы подвезли!

Kolmogorov-Arnold Transformer
Xingyi Yang, Xinchao Wang
https://arxiv.org/abs/2409.10594

В развитие темы про "не только трансформеры" и SSM. На The Gradient попалась неплохая статья с полезной интуицией про Мамбу:

https://thegradient.pub/mamba-explained/

Что-то всё-таки происходит

https://x.com/miramurati/status/1839025700009030027?t=VID0XQCRNrLXWmGXCAkiXw&s=19

Больше констатация факта, чем новый прорыв.

Гугл дал имя своей системе 2020 года для дизайна чипов на базе RL. Теперь это AlphaChip

https://deepmind.google/discover/blog/how-alphachip-transformed-computer-chip-design/

В небольшой статье-дополнении появилось чуть больше деталей, а также выложен один предобученный чекпойнт и инструкции по предобучению. Код уже был в опенсорсе.

За прошедшие годы Гугл активно использовал эту систему для дизайна собственных чипов, TPU и Axion. Также систему использовал тайваньский MediaTek.

TWIMC

Скоро стартует курс по AI Alignment.
https://aisafetyfundamentals.com/alignment/

Дедлайн подачи заявок до 6 октября

На всякий случай, вдруг вы пропустили, технооптимизм от Сэма Альтмана. В прошлом году был технооптимизм от Марка Андриссена.

"This may turn out to be the most consequential fact about all of history so far. It is possible that we will have superintelligence in a few thousand days (!); it may take longer, but I’m confident we’ll get there.

How did we get to the doorstep of the next leap in prosperity?

In three words: deep learning worked.

In 15 words: deep learning worked, got predictably better with scale, and we dedicated increasing resources to it.

That’s really it; humanity discovered an algorithm that could really, truly learn any distribution of data (or really, the underlying “rules” that produce any distribution of data). To a shocking degree of precision, the more compute and data available, the better it gets at helping people solve hard problems. I find that no matter how much time I spend thinking about this, I can never really internalize how consequential it is."

https://ia.samaltman.com/

"It won’t happen all at once, but we’ll soon be able to work with AI that helps us accomplish much more than we ever could without AI; eventually we can each have a personal AI team, full of virtual experts in different areas, working together to create almost anything we can imagine. Our children will have virtual tutors who can provide personalized instruction in any subject, in any language, and at whatever pace they need. We can imagine similar ideas for better healthcare, the ability to create any kind of software someone can imagine, and much more."

У OpenAI сегодня DevDay, вот один человек не поленился текстовый стрим сделать:

https://simonwillison.net/2024/Oct/1/openai-devday-2024-live-blog/

Из самого интересного для меня наверное Realtime API (пока для текста и аудио) через WebSocket.

Также дистилляция моделей, файнтюн vision моделей, и столь популярный нынче prompt caching. Я в начале года писал, что к этому должно прийти, вот все и пришли :)

Were RNNs All We Needed?
Leo Feng, Frederick Tung, Mohamed Osama Ahmed, Yoshua Bengio, Hossein Hajimirsadegh
Статья: https://arxiv.org/abs/2410.01201

Продолжение возрождения рекуррентных сетей. На сей раз снова классика (RNN/LSTM/GRU), а не новомодные SSM (которые ещё и не эквивалентны RNN, в смысле находятся в более простом классе сложности, см. https://www.youtube.com/watch?v=4-VXe1yPDjk).

RNN обладают фундаментальными преимуществами в виде требований к памяти. Они линейны (от размера последовательности) при обучении и константны при инференсе. Чего не скажешь про ванильные трансформеры, у которых квадратичная и линейная сложности соответственно. Один только большой минус был у RNN -- обучение не параллелилось. Обучались они последовательно через backpropagate through time (BPTT), что для длинных последовательностей очень медленно. Здесь преимущества трансформеров в виде параллелизации при всех их недостатках относительно сложности оказалось достаточно, чтобы их обучение скейлилось, и вот мы там где мы есть -- трансформеры вытеснили рекуррентные сети из своих экологических ниш и доминируют почти везде.

Работы последних пары лет устраняют этот недостаток RNN, на свет появились LRU, Griffin, RWKV, Mamba и прочие. Всё это современное разнообразие эффективно параллелится с помощью одного и того же алгоритма -- parallel prefix scan, он же parallel prefix sum (https://www.cs.cmu.edu/~guyb/papers/Ble93.pdf).

Авторы текущей работы адаптируют parallel scan для старых добрых LSTM/GRU, устраняя зависимость input, forget и update гейтов от скрытого состояния (H). Заодно и tanh нелинейность тоже убирают (привет, LRU!). Обычные ванильные RNN не рассматривают, ссылаясь на проблемы с затухающими и взрывающимися градиентами (но вспомним LRU от DeepMind, он как раз был вариацией обычной RNN, https://www.tgoop.com/gonzo_ML/1734).

У LSTM, кстати, тоже было 100500 разных вариантов, например, peephole connections с дополнительной зависимостью гейтов от содержимого ячейки памяти -- помните, у LSTM по факту две переменные состояния, внутреннее состояние ячейки памяти (C, не видно снаружи) и скрытое состояние (H, hidden state, которое снаружи видно). По LSTM, на мой взгляд, есть два фундаментальных источника информации кроме оригинальных статей. Один -- это PhD диссертация Феликса Герса (Felix Gers, http://www.felixgers.de/papers/phd.pdf), который и добавил в архитектуру forget gate (изначально было только два других гейта) + peephole connections. Второй -- PhD диссертация Алекса Грейвса (Alex Graves, https://www.cs.toronto.edu/~graves/phd.pdf), который придумал CTC loss и многомерные RNN. Сила хороших PhD. Ну да ладно.

Авторы получают минималистичные версии LSTM и GRU (minLSTM и minGRU соответственно), которые требуют меньше параметров, параллелятся при обучении и дают хорошее качество. Надо вспомнить, что было в истории много других заходов на рекуррентные сети с быстрым параллельным обучением. Например, QRNN (https://arxiv.org/abs/1611.01576, она более отличается благодаря наличию свёрток) или SRU (https://arxiv.org/abs/1709.02755).

По сути работы авторы посмотрели на оригинальные архитектуры LSTM и GRU и убрали вещи, которые мешали реализовать их обучение через parallel scan.

В GRU убралась зависимость update gate (z) и скрытого состояния (h) от предыдущего значения h. Reset gate ушёл совсем. Затем ушёл tanh при вычислении финального значения h. Теперь нужно O(2*d_h*d_x) параметров вместо O(3*d_h(d_x + d_h)) в оригинальном GRU.

В LSTM также ушла зависимость от предыдущего состяния h в forget и input гейтах, а также в содержимом ячейки памяти c. Из вычисления c также ушёл tanh, и в итоге дропнули output gate и саму ячейку c, осталось только h. minLSTM требует O(3*d_h*d_x) параметров вместо O(4*d_h(d_x + d_h)) в LSTM.

По времени вычисления новые модели minLSTM/minGRU сравнимы с Mamba и, например, на длине последовательности в 512 элементов быстрее оригиналов в LSTM/GRU в 235 и 175 раз. На больших длинах ещё солиднее.

Памяти требуется больше (на 88%), так как надо хранить большой вычислительный граф для распараллеливания. Mamba требует на 56% больше, чем minGRU.

На задачке Selective Copy из работы про мамбу модели minLSTM, minGRU и Mamba (S6) её решают. А S4, H3 и Hyena решают лишь частично (это взято из работы про мамбу).

На RL задачах с MuJoCo locomotion (HalfCheetah, Hopper, Walker) из бенчмарка D4RL сравнивали с вариантами Decision Transformer (https://www.tgoop.com/gonzo_ML/719), включая Decision S4, Decision Mamba и (Decision) Aaren. minLSTM и minGRU лучше Decision S4 и сравнимы с Decision Transformer, Aaren и Mamba.

Языковое моделирование проверили на character-level GPT (nanoGPT) на трудах Шекспира. Лосс на тест сете у minGRU, minLSTM, Mamba, и Transformers близки. Мамба чуть хуже других, зато обучается быстрее всего (400 шагов). minGRU/minLSTM выходят на оптимум за 575/625 шагов, а трансформер за 2000.

Вообще конечно, сравнивать полученные minLSTM/minGRU с трансформерами и SSM, но не сравнить с оптимизированными RNN типа LRU, SRU или QRNN -- это очень странно. Я не удивлюсь, если они в целом не хуже. Так же было бы интересно сравнить со свежей xLSTM (https://www.tgoop.com/gonzo_ML/2624), официальный код наконец появился (https://github.com/NX-AI/xlstm).

В целом мне кажется, что очень не хватает обзорной работы, где было бы сравнение кучи известных параллельных RNN. Никто не хочет?