SynJax: JAX向け構造化確率分布

SynJax: Structured Probability Distributions for JAX

August 7, 2023

著者: Miloš Stanojević, Laurent Sartran

cs.AI

要旨

ディープラーニングソフトウェアライブラリの開発は、ユーザーがモデリングに集中できるようにし、現代のハードウェアアクセラレータ向けの実行最適化という面倒で時間のかかるタスクをライブラリに任せることで、この分野に大きな進展をもたらしました。しかし、これはTransformerのような特定のタイプのディープラーニングモデルにのみ恩恵をもたらしました。これらのモデルのプリミティブは、ベクトル化された計算に容易にマッピングできるためです。一方で、木構造やセグメンテーションなどの構造化されたオブジェクトを明示的に考慮するモデルは、ベクトル化形式で実装することが難しいカスタムアルゴリズムを必要とするため、同等の恩恵を受けられませんでした。 SynJaxは、アラインメント、タグ付け、セグメンテーション、構成木、およびスパニングツリーをカバーする構造化分布の推論アルゴリズムの効率的なベクトル化実装を提供することで、この問題に直接取り組んでいます。SynJaxを使用することで、データ内の構造を明示的にモデル化する大規模な微分可能モデルを構築できます。コードは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.