Kron: A TensorFlow Implementation of the PSGD Kron Optimizer

The Kron optimizer provides an efficient implementation of the PSGD Kron algorithm, leveraging Kronecker-based preconditioning techniques to enhance the performance of stochastic gradient descent (SGD). This optimizer is specifically designed for large-scale models, where effective preconditioning can lead to significantly improved convergence rates while optimizing memory usage.

Key Features:

• Adaptive Gradient Direction: By maintaining a collection of per-parameter preconditioners that are probabilistically updated during training, Kron intelligently adjusts the gradient direction and scaling, facilitating faster training times.
• Memory Efficiency: Offers options to control memory consumption with specific settings for preconditioners, making it suitable for applications with limited resources.
• Customizability: Provides several hyperparameters that can be tuned to adapt the optimization process according to the specific needs of the model being trained.

Major Parameters:

• learning_rate (float, default=0.0003): Sets the step size for updating model parameters.
• weight_decay (float, default=0.0): Coefficient for weight decay regularization.
• b1 (float, default=0.9): Exponential decay rate for the momentum buffer updates.
• preconditioner_update_probability (callable or float, optional): Controls how often the preconditioner is updated, with a default schedule available.
• max_size_triangular (int, default=8192): Maximum size for applying a triangular preconditioner; larger dimensions use diagonal approximations.
• min_ndim_triangular (int, default=2): Minimum dimensions required to utilize a triangular preconditioner.
• memory_save_mode (str, optional): Modes to optimize memory utilization for preconditioners.
• clipnorm, clipvalue, and global_clipnorm: Parameters for gradient clipping to enhance stability during training.

Example Usage:

Integrating the Kron optimizer into a TensorFlow model is straightforward. Below is an example of how to implement the Kron optimizer within a Keras training workflow:

import tensorflow as tf

# Instantiate the Kron optimizer with default preconditioner schedule.
optimizer = Kron(
    learning_rate=0.0003,
    weight_decay=1e-4,
    b1=0.9,
    max_size_triangular=8192,
    min_ndim_triangular=2,
    memory_save_mode="smart_one_diag",
    momentum_into_precond_update=True,
    precond_lr=0.1,
    precond_init_scale=1.0,
    mu_dtype=tf.float32,
    precond_dtype=tf.float32,
)

# Compile a Keras model using the Kron optimizer.
model.compile(
    optimizer=optimizer,
    loss="sparse_categorical_crossentropy",
    metrics=["accuracy"]
)

# Train the model
model.fit(train_dataset, validation_data=val_dataset, epochs=10)

Kron is an essential tool for practitioners looking to optimize complex machine learning models with efficient parameter updates, saving both time and computational resources.