How does training time scale w.r.t. model size? about ferminet HOT 1 CLOSED

First, the JAX branch is actively developed and much easier (in our view) to use. The main branch is the version used in our first paper. (At some point soon, the JAX version will be made the main branch.) There is a Dockerfile for JAX but I haven't used it so don't know how well it works/how supported it is: https://github.com/google/jax/blob/main/build/Dockerfile

https://arxiv.org/abs/1909.02487 shows the cost per iteration for a fixed network scales as O(N^4). The number of iterations required and the size of the network to reach a desired accuracy is also likely to be system dependent. Something like benzene or caffeine will require a multi-GPU setup.

The speedup you get from running on GPU depends upon your GPU and CPU, but can easily be a factor of 10, even on a small system (and more on larger systems).

from ferminet.