Тензорное программирование: новые горизонты для ИИ и науки

Развитие языка программирования FORTRAN в 1957 году произвело революцию в сфере научных и инженерных вычислений. Сложные математические операции стали выражаться лаконичным, близким к математической записи кодом, благодаря использованию массивов – структур данных, упрощающих работу с большими объемами информации. Эта идея эволюционировала в современные «тензоры», которые лежат в основе передовых систем искусственного интеллекта и научных вычислений, благодаря таким фреймворкам, как NumPy и PyTorch.

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

Облака точек в 3D-сенсорике, геометрические модели в компьютерной графике или численные симуляции в физике часто оперируют в непрерывном пространстве, где координаты могут принимать любые действительные значения. Для работы с такими данными исследователям долгое время приходилось использовать сложные структуры данных и писать громоздкие, трудноподдерживаемые программы, отделенные от основных тензорных фреймворков.

Теперь исследователи из Лаборатории компьютерных наук и искусственного интеллекта (CSAIL) Массачусетского технологического института (MIT) нашли способ интегрировать эти непрерывные наборы данных в привычный мир тензоров. Их новая система, получившая название continuous tensor abstraction (CTA), позволяет хранить и получать доступ к данным по координатам, выраженным действительными числами. Это означает, что программисты смогут писать, например, «A[3.14]», а не только «A[3]».

Система также вводит новый язык – continuous Einsums. Это обобщение нотации Эйнштейна для суммирования, широко используемой в традиционном тензорном программировании и изначально предложенной Альбертом Эйнштейном. Continuous Einsums позволяет лаконично и в виде математических формул выражать вычисления между непрерывными тензорами. Результаты работы опубликованы в издании Proceedings of the ACM on Programming Languages.

Представление непрерывных данных в виде массива кажется невозможным – ведь действительных чисел бесконечно много, а массивы конечны. Команда MIT решила эту проблему с помощью концепции piecewise-constant tensors (почленно-постоянных тензоров). Этот подход делит непрерывное пространство на управляемые регионы, где значения остаются постоянными. Это похоже на создание большого коллажа из кусочков бумаги – сохраняется основная информация, но в форме, которую аппаратное обеспечение может обработать.

Такой подход позволяет выражать ранее сложные алгоритмы в одной компактной строке тензорного кода. Задачи, которые раньше требовали тысяч строк специализированной логики – от анализа 3D-сканов LiDAR до моделирования потоков жидкости или симуляции физических систем – теперь могут быть написаны на привычном высокоуровневом языке Einsum и эффективно выполняться на современных ускорителях.

«Программы, написание которых занимало 2000 строк кода, теперь могут быть реализованы в одной строке с помощью нашего языка», — говорит Саман Амарасингхе, профессор MIT и ведущий автор исследования. «Для любой задачи, связанной с непрерывными данными, наша работа показывает, что тензорное программирование может помочь представить такие программы гораздо более разнообразными способами, чем предполагали ранее ученые».

Еще один старший автор, профессор MIT Джоэл Эмер, добавляет: «Continuous Einsums интуитивно работают так же, как и набирающие популярность Einsums для тензоров, определенных на регулярных сетках, что делает новое расширение легким для освоения программистами».

Применение непрерывных тензоров

В одном из исследований команда CSAIL обнаружила, что CTA может помочь пользователям выполнять поиск в 2D-пространствах, подобных тем, что используются в географических информационных системах (ГИС), например, Google Maps. Исследователи использовали два метода запросов: поиск по прямоугольной области (box search) и поиск по радиусу (radius search).

В геолокационных тестах CTA сгенерировала программы поиска, содержащие в 62 раза меньше строк кода по сравнению с Python-инструментом Shapely. При поиске по радиусу система CSAIL оказалась примерно в девять раз быстрее.

CTA показала еще более впечатляющие результаты при программировании алгоритмов машинного обучения для анализа закономерностей в 3D-облаках точек. Реализация «Kernel Points Convolution», требовавшая более 2300 строк кода, с помощью подхода CSAIL заняла всего 23 строки – то есть, более чем в 100 раз короче.

Система также была протестирована на способность находить особенности в определенных участках хромосом человеческого генома. CTA позволила выполнять эти запросы, сгенерировав код, который был в 18 раз короче, при этом работая в среднем немного быстрее, чем три сопоставимые базовые системы.

Другое исследование продемонстрировало полезность CTA для задач 3D-глубокого обучения, помогающих роботам «видеть» окружение перед тем, как взять предмет. Система преуспела в вычислении точек данных внутри нейронного поля излучения (NeRF) – подхода глубокого обучения, который строит 3D-сцены на основе 2D-изображений. CTA оказалась почти вдвое быстрее сопоставимого инструмента PyTorch, написав при этом примерно на 70 строк кода меньше.

Каждый эксперимент подтвердил, что тензоры могут ускорить анализ разнообразных сложных приложений. «Раньше мир тензоров и нетензорный мир развивались практически изолированно, независимо разрабатывая свои алгоритмы и структуры данных», — отмечает ведущий автор Жаеён Вон. «Используя наш язык, мы обнаружили, что многие геометрические приложения могут быть лаконично и точно выражены в continuous Einsums, при этом пользуясь теми же оптимизациями производительности, что и в компиляторах для разреженных тензоров».

«Мы считаем, что это служит мостом между двумя мирами», — добавляет Вон. «Кроме того, мы стремимся изучить, как оптимизации из нетензорного мира – например, могут ли ядра трассировки лучей ускорить тензорные программы – могут быть переосмыслены и применены через призму непрерывной тензорной нотации».

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