optimizer that implements the adafactor algorithm adafactor is commonly used in nlp tasks and has the advantage of taking less memory because it only saves partial information of previous gradients the default argument setup is based on the original paper see reference when gradients are of dimension 2 adafactor optimizer will delete the last 2 dimensions separately in its accumulator variables args learningrate a float a keras optimizers schedules learningrateschedule instance or a callable that takes no arguments and returns the actual value to use the learning rate defaults to 0 001 beta2decay float defaults to 0 8 the decay rate of beta2 epsilon1 float defaults to 1e30 a small offset to keep denominator away from 0 epsilon2 float defaults to 1e3 a small offset to avoid learning rate becoming too small by time clipthreshold float defaults to 1 0 clipping threshold this is a part of adafactor algorithm independent from clipnorm clipvalue and globalclipnorm relativestep bool defaults to true if learningrate is a constant and relativesteptrue learning rate will be adjusted based on current iterations this is a default learning rate decay in adafactor baseoptimizerkeywordargs reference shazeer noam et al 2018https arxiv orgabs1804 04235 initialize optimizer variables adam optimizer has 3 types of variables momentums velocities and velocityhat only set when amsgrad is applied args varlist list of model variables to build adam variables on don t factor if variable is of dimension 2 but we still need to create dummy variables as placeholder always factor the last 2 dimensions update step given gradient and the associated model variable lr ops castlearningrate variable dtype gradient ops castgradient variable dtype epsilon2 ops castself epsilon2 variable dtype one ops cast1 0 variable dtype localstep ops castself iterations 1 variable dtype if not callableself learningrate and self relativestep lr ops minimumlr 1 ops sqrtlocalstep r self rself getvariableindexvariable c self cself getvariableindexvariable v self vself getvariableindexvariable rhot ops minimumlr 1 ops sqrtlocalstep alphat ops maximumepsilon2 self rmsvariable rhot regulatedgradsquare ops addops squaregradient self epsilon1 beta2t 1 ops powerlocalstep self beta2decay if lenvariable shape 2 r deletes the last dimension of gradient so it is of shape gradient shape 1 self assign r beta2t r 1 beta2t ops meanregulatedgradsquare axis1 c deletes the second last dimension of gradient so it is of shape gradient shape 2 gradient shape1 self assign c beta2t c 1 beta2t ops meanregulatedgradsquare axis2 self assign v ops expanddims r ops meanr axis1 keepdimstrue axis1 ops expanddimsc 2 else self assign v beta2t v 1 beta2t regulatedgradsquare converttotensor unifies the handling of sparse and dense grads ut ops dividegradient ops sqrtv uthat ops divide ut ops maximumone ops divideself rmsut self clipthreshold self assignsubvariable ops multiplyalphat uthat def getconfigself config super getconfig config update beta2decay self beta2decay epsilon1 self epsilon1 epsilon2 self epsilon2 clipthreshold self clipthreshold relativestep self relativestep return config adafactor doc adafactor doc replace baseoptimizerkeywordargs optimizer baseoptimizerkeywordargs optimizer that implements the adafactor algorithm adafactor is commonly used in nlp tasks and has the advantage of taking less memory because it only saves partial information of previous gradients the default argument setup is based on the original paper see reference when gradients are of dimension 2 adafactor optimizer will delete the last 2 dimensions separately in its accumulator variables args learning_rate a float a keras optimizers schedules learningrateschedule instance or a callable that takes no arguments and returns the actual value to use the learning rate defaults to 0 001 beta_2_decay float defaults to 0 8 the decay rate of beta_2 epsilon_1 float defaults to 1e 30 a small offset to keep denominator away from 0 epsilon_2 float defaults to 1e 3 a small offset to avoid learning rate becoming too small by time clip_threshold float defaults to 1 0 clipping threshold this is a part of adafactor algorithm independent from clipnorm clipvalue and global_clipnorm relative_step bool defaults to true if learning_rate is a constant and relative_step true learning rate will be adjusted based on current iterations this is a default learning rate decay in adafactor base_optimizer_keyword_args reference shazeer noam et al 2018 https arxiv org abs 1804 04235 initialize optimizer variables adam optimizer has 3 types of variables momentums velocities and velocity_hat only set when amsgrad is applied args var_list list of model variables to build adam variables on don t factor if variable is of dimension 2 but we still need to create dummy variables as placeholder always factor the last 2 dimensions update step given gradient and the associated model variable r deletes the last dimension of gradient so it is of shape gradient shape 1 c deletes the second last dimension of gradient so it is of shape gradient shape 2 gradient shape 1 convert_to_tensor unifies the handling of sparse and dense grads
from keras import backend
from keras import ops
from keras.api_export import keras_export
from keras.optimizers import optimizer


@keras_export(["keras.optimizers.Adafactor"])
class Adafactor(optimizer.Optimizer):
    

    def __init__(
        self,
        learning_rate=0.001,
        beta_2_decay=-0.8,
        epsilon_1=1e-30,
        epsilon_2=1e-3,
        clip_threshold=1.0,
        relative_step=True,
        weight_decay=None,
        clipnorm=None,
        clipvalue=None,
        global_clipnorm=None,
        use_ema=False,
        ema_momentum=0.99,
        ema_overwrite_frequency=None,
        name="adafactor",
        **kwargs,
    ):
        super().__init__(
            learning_rate=learning_rate,
            name=name,
            weight_decay=weight_decay,
            clipnorm=clipnorm,
            clipvalue=clipvalue,
            global_clipnorm=global_clipnorm,
            use_ema=use_ema,
            ema_momentum=ema_momentum,
            ema_overwrite_frequency=ema_overwrite_frequency,
            **kwargs,
        )
        self.beta_2_decay = beta_2_decay
        self.epsilon_1 = epsilon_1
        self.epsilon_2 = epsilon_2
        self.clip_threshold = clip_threshold
        self.relative_step = relative_step

    def build(self, var_list):
        
        if self.built:
            return
        super().build(var_list)
        self._r = []
        self._c = []
        self._v = []
        for var in var_list:
            if len(var.shape) < 2:
                
                
                with backend.name_scope(self.name, caller=self):
                    self._r.append(
                        backend.Variable(0, name=var.name, trainable=False)
                    )
                    self._c.append(
                        backend.Variable(0, name=var.name, trainable=False)
                    )
            else:
                
                r_shape = var.shape[:-1]
                c_shape = var.shape[:-2] + var.shape[-1]
                self._r.append(
                    self.add_variable(
                        shape=r_shape,
                        dtype=var.dtype,
                        name=var.name,
                    )
                )
                self._c.append(
                    self.add_variable(
                        shape=c_shape,
                        dtype=var.dtype,
                        name=var.name,
                    )
                )
            self._v.append(
                self.add_variable_from_reference(
                    reference_variable=var, name="velocity"
                )
            )

    def _rms(self, x):
        return ops.sqrt(ops.mean(ops.square(x)))

    def update_step(self, gradient, variable, learning_rate):
        

        lr = ops.cast(learning_rate, variable.dtype)
        gradient = ops.cast(gradient, variable.dtype)
        epsilon_2 = ops.cast(self.epsilon_2, variable.dtype)
        one = ops.cast(1.0, variable.dtype)
        local_step = ops.cast(self.iterations + 1, variable.dtype)
        if not callable(self._learning_rate) and self.relative_step:
            lr = ops.minimum(lr, 1 / ops.sqrt(local_step))

        r = self._r[self._get_variable_index(variable)]
        c = self._c[self._get_variable_index(variable)]
        v = self._v[self._get_variable_index(variable)]

        rho_t = ops.minimum(lr, 1 / ops.sqrt(local_step))
        alpha_t = ops.maximum(epsilon_2, self._rms(variable)) * rho_t
        regulated_grad_square = ops.add(ops.square(gradient), self.epsilon_1)
        beta_2_t = 1 - ops.power(local_step, self.beta_2_decay)

        if len(variable.shape) >= 2:
            
            
            self.assign(
                r,
                beta_2_t * r
                + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-1),
            )
            
            
            self.assign(
                c,
                beta_2_t * c
                + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-2),
            )
            self.assign(
                v,
                ops.expand_dims(
                    r / ops.mean(r, axis=-1, keepdims=True), axis=-1
                )
                * ops.expand_dims(c, -2),
            )
        else:
            self.assign(
                v, beta_2_t * v + (1 - beta_2_t) * regulated_grad_square
            )

        
        u_t = ops.divide(gradient, ops.sqrt(v))
        u_t_hat = ops.divide(
            u_t,
            ops.maximum(one, ops.divide(self._rms(u_t), self.clip_threshold)),
        )
        self.assign_sub(variable, ops.multiply(alpha_t, u_t_hat))

    def get_config(self):
        config = super().get_config()

        config.update(
            {
                "beta_2_decay": self.beta_2_decay,
                "epsilon_1": self.epsilon_1,
                "epsilon_2": self.epsilon_2,
                "clip_threshold": self.clip_threshold,
                "relative_step": self.relative_step,
            }
        )
        return config


Adafactor.__doc__ = Adafactor.__doc__.replace(
    "{{base_optimizer_keyword_args}}", optimizer.base_optimizer_keyword_args
)