SynJax: Distribuciones de Probabilidad Estructuradas para JAX

SynJax: Structured Probability Distributions for JAX

August 7, 2023

Autores: Miloš Stanojević, Laurent Sartran

cs.AI

Resumen

El desarrollo de bibliotecas de software de aprendizaje profundo permitió avances significativos en el campo al permitir que los usuarios se centraran en el modelado, mientras que la biblioteca se encargaba de la tarea tediosa y que consume tiempo de optimizar la ejecución para los aceleradores de hardware modernos. Sin embargo, esto ha beneficiado solo a tipos particulares de modelos de aprendizaje profundo, como los Transformers, cuyas primitivas se mapean fácilmente al cálculo vectorizado. Los modelos que explícitamente tienen en cuenta objetos estructurados, como árboles y segmentaciones, no se beneficiaron de igual manera porque requieren algoritmos personalizados que son difíciles de implementar en forma vectorizada. SynJax aborda directamente este problema al proporcionar una implementación vectorizada eficiente de algoritmos de inferencia para distribuciones estructuradas que cubren alineación, etiquetado, segmentación, árboles de constituyentes y árboles de expansión. Con SynJax podemos construir modelos diferenciables a gran escala que modelan explícitamente la estructura en los datos. El código está disponible en 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.