Skip to main content

Custom Training Loops

In research, standard training loops often aren't enough. You might need to track complex metrics, manage multiple optimizers, or implement non-standard update rules. Flax's design philosophy is to expose the training state explicitly, giving you full control.

Prerequisites

A research-grade guide and the foundation for most of this section. Comfortable with training loops and training best practices? Good — this page assumes you have written a basic step before and now want full control over it.

What you'll learn
  • Extend TrainState with batch_stats, an explicit dropout_rng, metrics, and ema_params
  • Maintain an exponential moving average of parameters for stable evaluation
  • Build an advanced JIT-compiled step with label smoothing and manual L2 regularization
  • Clip gradients by global norm and track the gradient scale as a metric

Why Custom Loops?​

Frameworks often abstract away the training loop (model.fit()), which is great for standard tasks but restricting for research. Custom loops allow you to:

  • Modify gradients before application (e.g., gradient clipping, noise injection)
  • Update auxiliary state like Batch Norm statistics or EMA inputs
  • Implement complex logic like GAN training (discriminator vs generator steps)
  • Debug easily by inspecting every intermediate tensor

Flexible Training State​

The TrainState pattern is central to Flax. It holds everything that changes during training: parameters, optimizer state, and mutable variables. For research, we often need to extend this.

Extending TrainState​

Here we create a CustomTrainState that includes:

  • batch_stats: For Batch Normalization.
  • dropout_rng: Explicit random key for dropout, ensuring reproducibility.
  • metrics: To track arbitrary values during the step.
  • ema_params: Exponential Moving Average of parameters for stable evaluation.
from flax import struct
from flax.training import train_state
import optax
import jax

@struct.dataclass
class CustomTrainState(train_state.TrainState):
    """Extended training state for research experiments."""
    
    # Batch statistics (e.g., for BatchNorm)
    batch_stats: dict
    
    # Random key for dropout (updated every step if needed)
    dropout_rng: jax.random.PRNGKey
    
    # Store arbitrary metrics for the current step
    metrics: dict
    
    # Exponential Moving Average of parameters
    ema_params: dict
    
    def apply_ema(self, decay=0.999):
        """
        Apply exponential moving average to parameters.
        
        New EMA = decay * Old EMA + (1 - decay) * Current Params
        """
        new_ema = jax.tree_map(
            lambda ema, p: decay * ema + (1 - decay) * p,
            self.ema_params,
            self.params
        )
        return self.replace(ema_params=new_ema)

Advanced Training Step​

A research-grade training step often involves more than just loss.backward(). The example below demonstrates:

  1. State Management: Handling mutable variables (batch_stats).
  2. Auxiliary Losses: Adding L2 regularization manually.
  3. Advanced Optimization: Gradient clipping and global norm tracking.
  4. EMA Updates: Maintaining a separate set of smoothed parameters.
import jax.numpy as jnp

@jax.jit
def advanced_train_step(state, batch, dropout_rng):
    """
    Execute a single training step with advanced features.
    
    Args:
        state: Current CustomTrainState.
        batch: Dictionary containing 'image' and 'label'.
        dropout_rng: PRNGKey for dropout.
        
    Returns:
        Updated state and metrics dictionary.
    """
    
    def loss_fn(params):
        # 1. Forward Pass
        # We must pass mutable=['batch_stats'] to capture updates
        variables = {'params': params, 'batch_stats': state.batch_stats}
        logits, new_variables = state.apply_fn(
            variables,
            batch['image'],
            training=True,
            rngs={'dropout': dropout_rng},
            mutable=['batch_stats']
        )
        
        # 2. Label Smoothing
        # Instead of hard 0/1 labels, we smooth them to prevent overconfidence
        labels_one_hot = jax.nn.one_hot(batch['label'], num_classes=10)
        labels_smooth = labels_one_hot * 0.9 + 0.1 / 10
        
        loss = optax.softmax_cross_entropy(
            logits=logits,
            labels=labels_smooth
        ).mean()
        
        # 3. Auxiliary Losses (L2 Regularization)
        # Manually compute L2 norm of all parameters
        l2_loss = 0.5 * sum(
            jnp.sum(jnp.square(p)) for p in jax.tree_util.tree_leaves(params)
        )
        total_loss = loss + 1e-4 * l2_loss
        
        return total_loss, (logits, new_variables['batch_stats'])
    
    # 4. Compute Gradients
    # has_aux=True tells JAX that loss_fn returns extra data (logits, new_stats)
    grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
    (loss, (logits, new_batch_stats)), grads = grad_fn(state.params)
    
    # 5. Gradient Clipping
    # Prevent gradients from exploding by clipping their global norm
    grads = optax.clip_by_global_norm(1.0).update(grads, state)
    
    # 6. Update State
    state = state.apply_gradients(
        grads=grads,
        batch_stats=new_batch_stats
    )
    
    # 7. Update EMA Parameters
    state = state.apply_ema(decay=0.999)
    
    # 8. Compute Metrics
    accuracy = jnp.mean(jnp.argmax(logits, -1) == batch['label'])
    metrics = {
        'loss': loss,
        'accuracy': accuracy,
        'grad_norm': optax.global_norm(grads),  # Track gradient scale
    }
    
    return state, metrics

Best Practices​

  • JIT Compilation: Always decorate your step function with @jax.jit for performance.
  • Pure Functions: Ensure your step function is pure (no side effects). Return the new state explicitly.
  • Metric Collection: Return a dictionary of scalar metrics. This makes logging to tools like TensorBoard or Weights & Biases trivial.
  • Debugging: If you hit NaNs, use jax.debug.print inside the JIT-compiled function or temporarily disable JIT with jax.disable_jit() context manager.

Next steps​