Overview:
The GaLore optimizer is a memory-efficient variant of the Adam optimizer that projects gradients into a low-rank subspace using orthogonal matrices derived from SVD. This reduces the memory footprint for storing momentum states, making it suitable for training large-scale models while maintaining performance.
Parameters:
learning_rate(float, default=1e-3): The step size for parameter updates.beta1(float, default=0.9): Exponential decay rate for the first moment estimates.beta2(float, default=0.999): Exponential decay rate for the second moment estimates.epsilon(float, default=1e-6): Small constant for numerical stability.weight_decay(float, default=0.0): Coefficient for weight decay.rank(int, optional, default=None): The rank of the low-rank projection; if None, no projection is applied.update_proj_gap(int, optional, default=None): The frequency (in steps) for updating the projection matrix; if None, no projection is applied.scale(float, optional, default=None): Scaling factor applied to the projected-back gradient.projection_type(str, optional, default=None): Type of projection scheme ('std', 'reverse_std', 'right', 'left', 'full', 'random'); if None, no projection is applied.clipnorm(float, optional): Clips gradients by norm.clipvalue(float, optional): Clips gradients by value.global_clipnorm(float, optional): Clips gradients by global norm.use_ema(bool, default=False): Whether to apply Exponential Moving Average to model weights.ema_momentum(float, default=0.99): Momentum for EMA.ema_overwrite_frequency(int, optional): Frequency for overwriting EMA weights.loss_scale_factor(float, optional): Factor for scaling the loss during gradient computation.gradient_accumulation_steps(int, optional): Steps for accumulating gradients.name(str, default="galore"): Name of the optimizer.
Example Usage:
import tensorflow as tf
from optimizers.galore import GaLore
# Instantiate optimizer
optimizer = GaLore(
learning_rate=1e-3,
rank=128,
update_proj_gap=50,
scale=1.0,
projection_type='std'
)
# Compile a model
model.compile(optimizer=optimizer, loss="sparse_categorical_crossentropy", metrics=["accuracy"])
# Train the model
model.fit(train_dataset, validation_data=val_dataset, epochs=10)