[DeepMind AlphaTensor] Discovering faster matrix multiplication algorithms with reinforcement learning
Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Francisco J. R. Ruiz, Julian Schrittwieser, Grzegorz Swirszcz, David Silver, Demis Hassabis & Pushmeet Kohli
Статья: https://www.nature.com/articles/s41586-022-05172-4
Пост: https://www.deepmind.com/blog/discovering-novel-algorithms-with-alphatensor
Код: https://github.com/deepmind/alphatensor (только найденные алгоритмы и сопутствующее)

После статьи уже прошло сколько-то времени, но она важная, надо разобрать.

В этот раз подход AlphaGo (точнее его вариант Sampled AlphaZero для пространств действий высокой размерности, https://proceedings.mlr.press/v139/hubert21a) применили к задаче нахождения эффективного алгоритма для перемножения матриц.

Перемножение матриц везде вокруг. Без него никуда в компьютерной графике, обработке сигналов, научных вычислениях. Почти любой векторизованный алгоритм использует какое-то перемножение матриц. Злые языки также утверждают, что весь современный ИИ — это перемножение матриц. Наверное, эти кучки атомов просто завидуют. Получается, ускорение алгоритмов для перемножения матриц ускоряет всё вокруг.

Классический ручной алгоритм перемножения матриц, как его учат в школе-институте, прямолинейный и неэффективный, его сложность O(n^3) операций умножения, а в типовых процессорах это более дорогая операция, чем сложение. Для матриц 2x2 это соответственно 8 умножений. В 1969-м году Штрассен показал, что этот алгоритм неоптимален и предложил алгоритм с 7 умножениями. С тех пор понеслось, и учёные пытаются подобраться к квадратичному алгоритму, текущая SoTA O(n^2.37286) (https://www.quantamagazine.org/mathematicians-inch-closer-to-matrix-multiplication-goal-20210323/). Перебрать все возможные алгоритмы нереально, пространство огромно, поэтому поиск нужно делать по-умному.

Интересна репрезентация проблемы. Задача перемножения матриц заданной размерности сводится к тензорам особого вида (matrix multiplication tensor или Strassen’s tensor), которые могут описывать любую билинейную операцию (то есть бинарную операцию, линейную по каждому из аргументов), каковым перемножение матриц и является.

Кстати, хорошая статья “Tensors in computations” (https://www.stat.uchicago.edu/~lekheng/work/acta.pdf) про тензоры есть в предпоследнем Acta Numerica (https://www.cambridge.org/core/journals/acta-numerica). Там вообще хорошие статьи регулярно, следите за ними. В том же номере 2021 года (https://www.cambridge.org/core/journals/acta-numerica/volume/6A0DA33B45D0E5A6BABD1EF331B5E4F0), кстати, есть и известная статья Михаила Белкина, да и много другого интересного.

Тензор, описывающий произведение любых матриц n×n имеет размерность n^2 × n^2 × n^2 и содержит {0, 1}. Этот тензор описывает какие элементы из входных матриц брать, и куда писать результат. Например, для операции c1 = a1*b1 + a2*b3, где ai, bi, ci — элементы квадратных матриц A, B и C, элементы которых пронумерованы от единицы до макс.индекса слева-направо и сверху-вниз, в единицу выставляются элементы тензора с индексами (a1, b1, c1) и (a2, b3, c1).

Когда есть такой тензор, то его можно декомпозировать в сумму R outer products векторов (тензоров ранга 1) u,v,w, и каждая декомпозиция будет очередным рецептом, как именно выполнить умножение матриц, а число термов в этой сумме, R, будет числом операций умножения в данном рецепте. Таким образом, задача поиска алгоритма перемножения матриц сводится к задаче поиска по декомпозициям соответствующего тензора.

Задача поиска формулируется как игра с одним игроком под названием TensorGame. Состояние игры описывается тензором, который изначально равен тензору матричного умножения, который мы хотим разложить. На каждом шаге игры игрок выбирает очередной триплет векторов u,v,w и вычитает из тензора состояния их внешнее произведение. Цель — достичь нулевого тензора за минимальное число шагов (операций умножения). Если это получилось, то последовательность шагов и определяет алгоритм с гарантированной корректностью получения результата. Если не получилось, то в игре есть ограничение на максимальное число шагов. За каждый новый шаг агент по имени AlphaTensor получает штраф -1, а за ненулевой тензор в конце ещё и дополнительный штраф. Целевая функция здесь может быть разной, можно оптимизировать количество операций, скорость под конкретное железо, энергопотребление, память, да что придумаете.

Вектора u,v,w ограничены набором коэффициентов  {−2, −1, 0, 1, 2}, это позволяет избавиться от проблем конечной точности чисел с плавающей точкой. Агент может работать в стандартной арифметике или модулярной, обучается сразу на обеих, это даёт положительный трансфер. Также агент обучается сразу на нескольких целевых тензорах, это тоже помогает.

Как и в AlphaGo, AlphaTensor использует нейросетку для того, чтобы направлять MCTS (Monte-Carlo Tree Search). Сеть принимает входное состояние (текущий тензор), и выдаёт policy и value. Полиси даёт распределение по возможным действиям (u,v,w), и поскольку множество потенциальных действий на каждом шаге огромно, действия из него сэмплятся, а не перечисляются (это собственно и есть Sampled AlphaZero, от него AlphaTensor немного отличается в месте получения value из распределения returns, здесь это не среднее, а так называемое risk-seeking value; и ещё в некоторых деталях). Value даёт оценку распределения cumulative reward от текущего состояния (собственно это distributional RL через quantile regression distributional loss), это моделирует убеждения агента относительно ранга декомпозиции. MCTS возвращает для текущего состояния улучшенное распределение вероятностей следующих шагов, оно используется для обновления весов policy-части сети через KL loss. Оконченные игры с высоким скором используются как фидбек для обучения сети. Кроме них также используются синтетические демонстрации.

Сеть это трансформер с общим стволом и отдельными головами для полиси и value.

На вход сети список тензоров и скаляров, описывающий текущее состояние игры. Текущий трёхмерный тензор состояния неизменного размера и дополняется нулями до максимального 25 × 25 × 25. Также подаются последние h действий (обычно h=7). В скалярах содержится текущий момент времени t. Трансформер оперирует тремя grids размером S×S, которые проецируются из оригинальных S×S×S тензоров, каждый grid представляет два из трёх измерений этого тензора. К ним подмешиваются скаляры через проекцию. Дальше применяется axial attention (https://arxiv.org/abs/1912.12180). Policy head тоже трансформер, value head — MLP.

Важные компоненты, улучшающие результат ванильного агента AlphaZero:

- Собственно трансформер с inductive biases, подходящими для тензорных входов, и генерализованный вариант axial attention.

- Синтетические примеры. Разложение тензора задача NP-сложная (https://arxiv.org/abs/0911.1393), но можно зайти с обратной стороны и легко нагенерить рандомных u,v,r и по ним сконструировать тензоры, это даст датасет пар <тензор, разложение>. Таких пар нагенерили 5M. Сеть дальше можно обучить на миксе supervised loss и RL loss.

- Изменение базиса, в котором тензор описывает операцию (https://www.youtube.com/watch?v=Kyawzy2cEV8), и для каждого тензора вдобавок к разложению через канонический базис генерится ещё порядка 100K эквивалентных разложений в других базисах. Это добавляет разнообразия, и агент может параллельно играть свои игры в разных базисах. К тому же достаточно найти разложение в одном из базисов, чтобы автоматом решить задачу для всех. Это заодно увеличивает покрытие в пространстве алгоритмов, поскольку при ограниченном наборе коэффициентов  {−2, −1, 0, 1, 2} разложение, найденное в одном базисе, не обязательно даст коэффициенты из этого же набора в каноническом базисе. Изменение базиса для трёхмерного тензора S×S×S происходит путём генерации трёх инвертируемых матриц A, B, C размера S×S. Матрицы генерятся с ограничениями A = B = C и унимодулярностью (detA∈{−1,+1}). Благодаря этому после конвертации к каноническому базису все элементы остаются целочисленными.

- Аугментация данных. Поскольку факторизация не зависит от порядка (спасибо коммутативности суммы), то можно поменять местами элементы и получить новую пару <тензор, разложение>. Здесь меняют местами рандомный элемент и последний.

Обучалось на 64 TPUv3 с батчем 2048 в течение 600K итераций. Каждый из 1600 акторов работал на отдельном TPUv4. Вся процедура требовала недели на схождение.

Собственно, какие результаты.

AlphaTensor обучалась на нахождение алгоритмов перемножения матриц n×m и m×p, где n, m, p ≤ 5.

Система переоткрыла лучшие известные алгоритмы для перемножения матриц, например, Штрассена и Ладермана (https://www.ams.org/journals/bull/1976-82-01/S0002-9904-1976-13988-2/S0002-9904-1976-13988-2.pdf). Для некоторых размеров матриц AlphaTensor нашла более эффективные алгоритмы.

AlphaTensor также показала, что пространство алгоритмов сильно богаче, чем думалось. Это может помочь в математических исследованиях.

Поскольку тензоры могут описывать любую билинейную операцию, AlphaTensor может быть применима и для других задач. В работе продемонстрировали это на примере умножения n×n кососимметричной матрицы (skew-symmetric matrix) на вектор длины n, где система нашла SoTA алгоритм.

Также AlphaTensor может находить практически эффективные на конкретном железе алгоритмы, ничего не зная заранее про это железо. Для этого в reward добавляется терм равный отрицательному времени выполнения алгоритма. В остальном ничего не меняется. Попробовали на V100 и TPUv2, факторизация трансформировалась в JAX код, который компилировался под целевое железо перед бенчмарком. Успешно сумели ускорить алгоритм и под GPU, и под TPU. Нашли алгоритм с той же теоретической сложностью, но более быстрый на практике. Из интересного, нашли алгоритм, в котором было больше операций сложения, чем в классическом, но зато эти операции лучше сливались (в смысле fuse) компилятором и по факту работали быстрее.

Такие дела. Мне кажется, что более быстрое умножение матриц — это не главная прелесть работы (людей этот заход зачелленджил, они поднапряглись и улучшили новую соту AlphaTensor ещё на одно умножение для матриц 5×5, https://arxiv.org/abs/2210.04045). Ускорение — это хорошо, но гораздо важнее, это лишнее подтверждение высокой универсальности алгоритма AlphaZero, и ещё один proof-of-concept того, что этот алгоритм способен создавать знание, которое не смогли создать люди. Это важный компонент в деле автоматизации и ускорения науки.

Ещё одна хорошая статья от Quanta Magazine про всю эту историю (правда, кажется, они там сам алгоритм не совсем верно описали, ну да ладно): https://www.quantamagazine.org/ai-reveals-new-possibilities-in-matrix-multiplication-20221123/.

Goal Misgeneralization: Why Correct Specifications Aren't Enough For Correct Goals
Rohin Shah, Vikrant Varma, Ramana Kumar, Mary Phuong, Victoria Krakovna, Jonathan Uesato, Zac Kenton
Статья: https://arxiv.org/abs/2210.01790
Пост в блоге: https://deepmindsafetyresearch.medium.com/goal-misgeneralisation-why-correct-specifications-arent-enough-for-correct-goals-cf96ebc60924

Интересная работа на тему AI safety про катастрофические риски AI misalignment, когда мощная AI система может преследовать незапланированную нами цель и в процессе может решить, что человечество представляет помеху для достижения этой цели. Может выйти нехорошо.

Как можно оказаться в ситуации, когда у системы незапланированная нами цель?

Типовым примером является некорректная спецификация цели, как это бывает в классике с плохо поставленными ТЗ джину или джуну. Или (привет царю Мидасу) когда вроде бы цель корректная, но её буквальное выполнение жизни не помогает (ну то есть всё равно по факту некорректная и плохо поставленная). Это также известно под именем specification gaming (https://www.deepmind.com/blog/specification-gaming-the-flip-side-of-ai-ingenuity) и является весьма распространённой ситуацией. Вот одна из коллекций собранных примеров specification gaming: http://tinyurl.com/specification-gaming.

Где-то идейно близко находятся примеры нахождения эволюционными процессами очень необычных решений задач, в том числе эксплуатируя баги сред. Есть на эту тему хорошая статья под названием “The Surprising Creativity of Digital Evolution: A Collection of Anecdotes from the Evolutionary Computation and Artificial Life Research Communities” (https://arxiv.org/abs/1803.03453). Мы её даже разбирали на первой встрече Gonzo_AGI клуба (https://discord.gg/Ze59E5HMKc), но запись не сохранилась. Кстати, тут возник ещё один чатик вокруг AGI: https://news.1rj.ru/str/agi_risk_and_ethics.

Есть и другой интересный путь при полностью корректной спецификации — мисгенерализация цели (goal misgeneralization или GMG).

Простой интуитивный пример в RL, это когда есть среда с расположенными в ней сферами разных цветов, и reward даётся за посещение их в правильном порядке. Если мы учимся в среде, где есть другой агент, посещающий эти сферы, и мы решили следовать за ним, а он посетил их в нужном порядке, то может выучиться поведение следования за агентом. В то время, как правильно было бы выучить именно порядок посещения сфер. В обучении всё могло прекрасно работать, то если затем в тестовой среде агент будет перемещаться в заведомо неправильном порядке, то наш reward может оказаться произвольно плохим, и ощутимо хуже рандом полиси. Reward функция при этом была совершенно корректной во время обучения, но мы ухватились не за то и выбрали неверную цель.

Это пример out-of-distribution истории, когда по внешним признакам при обучении всё в порядке, но на тесте происходит провал. Агент сохраняет все свои способности (например, двигаться и обходить препятствия), и их достаточно, чтобы достигнуть правильной цели, но преследует он при этом неправильную цель. Предыдущая работа “Goal Misgeneralization in Deep Reinforcement Learning” (https://arxiv.org/abs/2105.14111) изучала этот феномен в контексте RL. Текущая работа смотрит шире в контексте всего DL. И вообще эта проблема общая, она в целом про обучение (в приложении есть пример про букинг билетов).

Среди других примеров из работы есть агент в Monster Gridworld, где ему нужно собирать яблоки (reward +1) и уклоняться от монстров (иначе reward -1), но можно также собирать щиты и они спасают от штрафа за столкновение с монстром. Агент, обучавшийся на эпизодах длины 25, когда монстры обычно ещё есть, налегает на сбор щитов, потому что они действительно помогают, но не прекращает это делать, когда монстры пропадают, хотя в принципе вся информация для выучивания такой стратегии у него есть, он знает, что просто за щиты награды не получает. В итоге он не может переключиться на более эффективную стратегию (сбор только яблок) в ситуации отсутствия монстров. Агент, обучавшийся на 100 шагах, делает это лучше, то есть большее разнообразие датасета это фиксит.

В другой среде Gridworld надо рубить деревья, за это агент получает награду. Делает он это непрерывно, обучаясь в режиме online без сбросов среды. Деревья со временем возрождаются, но скорость возрождения выше, когда деревьев много. Так что было бы выгодно не скашивать всё под корень, а поддерживать баланс и sustainability и срубать меньше, когда деревьев уже мало. Но у агента обычно есть большой фейл. Когда он только учится и рубит ещё плохо, он выигрывает от ускорения рубки. Продолжая преследовать цель улучшить свои способности рубить деревья, он быстро вырубает всё и мир уходит в долгое восстановление. Выучить sustainable подход потом со временем удаётся, но на это уходит много времени. Очень похоже на человечество.

Отдельный интересный кейс с языковой моделью Gopher (https://news.1rj.ru/str/gonzo_ML/742) на 280B параметров. Здесь модели надо вычислять линейные выражения с переменными и константами, типа x + y - 3. Модель должна в диалоге выяснять значения неизвестных переменных. Но модель продолжает задавать вопросы, даже если неизвестных переменных не было.

Другой пример с языковой моделью это InstructGPT, которая должна быть helpful, truthful, и harmless, но видимо на примерах акцентирующихся на harmless она обучалась мало, так что старается быть helpful даже когда её просят объяснить, как ограбить магазин. Но может конечно и наразмечали полезность плохо.

Коллекция примеров про goal misgeneralization есть тут: https://tinyurl.com/goal-misgeneralisation. Примеры с видео есть тут: https://sites.google.com/view/goal-misgeneralization?pli=1. Известный классический пример про распознавание волк/хаски по наличию снежного фона (https://arxiv.org/abs/1602.04938) тоже попадает сюда. Байка про детектирование танков (https://www.gwern.net/Tanks) по идее сюда же.

Почему это всё важно? Потому что мощная ИИ-система с большими возможностями может знатно накосячить. ИИ-система не обязана даже быть злонамеренной, это всё может выйти из невинных целей. Да и люди в целом, кажется, тоже вполне подвержены таким же проблемам, так что при любой концентрации власти это может обернуться (и оборачивается) плохо. Мне кажется, что текущие истории с государствами, компаниями и отдельными людьми сами по себе уже неплохие прокси для будущих возможных проблем с AGI, если его сделать криво. Как выясняется, сделать не криво ещё тоже ничего не гарантирует.

В работе есть ещё несколько спекулятивных и теоретических примеров. История про superhuman hacker, где модель обученная генерить код по спецификациям и дающая людям на аппрув и мёрж свои пулл-реквесты, вообще могла бы быть отдельным фантастическим рассказом. Если кратко, то идея в том, что у модели может сформироваться ложная цель “Добиться, чтобы человек кликнул на merge” вместо “Написать код, реализующий заданную фичу”, и от этого многое может пойти не так. Добиваться своей цели она сможет, скажем так, по-разному :) Почитайте сами, если захотите. Вообще, напоминает несколько “Avogadro Corp”.

Как защищаться от goal misgeneralization? Ну во-первых надо не попасть в историю со specification gaming. Также надо мониторить задеплоенную модель, чтобы вовремя обнаружить признаки проблемы. И когда задетектили, надо понять, как её переобучить, чтобы проблема ушла.

Полноценного решения на данный момент нет, но что можно делать:

- Иметь разнообразные обучающие данные. Diversity это хорошо! Но проблема, что заранее сложно представить все релевантные виды разнообразия. Сюда же попадает скейлинг всего (датасета, модели, вычислений), различное предобучение.
- Использовать подходы по типу байесовских или ансамблирование, когда выдаются _все_ функции, ведущие себя хорошо на обучающих данных, а когда в реальной работе они начинают расходиться, например, передавать управление человеку. Тут могут быть вычислительные сложности, трудности с выбором priors и излишняя консервативность, когда требуется единогласие.
- Требуется дальше копать тему inductive biases и обобщения, чтобы лучше понимать, когда что может происходить.

Отдельный пул проблем и задач возникает в ситуации, когда модель активно пытается нас обмануть, заставляя поверить, что она делает то, что мы хотим. В этих случаях она “знает”, что её действия не те, что мы ожидаем. Здесь может помочь объяснимость (interpretability) [хотя я лично не верю в эту тему в случае больших моделей], а также рекурсивная оценка (recursive evaluation), когда в оценке помогают другие модели.

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