4 ms·

If you're talking about Jax, there's a couple different reasons to bother for research 1. Full numpy compatibility. 2. More efficient higher order gradients (

If you're talking about Jax, there's a couple different reasons to bother for research 1. Full numpy compatibility. 2. More efficient higher order gradients (because of forward mode auto diff). Naively it's asymptotic improvement, but I believe Pytorch uses some autodiff tricks to perform higher order gradients with backwards mode, at the cost of a decently high constant factor. 3. Some cool transformations like vmap. 4. Full code gen, which is neat especially for scientific computing purposes. 5. A neat API for higher order gradients. 2. and 5. are the most appealing for DL research, 1., 3. and 4. are appealing for those in the stats/scientific computing communities. PyTorch is working on all of these, to various degrees of effort, but Jax currently has an advantage in these points (and may have a fundamental advantage in design in some).

m0zg 7y ago
From a practitioner: 1. Meh. PyTorch is close enough to not worry about it, and is better in some places. 2. Meh. All the methods people use in practice for deep learning in particular do not use higher order gradients. Most higher order methods are prohibitively memory expensive, and memory is at a premium in acceleration hardware (and so is the bus bandwidth - so you can't "swap to RAM"). I do agree that higher order gradients are the next frontier in optimization though - current optimizer research seems to have stalled, so people focus on training with huge batches and stuff like that. Most SOTA models in my field are trained with SGD+momentum - super primitive stuff. I don't see how Jax would solve the memory problem though. You still have to store those Hessians somewhere, at least partially. 3. Do agree, that's cool if it actually parallelizes nontrivial stuff which e.g. tf.vectorized_map barfs on. Although in a lot of cases you can "vectorize" by concatenating input tensors into a higher dimensional tensor. 4. Meh. Not sure why I'd want that if I already have tracing and JIT. 5. This is #2 With PyTorch though, you get close enough to Numpy to feel at home in both, and there's so much code written for it already that you can usually find a good starting point for your research pretty easily on github and then build on top of that. If you need to deploy, there's also tracing and jit, which lets you load and serve models with libtorch. I see what you're saying regarding "advantages", I'm just pointing out that PyTorch might be "good enough" for most people. If I were on that team, I'd focus on providing comfortable transition from TF 2.x which is a dumpster fire (with the exception of TensorFlow Lite which is excellent). That, IMO, would be the only way for this project to achieve mainstream success unless PyTorch disintegrates over time.
- chillee 7y ago
  I agree from a practitioner standpoint - but you were talking about research :) 1. Much of the scientific computing/stats community is stuck in the past. Many are still using Matlab! As opposed to the CS community, who are used to learning new frameworks, offering the ability to "import jax.numpy as np" and having their scripts just run is valuable to that community. As is having an API that they've only just started to become familiar with (and has way more documentation about). 2. Once again, this is true for practitioners, but not research. Hessian vector products show up in a decent amount of places. For example, if you have an inner optimization loop (a la. most meta learning approaches or Deep Set Prediction Networks) you have a Hessian Vector Product! Perhaps not prevalent in models that practitioners run but definitely something to keep an eye on in research. 3. My understanding is that it actually does a pretty decent job. Enough that it's useful in the prototyping phase. 4. PyTorch JIT is neat, and is what I meant by Pytorch team is "working" on it. However, the JIT doesn't do full code gen (thus, significant operator overhead for say, scalar networks) and has significantly less man hours poured into compared to XLA. 5. I was specifically talking about how you call grad on a function to get a function that returns its gradient. It's a cleaner API than PyTorch's autograd. Jax is definitely not meant for deployment or industry usage, and I believe their developers hope they'll never be pushed along that direction :^) I definitely agree that PyTorch is "good enough" for most people. However, among researchers, there's a decent amount of subgroups it could gain favor in. You'd be surprised how many papers get submitted to ICML/Neurips that don't use PyTorch or TensorFlow at all, in favor of raw numpy, C++, or even MatLab! I think the numbers I had were something about 30% of papers don't use any ML framework. Jax could easily gain favor in this crowd. There's also the crowd that cares a lot about higher order gradients. Also, admittedly a specific subgroup, but growing. Meta learning people care a lot. So do Neural ODE people. All it takes is for one of these subfields to blow up for higher order gradients to all of a sudden become a lot more appealing. And finally, you have Google. Google researchers are never going to use PyTorch en masse (probably). If researchers at Google want to switch from TF, their only option is Jax. This is a pretty big subgroup of researchers :) I definitely agree that Jax has a difficult hill to climb. But, they have a solid foothold within Google, and several subfields very amenable to their advantages. PyTorch seems like the predominant research framework currently, but if any framework is going to erode their lead, I'd place my bets on Jax.

m0zg 7y ago

From a practitioner: 1. Meh. PyTorch is close enough to not worry about it, and is better in some places. 2. Meh. All the methods people use in practice for deep learning in particular do not use higher order gradients. Most higher order methods are prohibitively memory expensive, and memory is at a premium in acceleration hardware (and so is the bus bandwidth - so you can't "swap to RAM"). I do agree that higher order gradients are the next frontier in optimization though - current optimizer research seems to have stalled, so people focus on training with huge batches and stuff like that. Most SOTA models in my field are trained with SGD+momentum - super primitive stuff. I don't see how Jax would solve the memory problem though. You still have to store those Hessians somewhere, at least partially. 3. Do agree, that's cool if it actually parallelizes nontrivial stuff which e.g. tf.vectorized_map barfs on. Although in a lot of cases you can "vectorize" by concatenating input tensors into a higher dimensional tensor. 4. Meh. Not sure why I'd want that if I already have tracing and JIT. 5. This is #2 With PyTorch though, you get close enough to Numpy to feel at home in both, and there's so much code written for it already that you can usually find a good starting point for your research pretty easily on github and then build on top of that. If you need to deploy, there's also tracing and jit, which lets you load and serve models with libtorch. I see what you're saying regarding "advantages", I'm just pointing out that PyTorch might be "good enough" for most people. If I were on that team, I'd focus on providing comfortable transition from TF 2.x which is a dumpster fire (with the exception of TensorFlow Lite which is excellent). That, IMO, would be the only way for this project to achieve mainstream success unless PyTorch disintegrates over time.

chillee 7y ago
I agree from a practitioner standpoint - but you were talking about research :) 1. Much of the scientific computing/stats community is stuck in the past. Many are still using Matlab! As opposed to the CS community, who are used to learning new frameworks, offering the ability to "import jax.numpy as np" and having their scripts just run is valuable to that community. As is having an API that they've only just started to become familiar with (and has way more documentation about). 2. Once again, this is true for practitioners, but not research. Hessian vector products show up in a decent amount of places. For example, if you have an inner optimization loop (a la. most meta learning approaches or Deep Set Prediction Networks) you have a Hessian Vector Product! Perhaps not prevalent in models that practitioners run but definitely something to keep an eye on in research. 3. My understanding is that it actually does a pretty decent job. Enough that it's useful in the prototyping phase. 4. PyTorch JIT is neat, and is what I meant by Pytorch team is "working" on it. However, the JIT doesn't do full code gen (thus, significant operator overhead for say, scalar networks) and has significantly less man hours poured into compared to XLA. 5. I was specifically talking about how you call grad on a function to get a function that returns its gradient. It's a cleaner API than PyTorch's autograd. Jax is definitely not meant for deployment or industry usage, and I believe their developers hope they'll never be pushed along that direction :^) I definitely agree that PyTorch is "good enough" for most people. However, among researchers, there's a decent amount of subgroups it could gain favor in. You'd be surprised how many papers get submitted to ICML/Neurips that don't use PyTorch or TensorFlow at all, in favor of raw numpy, C++, or even MatLab! I think the numbers I had were something about 30% of papers don't use any ML framework. Jax could easily gain favor in this crowd. There's also the crowd that cares a lot about higher order gradients. Also, admittedly a specific subgroup, but growing. Meta learning people care a lot. So do Neural ODE people. All it takes is for one of these subfields to blow up for higher order gradients to all of a sudden become a lot more appealing. And finally, you have Google. Google researchers are never going to use PyTorch en masse (probably). If researchers at Google want to switch from TF, their only option is Jax. This is a pretty big subgroup of researchers :) I definitely agree that Jax has a difficult hill to climb. But, they have a solid foothold within Google, and several subfields very amenable to their advantages. PyTorch seems like the predominant research framework currently, but if any framework is going to erode their lead, I'd place my bets on Jax.