SynJax: Gestructureerde Waarschijnlijkheidsverdelingen voor JAX

SynJax: Structured Probability Distributions for JAX

August 7, 2023

Auteurs: Miloš Stanojević, Laurent Sartran

cs.AI

Samenvatting

De ontwikkeling van deep learning softwarebibliotheken heeft aanzienlijke vooruitgang in het veld mogelijk gemaakt door gebruikers zich te laten concentreren op modellering, terwijl de bibliotheek de vervelende en tijdrovende taak van het optimaliseren van de uitvoering voor moderne hardwareversnellers op zich neemt. Dit heeft echter alleen bepaalde typen deep learning-modellen ten goede gekomen, zoals Transformers, waarvan de primitieven eenvoudig te mappen zijn op de vectorberekening. Modellen die expliciet rekening houden met gestructureerde objecten, zoals bomen en segmentaties, hebben niet in gelijke mate geprofiteerd, omdat ze aangepaste algoritmen vereisen die moeilijk in een vectorvorm te implementeren zijn. SynJax pakt dit probleem direct aan door een efficiënte gevectoriseerde implementatie te bieden van inferentie-algoritmen voor gestructureerde verdelingen, waaronder uitlijning, tagging, segmentatie, constituentiebomen en overspannende bomen. Met SynJax kunnen we grootschalige differentieerbare modellen bouwen die expliciet structuur in de data modelleren. De code is beschikbaar op 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.