mlsanityData to 14 Oct 2026

SynJax: Structured Probability Distributions for JAX

arXiv 2308.03291 ↗First posted 7 Aug 2023cs.LG cs.AI cs.CLcited by 2

Not on the map. Why?

Abstract (arXiv)

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. This is done by exploiting the connection between algorithms for automatic differentiation and probabilistic inference. With SynJax we can build large-scale differentiable models that explicitly model structure in the data. The code is available at https://github.com/google-deepmind/synjax

Key facts. Cited by 2 papers in the map (data through 14 Oct 2026). Builds on 37 of its 87 references that are in the map. First posted on arXiv 7 Aug 2023.

Builds on

37 of 87 references.

References · 37

  1. 2016End-to-end Sequence Labeling via Bi-directional LSTM-CNNs-CRFcited by 341
  2. 2017Soft-DTW: a Differentiable Loss Function for Time-Seriescited by 271
  3. 2019Differentiation of Blackbox Combinatorial Solverscited by 226
  4. 2019Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyrocited by 191
  5. 2018Learning Latent Permutations with Gumbel-Sinkhorn Networkscited by 175
  6. 2021Equinox: neural networks in JAX via callable PyTrees and filtered transformationscited by 148
  7. 2019Interpretable Neural Predictions with Differentiable Binary Variablescited by 126
  8. 2018Multilingual Constituency Parsing with Self-Attention and Pre-Trainingcited by 117
  9. 2016Recurrent Neural Network Grammarscited by 109
  10. 2017TensorFlow Distributionscited by 106
  11. 2020Einsum Networks: Fast and Scalable Learning of Tractable Probabilistic Circuitscited by 103
  12. 2021Implicit MLE: Backpropagating Through Discrete Exponential Family Distributionscited by 98
25 more
  1. 2017Structured Attention Networkscited by 95
  2. 2013A Differential Approach to Inference in Bayesian Networkscited by 93
  3. 2019Compound Probabilistic Context-Free Grammars for Grammar Inductioncited by 76
  4. 2017Graph Convolutional Encoders for Syntax-aware Neural Machine Translationcited by 70
  5. 2020Gradient Estimation with Stochastic Softmax Trickscited by 67
  6. 2017A Minimal Span-Based Neural Constituency Parsercited by 59
  7. 2018Differentiable Dynamic Programming for Structured Prediction and Attentioncited by 58
  8. 2022Transformer Grammars: Augmenting Transformer Language Models with Syntactic Inductive Biases at Scalecited by 41
  9. 2020Torch-Struct: Deep Structured Prediction Librarycited by 40
  10. 2018AMR Parsing as Graph Prediction with Latent Alignmentcited by 36
  11. 2019Temporal Parallelization of Bayesian Smootherscited by 33
  12. 2019Cooperative Learning of Disjoint Syntax and Semanticscited by 27
  13. 2016Unsupervised Neural Hidden Markov Modelscited by 19
  14. 2018Differentiable Perturb-and-Parse: Semi-Supervised Parsing with a Structured Variational Autoencodercited by 19
  15. 2017Predicting Target Language CCG Supertags Improves Neural Machine Translationcited by 15
  16. 2020Efficient Computation of Expectations under Spanning Tree Distributionscited by 11
  17. 2020Please Mind the Root: Decoding Arborescences for Dependency Parsingcited by 11
  18. 2019LinearFold: linear-time approximate RNA folding by 5'-to-3' dynamic programming and beam searchcited by 10
  19. 2022Dynamic Programming in Rank Space: Scaling Structured Inference with Low-Rank HMMs and PCFGscited by 10
  20. 2020Latent Template Induction with Gumbel-CRFscited by 9
  21. 2016Segmental Recurrent Neural Networks for End-to-end Speech Recognitioncited by 7
  22. 2020Understanding the Mechanics of SPIGOT: Surrogate Gradients for Latent Structure Learningcited by 6
  23. 2022PGMax: Factor Graphs for Discrete Probabilistic Graphical Models and Loopy Belief Propagation in JAXcited by 6
  24. 2018Inducing Grammars with and for Neural Machine Translationcited by 3
  25. 2022Unbiased and Efficient Sampling of Dependency Treescited by 2

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 2.

2 papers.

  1. 2026Prefix Parsing is Just Parsingcited by 0
  2. 2026Syntactic Belief Update as the Driver of Garden Path Processing Difficultycited by 0