SynJax: Структурированные вероятностные распределения для JAX

SynJax: Structured Probability Distributions for JAX

August 7, 2023

Авторы: Miloš Stanojević, Laurent Sartran

cs.AI

Аннотация

Разработка библиотек для глубокого обучения способствовала значительному прогрессу в этой области, позволяя пользователям сосредоточиться на моделировании, в то время как библиотека берет на себя утомительную и трудоемкую задачу оптимизации выполнения для современных аппаратных ускорителей. Однако это принесло пользу лишь определенным типам моделей глубокого обучения, таким как Transformers, чьи примитивы легко отображаются на векторные вычисления. Модели, которые явно учитывают структурированные объекты, такие как деревья и сегментации, не получили равных преимуществ, поскольку они требуют специализированных алгоритмов, которые сложно реализовать в векторной форме. SynJax напрямую решает эту проблему, предоставляя эффективную векторную реализацию алгоритмов вывода для структурированных распределений, охватывающих выравнивание, разметку, сегментацию, деревья составляющих и остовные деревья. С помощью SynJax можно создавать крупномасштабные дифференцируемые модели, которые явно учитывают структуру данных. Код доступен по адресу https://github.com/deepmind/synjax.

English

The development of deep learning software libraries enabled significant progress in the field by allowing users to focus on modeling, while letting the library to take care of the tedious and time-consuming task of optimizing execution for modern hardware accelerators. However, this has benefited only particular types of deep learning models, such as Transformers, whose primitives map easily to the vectorized computation. The models that explicitly account for structured objects, such as trees and segmentations, did not benefit equally because they require custom algorithms that are difficult to implement in a vectorized form. SynJax directly addresses this problem by providing an efficient vectorized implementation of inference algorithms for structured distributions covering alignment, tagging, segmentation, constituency trees and spanning trees. With SynJax we can build large-scale differentiable models that explicitly model structure in the data. The code is available at https://github.com/deepmind/synjax.