Optax Documentation

Optax provides a wide variety of gradient transformations that can be composed to create complex optimizers. These transformations modify gradients before they are applied to parameters.

Common categories of transformations include:

• Scaling: Adjusting gradient magnitude using methods like scale_by_adam, scale_by_lion, scale_by_rms, or scale_by_learning_rate.
• Clipping: Preventing exploding gradients using clip, clip_by_global_norm, or adaptive_grad_clip.
• Noise/Regularization: Adding stochasticity via add_noise or managing weight decay with add_decayed_weights.
• State Management: Using ema (Exponential Moving Average) or snapshot to track statistics.
• Masking/Conditional: Using conditionally_mask or selective_transform to apply updates only to specific parameters.

Each transformation typically returns an OptState (or a specific state like ScaleByAdamState) which must be maintained and passed back into the optimizer during training loops.

Optax provides several wrapper functions to modify the behavior of existing optimizers. These wrappers allow you to add logic such as skipping updates when gradients are not finite, applying lookahead mechanisms, or masking certain updates.

Commonly used wrappers include:

• apply_if_finite: Only applies the optimizer update if the gradients are finite.
• lookahead: Implements the lookahead optimization algorithm.
• masked: Applies the optimizer update only to specific indices or elements.
• skip_large_updates: Skips updates if the gradient magnitude exceeds a certain threshold.
• skip_not_finite: Skips updates if gradients contain NaNs or Infs.
• MultiSteps: Wraps an optimizer to perform multiple steps of the inner optimizer for every single step of the outer optimizer.

When using these wrappers, you will typically interact with their associated State objects (e.g., ApplyIfFiniteState, LookaheadState, MaskedState, MultiStepsState) to manage the internal statistics required by the wrapper.

Optax provides projection functions in optax.projections to perform constrained optimization. A projection onto a set $\mathcal{C}$ finds the closest point $v$ in that set to a given point $u$ by minimizing the squared Euclidean distance: $\text{proj}{\mathcal{C}}(u) := \text{argmin}{v} |u - v|^2_2$ subject to $v \in \mathcal{C}$.

In a typical training loop, you apply the projection to your parameters immediately after applying the optimizer updates to ensure the parameters remain within the desired constraint set.

import optax
import jax
import jax.numpy as jnp

# Example: Projecting parameters to the non-negative orthant
num_weights = 2
xs = jnp.array([[-1.8, 2.2], [-2.0, 1.2]])
ys = jnp.array([0.5, 0.8])
optimizer = optax.adam(learning_rate=1e-3)
params = {'w': jnp.zeros(num_weights)}
opt_state = optimizer.init(params)

loss = lambda params, x, y: jnp.mean((params['w'].dot(x) - y) ** 2)

# Standard update step
grads = jax.grad(loss)(params, xs, ys)
updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)

# Apply projection to enforce non-negativity
params = optax.projections.projection_non_negative(params)

Optax provides a variety of schedules in optax.schedules to control how hyperparameters (like learning rates) change over time. These schedules are typically used as arguments to optimizer functions.

Commonly available schedules include:

• Linear/Polynomial: linear_schedule, polynomial_schedule.
• Decay-based: exponential_decay, cosine_decay_schedule.
• Warmup-based: warmup_constant_schedule, warmup_cosine_decay_schedule, warmup_exponential_decay_schedule.
• Cyclic/One-cycle: cosine_onecycle_schedule, linear_onecycle_schedule.
• Piecewise: piecewise_constant_schedule, piecewise_interpolate_schedule.
• Stochastic/Advanced: sgdr_schedule (Stochastic Gradient Descent with Restarts).
• Composition: join_schedules allows combining multiple schedules.

All schedules follow the Schedule abstraction, which typically takes a step count as input and returns the scheduled value.

The Optax example gallery provides a collection of practical implementations and research-oriented use cases. These examples demonstrate how to use Optax for standard deep learning tasks, integration with other libraries like Flax, and advanced optimization techniques.

Core Examples include:

• Standard Training: MLP classifier on MNIST, ResNet on CIFAR10 (with Flax), and Character-level Transformers.
• Advanced Optimization: Gradient Accumulation, Meta-Learning, Optimistic Gradient Descent (OGDA), LBFGS with linesearch, and Lookahead wrappers.
• Specialized Tasks: Adversarial training, solving linear assignment problems, and differentiable functions with perturbations.

Contribution-based Examples (using optax.contrib):

• Differential Privacy: Differentially private SGD for CNNs on MNIST.
• Learning Rate Schedulers: Usage of reduce_on_plateau.
• Advanced Algorithms: Sharpness-Aware Minimization (SAM) and AdEMAMix.