SynJax: Distribuições de Probabilidade Estruturadas para JAX

SynJax: Structured Probability Distributions for JAX

August 7, 2023

Autores: Miloš Stanojević, Laurent Sartran

cs.AI

Resumo

O desenvolvimento de bibliotecas de software para aprendizado profundo possibilitou avanços significativos na área, permitindo que os usuários se concentrassem na modelagem, enquanto a biblioteca se encarregava da tarefa tediosa e demorada de otimizar a execução para aceleradores de hardware modernos. No entanto, isso beneficiou apenas tipos específicos de modelos de aprendizado profundo, como os Transformers, cujas primitivas se mapeiam facilmente para a computação vetorizada. Os modelos que explicitamente consideram objetos estruturados, como árvores e segmentações, não se beneficiaram da mesma forma, pois exigem algoritmos personalizados que são difíceis de implementar de forma vetorizada. O SynJax aborda diretamente esse problema ao fornecer uma implementação eficiente e vetorizada de algoritmos de inferência para distribuições estruturadas, abrangendo alinhamento, etiquetagem, segmentação, árvores de constituintes e árvores de abrangência. Com o SynJax, podemos construir modelos diferenciáveis em larga escala que explicitamente modelam a estrutura nos dados. O código está disponível em 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.