Новый алгоритм перемножил разреженные тензоры в сто раз быстрее

Исследователи из MIT разработали библиотеку для C++, которая автоматически генерирует код, оптимизированный для выполнения операций над разреженными тензорами произвольного ранга. При вычислениях с матрицами время выполнения этого кода было сопоставимо с оптимизированными вручную программами, а для тензоров более высокого ранга она работала намного быстрее (почти в сто раз). Статья опубликована в Proceedings of the ACM on Programming Languages, посмотреть на работу программы онлайн можно на посвященном ей сайте.

При работе с данными их часто удобно представлять в виде матриц или тензоров. С помощью подобной матрицы можно описать связи между страничками в Facebook или сопоставить покупателей и их отзывы к товарам на Amazon. При этом большая часть информации, записанной таким образом, может оказаться «бесполезной». Например, тензор отзывов покупателей на Amazon содержит примерно 10¹⁹ компонент, однако отличны от нуля только 10⁹ из них. Такие тензоры называются разреженными.

Вычисления с разреженными тензорами можно значительно ускорить, если учесть тот факт, что при умножении любого числа на ноль всегда получается ноль. Кроме того, хранить результаты промежуточных вычислений неэффективно при больших объемах данных, в связи с чем нужно провести некоторую дополнительную оптимизацию. Например, вычисление выражения Σ_i,j A_ij^kB^ij + C^k должно выполняться за один шаг, а не разбиваться на отдельные операции умножения и сложения. Обычно такая оптимизация выполняется вручную, и для каждого типа операций разрабатывается собственный шаблон (так называемое ядро, kernel). 

В этой статье ученые предлагают новый инструмент оптимизации произвольных выражений с тензорами, генерирующий программный код автоматически на основе анализа структуры данных. Пользователю нужно указать только формат хранения данных и последовательность операций, которую надо оптимизировать. Чтобы продемонстрировать их подход, ученые разработали библиотеку C++, которую они назвали taco (Tensor Algebra COmpiler, Компилятор Тензорной Алгебры).

Для этого ученые разработали удобный формат хранения данных, основанный на представлении тензора в виде дерева, у которого число уровней равно рангу тензора (если не считать корень дерева за отдельный уровень). Затем они определили, как операции над тензорами (умножение и сложение) реализуются в виде операций над такими деревьями. Так, умножение двух тензоров выполняется с помощью графов, указывающих направления обхода деревьев, а суммирование реализуется слиянием двух деревьев. Наконец, ученые описали рекурсивный алгоритм, который использует эти идеи для генерации шаблонов кода.

Затем исследователи сравнили работу разработанной ими библиотеки с существующими программами. Оказалось, что с вычислениями над разреженными матрицами она справляется так же хорошо, как и все оптимизированные вручную алгоритмы, а в некоторых случаях время вычислений с помощью taco оказывалось даже меньше. При работе же с тензорами более высокого ранга taco превосходил существующие аналоги во много раз. Например, тензоры третьего ранга, сформированные на основе анализа постов в локальной сети Facebook Нового Орлеана, новый алгоритм перемножал в 114 раз быстрее, чем MATLAB Tensor Toolbox.

Ранее мы писали о том, как исследователи из MIT разработали программу, которая автоматически ищет и исправляет ошибки в исходном коде других программ. А ученые из Кембриджа научили искусственный интеллект воровать код и собирать из него собственные проекты.

Дмитрий Трунин