How do you keep a JAX training loop’s moving parts together? Flax’s TrainState packages the model’s apply function, parameters, optimizer transformation, optimizer state, and step information into one functional training value. section 11.1.3 section 11.4
Create a state with TrainState.create(), supplying model.apply, initialized parameters, and an Optax transformation. Initialization of the optimizer state happens inside create(). During each update, call state.apply_fn to run the model and state.apply_gradients(grads=grads) to update parameters and optimizer state. Because the updated state is returned, the loop remains compatible with JAX’s functional transformations. section 11.1.3 section 11.4
A basic TrainState provides:
- apply_fn, usually model.apply, for applying the network to variables and inputs. section 11.1.3
- params for the trainable model parameters updated by the optimizer. section 11.1.3
- tx for an Optax gradient transformation such as SGD with momentum. section 11.1.3
- A single-optimizer design suited to the common case; add metrics or other state by subclassing TrainState. section 11.1.3
TrainState is extensible rather than universal. A metrics subclass can hold CLU collections and update them with replace(). Models with internal state, such as BatchNorm, require another field—often model_state—and explicit mutable collections during apply. section 11.2.1
