SynJax: Strukturierte Wahrscheinlichkeitsverteilungen für JAX

SynJax: Structured Probability Distributions for JAX

August 7, 2023

papers.authors: Miloš Stanojević, Laurent Sartran

cs.AI

papers.abstract

Die Entwicklung von Deep-Learning-Softwarebibliotheken ermöglichte bedeutende Fortschritte auf diesem Gebiet, indem sie Nutzern erlaubte, sich auf die Modellierung zu konzentrieren, während die Bibliothek die mühsame und zeitaufwendige Aufgabe der Optimierung der Ausführung für moderne Hardwarebeschleuniger übernahm. Dies hat jedoch nur bestimmte Arten von Deep-Learning-Modellen begünstigt, wie beispielsweise Transformers, deren Grundoperationen sich leicht auf vektorisierte Berechnungen abbilden lassen. Modelle, die strukturierte Objekte wie Bäume und Segmentierungen explizit berücksichtigen, profitierten nicht in gleichem Maße, da sie spezielle Algorithmen erfordern, die schwierig in vektorisierter Form zu implementieren sind. SynJax geht dieses Problem direkt an, indem es eine effiziente vektorisierte Implementierung von Inferenzalgorithmen für strukturierte Verteilungen bereitstellt, die Ausrichtungen, Tagging, Segmentierungen, Konstituentenbäume und Spannbäume abdecken. Mit SynJax können wir groß skalierbare differenzierbare Modelle entwickeln, die die Struktur in den Daten explizit modellieren. Der Code ist verfügbar unter 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.