SynJax : Distributions de probabilité structurées pour JAX

SynJax: Structured Probability Distributions for JAX

August 7, 2023

papers.authors: Miloš Stanojević, Laurent Sartran

cs.AI

papers.abstract

Le développement de bibliothèques logicielles pour l'apprentissage profond a permis des avancées significatives dans le domaine en permettant aux utilisateurs de se concentrer sur la modélisation, tout en laissant la bibliothèque gérer la tâche fastidieuse et chronophage d'optimiser l'exécution pour les accélérateurs matériels modernes. Cependant, cela n'a profité qu'à certains types de modèles d'apprentissage profond, tels que les Transformers, dont les primitives se prêtent facilement au calcul vectorisé. Les modèles qui prennent explicitement en compte des objets structurés, comme les arbres et les segmentations, n'ont pas bénéficié de la même manière, car ils nécessitent des algorithmes personnalisés difficiles à implémenter sous forme vectorisée. SynJax s'attaque directement à ce problème en fournissant une implémentation vectorisée efficace d'algorithmes d'inférence pour des distributions structurées couvrant les alignements, l'étiquetage, la segmentation, les arbres de constituants et les arbres couvrants. Avec SynJax, nous pouvons construire des modèles différentiables à grande échelle qui modélisent explicitement la structure des données. Le code est disponible à l'adresse suivante : 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.