SynJax: Distribuzioni di Probabilità Strutturate per JAX

SynJax: Structured Probability Distributions for JAX

August 7, 2023

Autori: Miloš Stanojević, Laurent Sartran

cs.AI

Abstract

Lo sviluppo di librerie software per il deep learning ha permesso progressi significativi nel campo, consentendo agli utenti di concentrarsi sulla modellazione, mentre la libreria si occupa del compito noioso e dispendioso in termini di tempo di ottimizzare l'esecuzione per gli acceleratori hardware moderni. Tuttavia, ciò ha beneficiato solo particolari tipi di modelli di deep learning, come i Transformers, le cui primitive si mappano facilmente al calcolo vettorizzato. I modelli che tengono esplicitamente conto di oggetti strutturati, come alberi e segmentazioni, non hanno beneficiato allo stesso modo perché richiedono algoritmi personalizzati difficili da implementare in forma vettorizzata. SynJax affronta direttamente questo problema fornendo un'implementazione vettorizzata efficiente di algoritmi di inferenza per distribuzioni strutturate che coprono allineamento, etichettatura, segmentazione, alberi di costituenza e alberi di copertura. Con SynJax possiamo costruire modelli differenziabili su larga scala che modellano esplicitamente la struttura nei dati. Il codice è disponibile all'indirizzo 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.