Введение

Умножение матриц — фундаментальная проблема вычислений, настолько же разнообразная, насколько и простая. Мы видим, как матричное умножение используется для анализа фотографий со смартфона, распознавания словесных команд, создания визуальных эффектов для компьютерных игр, моделирования погоды, сжатия данных и фильмов для обмена в Интернете и многого другого. Многие компании вложили миллионы в создание оборудования, которое может более эффективно вычислять умножение матриц. Однако со времени прорыва Штрассена в 1969 году не произошло каких-либо существенных улучшений в оптимизации умножения матриц — до сих пор.
DeepMind недавно опубликовала AlphaTensor — первая система искусственного интеллекта (ИИ) для разработки инновационных, эффективных и доказуемо правильных алгоритмов для фундаментальных задач, таких как умножение матриц.

Автоматизация алгоритмического обнаружения

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

Эта игра чрезвычайно сложна — по сравнению с Го, которую решила AlphaZero, эта игра на порядок больше на 30 величин. Для решения этой задачи было разработано несколько различных компонентов, включая новую архитектуру нейронной сети, которая включает индуктивные смещения для конкретных задач, процедуру для создания полезных синтетических данных и рецепт использования симметрии проблемы.

Алгоритмы как тензорная декомпозиция

Билинейная функция определяется как «функция, объединяющая элементы двух векторных пространств для получения элемента третьего векторного пространства, и линейная по каждому из своих аргументов». Например, пусть V, W и X — три векторных пространства над одним и тем же базовым полем F. Билинейная карта — это функция B такая, что:
B: V x W → X
Когда мы удерживаем V фиксировать и варьировать W, то должно быть линейное отображение в X. Аналогично это верно для W, удерживаемого фиксированным. Умножение матриц является примером билинейной функции и поэтому может быть полностью представлено трехмерным тензором. Тензор T_n описывает умножение n x n, где T_n не зависит от умножаемых матриц. Элементы тензора равны {0,1} и имеют размер n² x n² x n². Это может быть расширено до более общего умножения матриц n x m и m x p. Разложение T_n на R членов ранга один мы получаем [1]:

где x обозначает внешнее произведение, а u^r, v^r и w^r — все векторы. Пример умножения матрицы 2 x 2 показан ниже:

Разложение T_n на R членов первого ранга обеспечивает алгоритм умножения произвольных матриц размера n x n с использованием R скалярных умножений. Это показано в алгоритме ниже:

DRL для обнаружения алгоритмов

Как было сказано ранее, проблема поиска эффективных алгоритмов умножения матриц рассматривается как задача обучения с подкреплением — TensorGame. Шаг игры описывается тензором T_t, который изначально устанавливается равным целевому тензору (результат умножения матриц). На каждом шаге t игрок выбирает триплет (u^t, v^t, w^t), и тензор T_t обновляется путем вычитания полученного тензора [2]:

Конечная цель состоит в том, чтобы достичь нулевого тензора, поскольку к тому времени, возвращаясь назад через использованные триплеты, мы получаем последовательность, которая удовлетворяет:

где R - количество ходов. Чтобы избежать излишне длинных игр, существует максимальное количество ходов, R_limit.

За каждый сделанный шаг предусмотрена награда в размере -1, чтобы поощрять поиск кратчайшего пути к нулевому тензору. Если после шагов R_limit игра завершается с ненулевым тензором, агент получает дополнительное терминальное вознаграждение, равное -\gamma(T_{R_{limit}}), которое зависит от ранга этого конечного тензора. Кроме того, чтобы найти точное умножение матриц (u ^ {(t)}, v ^ {(t)}, w ^ {(t)}) ограничиваются дискретным набором коэффициентов, чтобы избежать точности с плавающей запятой [2] .

Чтобы играть в эту игру, AlphaTensor использует глубокую нейронную сеть для руководства процедурой планирования поиска по дереву Монте-Карло (MCTS). Сеть принимает тензор T_t для разложения и выводит политику и значение. Политика обеспечивает распределение возможных действий. Значение обеспечивает оценку распределения доходности, начиная с текущего состояния T_t. В приведенной выше схеме вознаграждения распределение моделирует мнение агента о ранге тензора T_t. Чтобы играть в игру, AlphaTensor начинает с T_n и использует MCTS на каждом этапе, чтобы выбрать следующее действие. Любые готовые игры используются в качестве обратной связи с сетью для улучшения параметров сети [2].
Поиск по дереву Монте-Карло (для заинтересованных)
Процесс MCTS можно разбить на четыре отдельных этапа [3]:

  1. Выбор. В этом процессе алгоритм MCTS проходит по текущему дереву от корневого узла, используя определенную стратегию. При обходе узел выбирается на основе некоторых параметров, которые возвращают максимальное значение. Если во время обхода найден дочерний узел, который также является конечным узлом, MCTS переходит к шагу расширения.
  2. Расширение: в этом процессе к дереву добавляется новый дочерний узел к узлу, который был достигнут во время выбора.
  3. Моделирование: в этом процессе выполняется моделирование путем выбора ходов или стратегий до тех пор, пока не будет достигнут результат или предопределенное состояние.
  4. Обратное распространение: после определения значения вновь добавленного узла необходимо обновить оставшееся дерево. Затем дерево выполняет обратное распространение от нового узла к корневому узлу. В ходе этого процесса количество симуляций, хранящихся в каждом узле, увеличивается. Кроме того, если симуляция нового узла приводит к выигрышу, количество выигрышей также увеличивается.

Повышение производительности

Основная проблема, связанная с приведенным выше описанием TensorGame, заключается в огромном пространстве действий, а именно в возможных ходах в любой момент времени t. Все компоненты, описанные ниже, улучшают общую производительность по сравнению с обычным агентом AlphaZero.

Синтетические демонстрации

Хотя тензорное разложение очень сложно, задача построения тензора из его факторов первого ранга элементарна. Мы можем сгенерировать большой набор данных пар тензор-факторизация путем случайной выборки (u^t, v^t, w^t), а затем построения тензора. Обучая сеть на комбинации контролируемой потери (имитация синтетических демонстраций) и стандартной потери обучения с подкреплением (ранее описанный процесс), мы превосходим каждую стратегию обучения по отдельности. Отметим, что это оптимизация, специфичная для билинейных функций, таких как умножение матриц, поскольку в противном случае тензор не сможет быть сформирован.

Изменение основы

T_n — тензор, представляющий билинейную операцию матричного умножения, являющуюся каноническим базисом. Та же операция может быть выражена в других базисах, что приведет к другим тензорам. Эти разные тензоры эквивалентны, поскольку имеют одинаковый ранг, а их разложения можно сопоставить с каноническим базисом. Используя изменение базиса в начале каждой игры, применяемое к T_n, в игры, в которые играет агент, вносится больше разнообразия.

AlphaTensor удалось улучшить алгоритм Штрассена. Эту эффективность можно увидеть в умножении матриц 4 x 5 на 5 x 5. Традиционный метод потребовал бы около 100 операций, метод Штрассена — 80, а новый AlphaTensor — всего 76 [2].

Улучшения AlphaTensor видны не только в небольших матрицах, но также показывают заметное улучшение в умножении больших матриц, которое чаще встречается в ML. AlphaTensor предоставил не один, а тысячи различных ранее неизвестных методов умножения матриц.

В заключение, хотя AlphaTensor был разработан для умножения матриц, следует помнить, что предпосылка AlphaTensor была построена на создании тензора из билинейной функции. Это означает, что любые билинейные операции, такие как умножение матриц, полиномиальное умножение или другие пользовательские билинейные операции, используемые в машинном обучении, могут быть представлены в виде тензоров и решаться аналогичным образом.

Ссылки

  1. Открывая новые алгоритмы с Alphatensor. RSS. (н.д.). Получено 10 ноября 2022 г. с https://www.deepmind.com/blog/discovering-novel-algorithms-with-alphatensor.
  2. Фавзи А., Балог М., Хуанг А., Хьюберт Т., Ромера-Паредес Б., Барекатаин М., Новиков А., Р. Руис Ф. Дж., Шриттвизер Дж., Свирщ, Г., Сильвер, Д., Хассабис, Д., Кохли, П. (2022). Обнаружение более быстрых алгоритмов умножения матриц с обучением с подкреплением. Природа, 610 (7930), 47–53. https://doi.org/10.1038/s41586-022-05172-4
  3. ML: поиск по дереву Монте-Карло (MCTS). Гикс для гиков. (2022, 5 июля). Получено 10 ноября 2022 г. с https://www.geeksforgeeks.org/ml-monte-carlo-tree-search-mcts/.