[DeepMind AlphaDev] Faster sorting algorithms discovered using deep reinforcement learning
Daniel J. Mankowitz, Andrea Michi, Anton Zhernov, Marco Gelmi, Marco Selvi, Cosmin Paduraru, Edouard Leurent, Shariq Iqbal, Jean-Baptiste Lespiau, Alex Ahern, Thomas Köppe, Kevin Millikin, Stephen Gaffney, Sophie Elster, Jackson Broshear, Chris Gamble, Kieran Milan, Robert Tung, Minjae Hwang, Taylan Cemgil, Mohammadamin Barekatain, Yujia Li, Amol Mandhane, Thomas Hubert, Julian Schrittwieser, Demis Hassabis, Pushmeet Kohli, Martin Riedmiller, Oriol Vinyals & David Silver
Статья: https://www.nature.com/articles/s41586-023-06004-9
Пост: https://www.deepmind.com/blog/alphadev-discovers-faster-sorting-algorithms
(Псевдо)код: https://github.com/deepmind/alphadev
Это в каком-то смысле продолжение истории про AlphaTensor (https://hottg.com/gonzo_ML/1146). Там семейство алгоритмов AlphaZero применяли для поиска более быстрого алгоритма матричного умножения, а здесь его применили для поиска алгоритмов сортировки. Полученная система или агент, AlphaDev, нашла новый оптимизированный алгоритм сортировки. Код найденного алгоритма интегрировали в стандартную C++ библиотеку алгоритмов в LLVM (https://reviews.llvm.org/D118029). В этой части библиотеки не было изменений больше десятка лет и это там первый алгоритм, найденный RL. Делаем ставки, через сколько лет в libc++ не останется кода, написанного человеком.
Что интересно, работа эта не новая. Код в libc++ засабмитили ещё в начале прошлого года, а сама работа видимо сделана была и того раньше. В Nature статью засабмитили в июле 2022. AlphaTensor тоже, кстати, был засабмичен за год до самой статьи, ещё в октябре 2021.
Авторы фокусировались на алгоритмах для коротких последовательностей (3-5 элементов) и ищут функции для fixed sort (строго заданное количество элементов, 3, 4 или 5) и variable sort (любое количество элементов от нуля до 3-5). Поиск идёт на уровне ассемблерных инструкций (подмножество инструкций x86: mov/cmov/cmp/jump) с нуля, не с какого-то готового алгоритма. Как и AlphaTensor, AlphaDev основан на AlphaZero. Для использования RL задачу оформили в виде игры c одним игроком AssemblyGame (в случае AlphaTensor был TensorGame).
На каждом ходу AlphaDev смотрит на уже сгенерированный алгоритм (его репрезентацию P_t) и на состояние процессора (регистры и память, Z_t), и делает ход, выбирая инструкцию для добавления в алгоритм. Это делается через MCTS (Monte-Carlo Tree Search), которая по входному состоянию S_t = <P_t, Z_t> выдаёт полиси (распределение действий) и value function предсказывающую cumulative returns из текущего состояния. У value function две головы, одна для latency, другая для корректности.
Пространство решений здесь огромно, поэтому применялись правила для прунинга действий (типа, запрещены последовательные инструкции сравнения, ячейки памяти читаются строго в инкрементальном порядке, регистры аллоцируются также, и т.д.).
Также заявлено, что разрешены и другие трансформации типа перестановки двух инструкций, случайные замены опкода и/или операнда. Гиперпараметрами оптимизировали веса этих других трансформаций.
Алгоритм валидируется на тестовых кейсах, где его выход должен совпасть с эталонным, и вознаграждение агент получает за корректность и скорость (либо длину как прокси для latency) алгоритма.
Для вычисления репрезентации собрали отдельную сеть из двух энкодеров. Один генерит эмбеддинг алгоритма, второй эмбеддинг состояния процессора и памяти. Оба эмбеддинга затем объединяются в итоговую репрезентацию.
Эмбеддинг алгоритма делается трансформером Multi-Query Attention (в котором K и V шарятся между головами, https://arxiv.org/abs/1911.02150), 6 слоёв, 4 головы. Инструкции кодируются опкодом и операндами через one-hot и дают входную последовательность. В доп.материалах также описывается графовый энкодер через GNN.
Эмбеддинг состояния делается двуслойным MLP с relu, от состояния всех регистров и памяти.
>>Click here to continue<<
