Compiling machine learning programs via high-level tracing

By Roy Frostig, Matthew Johnson, Chris Leary. 2018

441 citations recorded by the source corpus.

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Natural Language Processing

Citations in the corpus

Rethinking Attention with Performers (2020, methodology): “We implemented our setup on top of pre-existing Transformer training code in Jax (Frostig et al., 2018) optimized with just-in-time ( jax.jit ) compilation, and complement our theory with empirical evidence to demonstrate the practicality of FAVOR+ in multiple settings.”
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding (2020, methodology): “Data parallelism [57] is the most commonly used approach and is supported by major frameworks (TensorFlow [21], PyTorch [22], JAX [79, 80]), where devices run the same program with different input data and combine their local gradients before the weight updates.”
Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning (2020, methodology): “Our federated learning simulation code is written in Jax [7].”
Weight Agnostic Neural Networks (2019, methodology): “We explored the use of autograd packages such as JAX [24] to fine-tune individual weights of WANNs for the MNIST experiment.”
Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers (2020, methodology): “• We show that our mechanism, implemented in Jax [22], is API-compatible with the regular Transformer, whose standard dot-product attention can be replaced by FAVOR with all other components of the architecture intact.”
End-to-End Differentiable , Hardware Accelerated , Molecular Dynamics in Pure Python (2019, methodology): “Moreover, JAX MD is based on JAX (Bradbury et al., 2018; Frostig et al., 2018) which has a strong neural network ecosystem that can be used seamlessly with simulations.”
Large-Scale Differentially Private BERT (2021, methodology): “We build off the recent advances in software (XLA team and collaborators, 2017; Bradbury et al., 2018; Frostig et al., 2018) and hardware (Jouppi et al., 2017) and establish a baseline for BERT-Large pre-training with DP.”
Reverb: A Framework For Experience Replay (2021, methodology): “Beyond TensorFlow, it has also been commonly used in JAX (Frostig et al., 2018) as part of training workﬂows.”
TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learning (2019, methodology): “JAX (Frostig et al., 2018), a tracing-JIT compiler that generates code for heterogeneous devices via XLA (The XLA team, 2017), provides a similar programming paradigm; MXNet and Gluon also let users interpolate between imperative and staged computations, but at a level of abstraction that is higher ”
Scalable Gradients for Stochastic Differential Equations (2020, background): “These vector-Jacobian products have the same asymptotic time cost as evaluating the drift and diﬀusion functions, and can be easily computed by automatic diﬀerentiation [1, 15, 47, 56].”
Data Movement Is All You Need: A Case Study on Optimizing Transformers (2020, methodology): “from data optimization passes include Halide [131], JAX [82], [83], Legion [132], Lift [133], and MLIR [134].”
Efficient and Modular Implicit Differentiation (2021, unknown): “In this experiment, we revisit an example from JAX-MD [76], the problem of ﬁnding energy minimizing conﬁgurations to a system of k packed particles in a 2 -dimensional box of size U ( x i,j , θ ) , where x (cid:63) ( θ ) ∈ R k × 2 are the optimal coordinates of the k particles, U ( x i,j , θ ) is th”